<commit_before>package ibclient

import (
	"fmt"
	"net"
)

func (objMgr *ObjectManager) CreateHostRecord(
	enabledns bool,
	enableDhcp bool,
	recordName string,
	netview string,
	dnsview string,
	ipv4cidr string,
	ipv6cidr string,
	ipv4Addr string,
	ipv6Addr string,
	macAddr string,
	duid string,
	useTtl bool,
	ttl uint32,
	comment string,
	eas EA,
	aliases []string) (*HostRecord, error) {

	if ipv4Addr == "" && ipv4cidr != "" {
		if netview == "" {
			netview = "default"
		}
		ipv4Addr = fmt.Sprintf("func:nextavailableip:%s,%s", ipv4cidr, netview)
	}
	if ipv6Addr == "" && ipv6cidr != "" {
		if netview == "" {
			netview = "default"
		}
		ipv6Addr = fmt.Sprintf("func:nextavailableip:%s,%s", ipv6cidr, netview)
	}
	recordHost := NewEmptyHostRecord()
	recordHostIpv6AddrSlice := []HostRecordIpv6Addr{}
	recordHostIpv4AddrSlice := []HostRecordIpv4Addr{}
	if ipv6Addr != "" {
		enableDhcpv6 := false
		if enableDhcp && duid != "" {
			enableDhcpv6 = true
		}
		recordHostIpv6Addr := NewHostRecordIpv6Addr(ipv6Addr, duid, enableDhcpv6, "")
		recordHostIpv6AddrSlice = []HostRecordIpv6Addr{*recordHostIpv6Addr}
	}
	if ipv4Addr != "" {
		enableDhcpv4 := false
		if enableDhcp && macAddr != "" && macAddr != MACADDR_ZERO {
			enableDhcpv4 = true
		}
		recordHostIpAddr := NewHostRecordIpv4Addr(ipv4Addr, macAddr, enableDhcpv4, "")
		recordHostIpv4AddrSlice = []HostRecordIpv4Addr{*recordHostIpAddr}
	}
	recordHost = NewHostRecord(
		netview, recordName, "", "", recordHostIpv4AddrSlice, recordHostIpv6AddrSlice,
		eas, enabledns, dnsview, "", "", useTtl, ttl, comment, aliases)
	ref, err := objMgr.connector.CreateObject(recordHost)
	if err != nil {
		return nil, err
	}
	recordHost.Ref = ref
	err = objMgr.connector.GetObject(
		recordHost, ref, NewQueryParams(false, nil), &recordHost)
	return recordHost, err
}

func (objMgr *ObjectManager) GetHostRecordByRef(ref string) (*HostRecord, error) {
	recordHost := NewEmptyHostRecord()
	err := objMgr.connector.GetObject(
		recordHost, ref, NewQueryParams(false, nil), &recordHost)

	return recordHost, err
}

func (objMgr *ObjectManager) SearchHostRecordByAltId(
	internalId string, ref string, eaNameForInternalId string) (*HostRecord, error) {

	if internalId == "" {
		return nil, fmt.Errorf("internal ID must not be empty")
	}

	recordHost := NewEmptyHostRecord()
	if ref != "" {
		if err := objMgr.connector.GetObject(recordHost, ref, NewQueryParams(false, nil), &recordHost); err != nil {
			return nil, err
		}
	}
	if recordHost.Ref != "" {
		return recordHost, nil
	}
	recordHost.eaSearch = map[string]interface{}{
		eaNameForInternalId: internalId,
	}

	res := make([]HostRecord, 0)
	err := objMgr.connector.GetObject(recordHost, "", NewQueryParams(false, nil), &res)

	if err != nil {
		return nil, err
	}

	if res == nil || len(res) == 0 {
		return nil, NewNotFoundError("host record not found")
	}

	result := res[0]

	return &result, nil
}

func (objMgr *ObjectManager) GetHostRecord(netview string, dnsview string, recordName string, ipv4addr string, ipv6addr string) (*HostRecord, error) {
	var res []HostRecord

	recordHost := NewEmptyHostRecord()

	sf := map[string]string{
		"name": recordName,
	}
	if netview != "" {
		sf["network_view"] = netview
	}
	if dnsview != "" {
		sf["view"] = dnsview
	}
	if ipv4addr != "" {
		sf["ipv4addr"] = ipv4addr
	}
	if ipv6addr != "" {
		sf["ipv6addr"] = ipv6addr
	}
	queryParams := NewQueryParams(false, sf)
	err := objMgr.connector.GetObject(recordHost, "", queryParams, &res)

	if err != nil || res == nil || len(res) == 0 {
		return nil, err
	}
	return &res[0], err

}

func (objMgr *ObjectManager) GetIpAddressFromHostRecord(host HostRecord) (string, error) {
	err := objMgr.connector.GetObject(
		&host, host.Ref, NewQueryParams(false, nil), &host)
	return host.Ipv4Addrs[0].Ipv4Addr, err
}

func (objMgr *ObjectManager) UpdateHostRecord(
	hostRref string,
	enabledns bool,
	enableDhcp bool,
	name string,
	netView string,
	ipv4cidr string,
	ipv6cidr string,
	ipv4Addr string,
	ipv6Addr string,
	macAddr string,
	duid string,
	useTtl bool,
	ttl uint32,
	comment string,
	eas EA,
	aliases []string) (*HostRecord, error) {

	recordHostIpv4AddrSlice := []HostRecordIpv4Addr{}
	recordHostIpv6AddrSlice := []HostRecordIpv6Addr{}

	enableDhcpv4 := false
	enableDhcpv6 := false
	if ipv6Addr != "" {
		if enableDhcp && duid != "" {
			enableDhcpv6 = true
		}
	}
	if ipv4Addr != "" {
		if enableDhcp && macAddr != "" && macAddr != MACADDR_ZERO {
			enableDhcpv4 = true
		}
	}

	if ipv4Addr == "" {
		if ipv4cidr != "" {
			ip, _, err := net.ParseCIDR(ipv4cidr)
			if err != nil {
				return nil, fmt.Errorf("cannot parse CIDR value: %s", err.Error())
			}
			if ip.To4() == nil {
				return nil, fmt.Errorf("CIDR value must be an IPv4 CIDR, not an IPv6 one")
			}
			if netView == "" {
				netView = "default"
			}
			newIpAddr := fmt.Sprintf("func:nextavailableip:%s,%s", ipv4cidr, netView)
			recordHostIpAddr := NewHostRecordIpv4Addr(newIpAddr, macAddr, enableDhcpv4, "")
			recordHostIpv4AddrSlice = []HostRecordIpv4Addr{*recordHostIpAddr}
		}
	} else {
		ip := net.ParseIP(ipv4Addr)
		if ip == nil {
			return nil, fmt.Errorf("'IP address for the record is not valid")
		}
		if ip.To4() == nil {
			return nil, fmt.Errorf("IP address must be an IPv4 address, not an IPv6 one")
		}
		recordHostIpAddr := NewHostRecordIpv4Addr(ipv4Addr, macAddr, enableDhcpv4, "")
		recordHostIpv4AddrSlice = []HostRecordIpv4Addr{*recordHostIpAddr}
	}
	if ipv6Addr == "" {
		if ipv6cidr != "" {
			ip, _, err := net.ParseCIDR(ipv6cidr)
			if err != nil {
				return nil, fmt.Errorf("cannot parse CIDR value: %s", err.Error())
			}
			if ip.To4() != nil || ip.To16() == nil {
				return nil, fmt.Errorf("CIDR value must be an IPv6 CIDR, not an IPv4 one")
			}
			if netView == "" {
				netView = "default"
			}
			newIpAddr := fmt.Sprintf("func:nextavailableip:%s,%s", ipv6cidr, netView)
			recordHostIpAddr := NewHostRecordIpv6Addr(newIpAddr, duid, enableDhcpv6, "")
			recordHostIpv6AddrSlice = []HostRecordIpv6Addr{*recordHostIpAddr}
		}
	} else {
		ip := net.ParseIP(ipv6Addr)
		if ip == nil {
			return nil, fmt.Errorf("IP address for the record is not valid")
		}
		if ip.To4() != nil || ip.To16() == nil {
			return nil, fmt.Errorf("IP address must be an IPv6 address, not an IPv4 one")
		}
		recordHostIpAddr := NewHostRecordIpv6Addr(ipv6Addr, duid, enableDhcpv6, "")
		recordHostIpv6AddrSlice = []HostRecordIpv6Addr{*recordHostIpAddr}
	}
	updateHostRecord := NewHostRecord(
		"", name, "", "", recordHostIpv4AddrSlice, recordHostIpv6AddrSlice,
		eas, enabledns, "", "", hostRref, useTtl, ttl, comment, aliases)
	ref, err := objMgr.connector.UpdateObject(updateHostRecord, hostRref)
	if err != nil {
		return nil, err
	}

	updateHostRecord, err = objMgr.GetHostRecordByRef(ref)
	if err != nil {
		return nil, err
	}
	return updateHostRecord, nil
}

func (objMgr *ObjectManager) DeleteHostRecord(ref string) (string, error) {
	return objMgr.connector.DeleteObject(ref)
}
<commit_msg>fixing ObjectManager.SearchHostRecordByAltId()<commit_after>package ibclient

import (
	"fmt"
	"net"
)

func (objMgr *ObjectManager) CreateHostRecord(
	enabledns bool,
	enableDhcp bool,
	recordName string,
	netview string,
	dnsview string,
	ipv4cidr string,
	ipv6cidr string,
	ipv4Addr string,
	ipv6Addr string,
	macAddr string,
	duid string,
	useTtl bool,
	ttl uint32,
	comment string,
	eas EA,
	aliases []string) (*HostRecord, error) {

	if ipv4Addr == "" && ipv4cidr != "" {
		if netview == "" {
			netview = "default"
		}
		ipv4Addr = fmt.Sprintf("func:nextavailableip:%s,%s", ipv4cidr, netview)
	}
	if ipv6Addr == "" && ipv6cidr != "" {
		if netview == "" {
			netview = "default"
		}
		ipv6Addr = fmt.Sprintf("func:nextavailableip:%s,%s", ipv6cidr, netview)
	}
	recordHost := NewEmptyHostRecord()
	recordHostIpv6AddrSlice := []HostRecordIpv6Addr{}
	recordHostIpv4AddrSlice := []HostRecordIpv4Addr{}
	if ipv6Addr != "" {
		enableDhcpv6 := false
		if enableDhcp && duid != "" {
			enableDhcpv6 = true
		}
		recordHostIpv6Addr := NewHostRecordIpv6Addr(ipv6Addr, duid, enableDhcpv6, "")
		recordHostIpv6AddrSlice = []HostRecordIpv6Addr{*recordHostIpv6Addr}
	}
	if ipv4Addr != "" {
		enableDhcpv4 := false
		if enableDhcp && macAddr != "" && macAddr != MACADDR_ZERO {
			enableDhcpv4 = true
		}
		recordHostIpAddr := NewHostRecordIpv4Addr(ipv4Addr, macAddr, enableDhcpv4, "")
		recordHostIpv4AddrSlice = []HostRecordIpv4Addr{*recordHostIpAddr}
	}
	recordHost = NewHostRecord(
		netview, recordName, "", "", recordHostIpv4AddrSlice, recordHostIpv6AddrSlice,
		eas, enabledns, dnsview, "", "", useTtl, ttl, comment, aliases)
	ref, err := objMgr.connector.CreateObject(recordHost)
	if err != nil {
		return nil, err
	}
	recordHost.Ref = ref
	err = objMgr.connector.GetObject(
		recordHost, ref, NewQueryParams(false, nil), &recordHost)
	return recordHost, err
}

func (objMgr *ObjectManager) GetHostRecordByRef(ref string) (*HostRecord, error) {
	recordHost := NewEmptyHostRecord()
	err := objMgr.connector.GetObject(
		recordHost, ref, NewQueryParams(false, nil), &recordHost)

	return recordHost, err
}

func (objMgr *ObjectManager) SearchHostRecordByAltId(
	internalId string, ref string, eaNameForInternalId string) (*HostRecord, error) {

	if internalId == "" {
		return nil, fmt.Errorf("internal ID must not be empty")
	}

	recordHost := NewEmptyHostRecord()
	if ref != "" {
		if err := objMgr.connector.GetObject(recordHost, ref, NewQueryParams(false, nil), &recordHost); err != nil {
			return nil, err
		}
	}
	if recordHost.Ref != "" {
		return recordHost, nil
	}
	sf := map[string]string{
		fmt.Sprintf("*%s", eaNameForInternalId): internalId,
	}

	res := make([]HostRecord, 0)
	err := objMgr.connector.GetObject(recordHost, "", NewQueryParams(false, sf), &res)

	if err != nil {
		return nil, err
	}

	if res == nil || len(res) == 0 {
		return nil, NewNotFoundError("host record not found")
	}

	result := res[0]

	return &result, nil
}

func (objMgr *ObjectManager) GetHostRecord(netview string, dnsview string, recordName string, ipv4addr string, ipv6addr string) (*HostRecord, error) {
	var res []HostRecord

	recordHost := NewEmptyHostRecord()

	sf := map[string]string{
		"name": recordName,
	}
	if netview != "" {
		sf["network_view"] = netview
	}
	if dnsview != "" {
		sf["view"] = dnsview
	}
	if ipv4addr != "" {
		sf["ipv4addr"] = ipv4addr
	}
	if ipv6addr != "" {
		sf["ipv6addr"] = ipv6addr
	}
	queryParams := NewQueryParams(false, sf)
	err := objMgr.connector.GetObject(recordHost, "", queryParams, &res)

	if err != nil || res == nil || len(res) == 0 {
		return nil, err
	}
	return &res[0], err

}

func (objMgr *ObjectManager) GetIpAddressFromHostRecord(host HostRecord) (string, error) {
	err := objMgr.connector.GetObject(
		&host, host.Ref, NewQueryParams(false, nil), &host)
	return host.Ipv4Addrs[0].Ipv4Addr, err
}

func (objMgr *ObjectManager) UpdateHostRecord(
	hostRref string,
	enabledns bool,
	enableDhcp bool,
	name string,
	netView string,
	ipv4cidr string,
	ipv6cidr string,
	ipv4Addr string,
	ipv6Addr string,
	macAddr string,
	duid string,
	useTtl bool,
	ttl uint32,
	comment string,
	eas EA,
	aliases []string) (*HostRecord, error) {

	recordHostIpv4AddrSlice := []HostRecordIpv4Addr{}
	recordHostIpv6AddrSlice := []HostRecordIpv6Addr{}

	enableDhcpv4 := false
	enableDhcpv6 := false
	if ipv6Addr != "" {
		if enableDhcp && duid != "" {
			enableDhcpv6 = true
		}
	}
	if ipv4Addr != "" {
		if enableDhcp && macAddr != "" && macAddr != MACADDR_ZERO {
			enableDhcpv4 = true
		}
	}

	if ipv4Addr == "" {
		if ipv4cidr != "" {
			ip, _, err := net.ParseCIDR(ipv4cidr)
			if err != nil {
				return nil, fmt.Errorf("cannot parse CIDR value: %s", err.Error())
			}
			if ip.To4() == nil {
				return nil, fmt.Errorf("CIDR value must be an IPv4 CIDR, not an IPv6 one")
			}
			if netView == "" {
				netView = "default"
			}
			newIpAddr := fmt.Sprintf("func:nextavailableip:%s,%s", ipv4cidr, netView)
			recordHostIpAddr := NewHostRecordIpv4Addr(newIpAddr, macAddr, enableDhcpv4, "")
			recordHostIpv4AddrSlice = []HostRecordIpv4Addr{*recordHostIpAddr}
		}
	} else {
		ip := net.ParseIP(ipv4Addr)
		if ip == nil {
			return nil, fmt.Errorf("'IP address for the record is not valid")
		}
		if ip.To4() == nil {
			return nil, fmt.Errorf("IP address must be an IPv4 address, not an IPv6 one")
		}
		recordHostIpAddr := NewHostRecordIpv4Addr(ipv4Addr, macAddr, enableDhcpv4, "")
		recordHostIpv4AddrSlice = []HostRecordIpv4Addr{*recordHostIpAddr}
	}
	if ipv6Addr == "" {
		if ipv6cidr != "" {
			ip, _, err := net.ParseCIDR(ipv6cidr)
			if err != nil {
				return nil, fmt.Errorf("cannot parse CIDR value: %s", err.Error())
			}
			if ip.To4() != nil || ip.To16() == nil {
				return nil, fmt.Errorf("CIDR value must be an IPv6 CIDR, not an IPv4 one")
			}
			if netView == "" {
				netView = "default"
			}
			newIpAddr := fmt.Sprintf("func:nextavailableip:%s,%s", ipv6cidr, netView)
			recordHostIpAddr := NewHostRecordIpv6Addr(newIpAddr, duid, enableDhcpv6, "")
			recordHostIpv6AddrSlice = []HostRecordIpv6Addr{*recordHostIpAddr}
		}
	} else {
		ip := net.ParseIP(ipv6Addr)
		if ip == nil {
			return nil, fmt.Errorf("IP address for the record is not valid")
		}
		if ip.To4() != nil || ip.To16() == nil {
			return nil, fmt.Errorf("IP address must be an IPv6 address, not an IPv4 one")
		}
		recordHostIpAddr := NewHostRecordIpv6Addr(ipv6Addr, duid, enableDhcpv6, "")
		recordHostIpv6AddrSlice = []HostRecordIpv6Addr{*recordHostIpAddr}
	}
	updateHostRecord := NewHostRecord(
		"", name, "", "", recordHostIpv4AddrSlice, recordHostIpv6AddrSlice,
		eas, enabledns, "", "", hostRref, useTtl, ttl, comment, aliases)
	ref, err := objMgr.connector.UpdateObject(updateHostRecord, hostRref)
	if err != nil {
		return nil, err
	}

	updateHostRecord, err = objMgr.GetHostRecordByRef(ref)
	if err != nil {
		return nil, err
	}
	return updateHostRecord, nil
}

func (objMgr *ObjectManager) DeleteHostRecord(ref string) (string, error) {
	return objMgr.connector.DeleteObject(ref)
}
<|endoftext|>


<commit_before>package connections

import (
	"errors"
	"sync"

	"github.com/ivan1993spb/snake-server/game"
)

type ConnectionGroup struct {
	limit   int
	counter int
	mutex   *sync.RWMutex
	game    *game.Game
}

func NewConnectionGroup(connectionLimit int, g *game.Game) (*ConnectionGroup, error) {
	if connectionLimit > 0 {
		return &ConnectionGroup{
			limit: connectionLimit,
			mutex: &sync.RWMutex{},
			game:  g,
		}, nil
	}

	return nil, errors.New("cannot create connection group: invalid connection limit")
}

func (cg *ConnectionGroup) GetLimit() int {
	return cg.limit
}

func (cg *ConnectionGroup) GetCount() int {
	cg.mutex.RLock()
	defer cg.mutex.RUnlock()
	return cg.counter
}

// unsafeIsFull returns true if group is full
func (cg *ConnectionGroup) unsafeIsFull() bool {
	return cg.counter == cg.limit
}

func (cg *ConnectionGroup) IsFull() bool {
	cg.mutex.RLock()
	defer cg.mutex.RUnlock()
	return cg.unsafeIsFull()
}

// unsafeIsEmpty returns true if group is empty
func (cg *ConnectionGroup) unsafeIsEmpty() bool {
	return cg.counter == 0
}

func (cg *ConnectionGroup) IsEmpty() bool {
	cg.mutex.RLock()
	defer cg.mutex.RUnlock()
	return cg.unsafeIsEmpty()
}

func (cg *ConnectionGroup) Run(f func(game *game.Game)) error {
	cg.mutex.Lock()
	if cg.unsafeIsFull() {
		cg.mutex.Unlock()
		return errors.New("add connection to group: group is full")
	}
	cg.counter += 1
	cg.mutex.Unlock()

	f(cg.game)

	cg.mutex.Lock()
	cg.counter -= 1
	cg.mutex.Unlock()
	return nil
}
<commit_msg>Rewrite run method in connection groups<commit_after>package connections

import (
	"errors"
	"sync"

	"github.com/ivan1993spb/snake-server/game"
)

type ConnectionGroup struct {
	limit   int
	counter int
	mutex   *sync.RWMutex
	game    *game.Game
}

func NewConnectionGroup(connectionLimit int, g *game.Game) (*ConnectionGroup, error) {
	if connectionLimit > 0 {
		return &ConnectionGroup{
			limit: connectionLimit,
			mutex: &sync.RWMutex{},
			game:  g,
		}, nil
	}

	return nil, errors.New("cannot create connection group: invalid connection limit")
}

func (cg *ConnectionGroup) GetLimit() int {
	return cg.limit
}

func (cg *ConnectionGroup) GetCount() int {
	cg.mutex.RLock()
	defer cg.mutex.RUnlock()
	return cg.counter
}

// unsafeIsFull returns true if group is full
func (cg *ConnectionGroup) unsafeIsFull() bool {
	return cg.counter == cg.limit
}

func (cg *ConnectionGroup) IsFull() bool {
	cg.mutex.RLock()
	defer cg.mutex.RUnlock()
	return cg.unsafeIsFull()
}

// unsafeIsEmpty returns true if group is empty
func (cg *ConnectionGroup) unsafeIsEmpty() bool {
	return cg.counter == 0
}

func (cg *ConnectionGroup) IsEmpty() bool {
	cg.mutex.RLock()
	defer cg.mutex.RUnlock()
	return cg.unsafeIsEmpty()
}

type ErrRunConnection struct {
	Err error
}

func (e *ErrRunConnection) Error() string {
	return "run connection error: " + e.Err.Error()
}

func (cg *ConnectionGroup) Run(connection *Connection) *ErrRunConnection {
	cg.mutex.Lock()
	if cg.unsafeIsFull() {
		cg.mutex.Unlock()
		return &ErrRunConnection{
			Err: errors.New("group is full"),
		}
	}
	cg.counter += 1
	cg.mutex.Unlock()

	defer func() {
		cg.mutex.Lock()
		cg.counter -= 1
		cg.mutex.Unlock()
	}()

	if err := connection.Run(cg.game); err != nil {
		return &ErrRunConnection{
			Err: err,
		}
	}

	return nil
}
<|endoftext|>


<commit_before>package consumergroup

import (
	"errors"
	"fmt"
	"gopkg.in/Shopify/sarama.v1"
	"sync"
	"time"
)

// OffsetManager is the main interface consumergroup requires to manage offsets of the consumergroup.
type OffsetManager interface {

	// InitializePartition is called when the consumergroup is starting to consume a
	// partition. It should return the last processed offset for this partition. Note:
	// the same partition can be initialized multiple times during a single run of a
	// consumer group due to other consumer instances coming online and offline.
	InitializePartition(topic string, partition int32) (int64, error)

	// MarkAsProcessed tells the offset manager than a certain message has been successfully
	// processed by the consumer, and should be committed. The implementation does not have
	// to store this offset right away, but should return true if it intends to do this at
	// some point.
	//
	// Offsets should generally be increasing if the consumer
	// processes events serially, but this cannot be guaranteed if the consumer does any
	// asynchronous processing. This can be handled in various ways, e.g. by only accepting
	// offsets that are higehr than the offsets seen before for the same partition.
	MarkAsProcessed(topic string, partition int32, offset int64) (bool, error)

	// FinalizePartition is called when the consumergroup is done consuming a
	// partition. In this method, the offset manager can flush any remaining offsets to its
	// backend store. It should return an error if it was not able to commit the offset.
	// Note: it's possible that the consumergroup instance will start to consume the same
	// partition again after this function is called.
	FinalizePartition(topic string, partition int32, lastOffset int64, timeout time.Duration) error

	// Close is called when the consumergroup is shutting down. In normal circumstances, all
	// offsets are committed because FinalizePartition is called for all the running partition
	// consumers. You may want to check for this to be true, and try to commit any outstanding
	// offsets. If this doesn't succeed, it should return an error.
	Close() error
}

var (
	UncleanClose    = errors.New("Not all offsets were committed before shutdown was completed")
	PartitionClosed = errors.New("Tried to mark offset as processed on closed partition.")
)

// OffsetManagerConfig holds configuration setting son how the offset manager should behave.
type OffsetManagerConfig struct {
	CommitInterval time.Duration // Interval between offset flushes to the backend store.
	VerboseLogging bool          // Whether to enable verbose logging.
}

// NewOffsetManagerConfig returns a new OffsetManagerConfig with sane defaults.
func NewOffsetManagerConfig() *OffsetManagerConfig {
	return &OffsetManagerConfig{
		CommitInterval: 10 * time.Second,
	}
}

type (
	topicOffsets    map[int32]*partitionOffsetTracker
	offsetsMap      map[string]topicOffsets
	offsetCommitter func(int64) error
)

type partitionOffsetTracker struct {
	l                      sync.Mutex
	waitingForOffset       int64
	highestProcessedOffset int64
	lastCommittedOffset    int64
	done                   chan struct{}
}

type zookeeperOffsetManager struct {
	config  *OffsetManagerConfig
	l       sync.RWMutex
	offsets offsetsMap
	cg      *ConsumerGroup

	closing, closed chan struct{}
}

// NewZookeeperOffsetManager returns an offset manager that uses Zookeeper
// to store offsets.
func NewZookeeperOffsetManager(cg *ConsumerGroup, config *OffsetManagerConfig) OffsetManager {
	if config == nil {
		config = NewOffsetManagerConfig()
	}

	zom := &zookeeperOffsetManager{
		config:  config,
		cg:      cg,
		offsets: make(offsetsMap),
		closing: make(chan struct{}),
		closed:  make(chan struct{}),
	}

	go zom.offsetCommitter()

	return zom
}

func (zom *zookeeperOffsetManager) InitializePartition(topic string, partition int32) (int64, error) {
	zom.l.Lock()
	defer zom.l.Unlock()

	if zom.offsets[topic] == nil {
		zom.offsets[topic] = make(topicOffsets)
	}

	nextOffset, err := zom.cg.group.FetchOffset(topic, partition)
	if err != nil {
		return 0, err
	}

	zom.offsets[topic][partition] = &partitionOffsetTracker{
		highestProcessedOffset: nextOffset - 1,
		lastCommittedOffset:    nextOffset - 1,
		done:                   make(chan struct{}),
	}

	return nextOffset, nil
}

func (zom *zookeeperOffsetManager) FinalizePartition(topic string, partition int32, lastOffset int64, timeout time.Duration) error {
	zom.l.RLock()
	tracker := zom.offsets[topic][partition]
	zom.l.RUnlock()

	if lastOffset >= 0 {
		if lastOffset-tracker.highestProcessedOffset > 0 {
			zom.cg.Logf("%s/%d :: Last processed offset: %d. Waiting up to %ds for another %d messages to process...", topic, partition, tracker.highestProcessedOffset, timeout/time.Second, lastOffset-tracker.highestProcessedOffset)
			if !tracker.waitForOffset(lastOffset, timeout) {
				return fmt.Errorf("TIMEOUT waiting for offset %d. Last committed offset: %d", lastOffset, tracker.lastCommittedOffset)
			}
		}

		if err := zom.commitOffset(topic, partition, tracker); err != nil {
			return fmt.Errorf("FAILED to commit offset %d to Zookeeper. Last committed offset: %d", tracker.highestProcessedOffset, tracker.lastCommittedOffset)
		}
	}

	zom.l.Lock()
	delete(zom.offsets[topic], partition)
	zom.l.Unlock()

	return nil
}

func (zom *zookeeperOffsetManager) MarkAsProcessed(topic string, partition int32, offset int64) (bool, error) {
	zom.l.RLock()
	defer zom.l.RUnlock()
	if p, ok := zom.offsets[topic][partition]; ok {
		result, err := p.markAsProcessed(offset)
		if err != nil {
			zom.cg.Logf("FAILED to mark offset %d for %s/%d as processed!", offset, topic, partition)
		}
		return result, err
	} else {
		return false, fmt.Errorf("no partition %d for topic %v", partition, topic)
	}
}

func (zom *zookeeperOffsetManager) Close() error {
	close(zom.closing)
	<-zom.closed

	zom.l.Lock()
	defer zom.l.Unlock()

	var closeError error
	for _, partitionOffsets := range zom.offsets {
		if len(partitionOffsets) > 0 {
			closeError = UncleanClose
		}
	}

	return closeError
}

func (zom *zookeeperOffsetManager) offsetCommitter() {
	commitTicker := time.NewTicker(zom.config.CommitInterval)
	defer commitTicker.Stop()

	for {
		select {
		case <-zom.closing:
			close(zom.closed)
			return
		case <-commitTicker.C:
			if err := zom.commitOffsets(); err != nil {
				zom.cg.errors <- &sarama.ConsumerError{
					Topic:     "",
					Partition: -1,
					Err:       err,
				}
			}
		}
	}
}

func (zom *zookeeperOffsetManager) commitOffsets() error {
	zom.l.RLock()
	defer zom.l.RUnlock()

	var returnErr error
	for topic, partitionOffsets := range zom.offsets {
		for partition, offsetTracker := range partitionOffsets {
			err := zom.commitOffset(topic, partition, offsetTracker)
			switch err {
			case nil:
				// noop
			default:
				returnErr = err
			}
		}
	}
	return returnErr
}

func (zom *zookeeperOffsetManager) commitOffset(topic string, partition int32, tracker *partitionOffsetTracker) error {
	err := tracker.commit(func(offset int64) error {
		if offset >= 0 {
			return zom.cg.group.CommitOffset(topic, partition, offset+1)
		} else {
			return nil
		}
	})

	if err != nil {
		zom.cg.Logf("FAILED to commit offset %d for %s/%d!", tracker.highestProcessedOffset, topic, partition)
	} else if zom.config.VerboseLogging {
		zom.cg.Logf("Committed offset %d for %s/%d!", tracker.lastCommittedOffset, topic, partition)
	}

	return err
}

// MarkAsProcessed marks the provided offset as highest processed offset if
// it's higehr than any previous offset it has received.
func (pot *partitionOffsetTracker) markAsProcessed(offset int64) (bool, error) {
	if pot == nil {
		return false, PartitionClosed
	}
	pot.l.Lock()
	defer pot.l.Unlock()
	if offset > pot.highestProcessedOffset {
		pot.highestProcessedOffset = offset
		if pot.waitingForOffset == pot.highestProcessedOffset {
			close(pot.done)
		}
		return true, nil
	} else {
		return false, nil
	}
}

// Commit calls a committer function if the highest processed offset is out
// of sync with the last committed offset.
func (pot *partitionOffsetTracker) commit(committer offsetCommitter) error {
	pot.l.Lock()
	defer pot.l.Unlock()

	if pot.highestProcessedOffset > pot.lastCommittedOffset {
		if err := committer(pot.highestProcessedOffset); err != nil {
			return err
		}
		pot.lastCommittedOffset = pot.highestProcessedOffset
		return nil
	} else {
		return nil
	}
}

func (pot *partitionOffsetTracker) waitForOffset(offset int64, timeout time.Duration) bool {
	pot.l.Lock()
	if offset > pot.highestProcessedOffset {
		pot.waitingForOffset = offset
		pot.l.Unlock()
		select {
		case <-pot.done:
			return true
		case <-time.After(timeout):
			return false
		}
	} else {
		pot.l.Unlock()
		return true
	}
}
<commit_msg>Updates to github sarama package<commit_after>package consumergroup

import (
	"errors"
	"fmt"
	"github.com/Shopify/sarama"
	"sync"
	"time"
)

// OffsetManager is the main interface consumergroup requires to manage offsets of the consumergroup.
type OffsetManager interface {

	// InitializePartition is called when the consumergroup is starting to consume a
	// partition. It should return the last processed offset for this partition. Note:
	// the same partition can be initialized multiple times during a single run of a
	// consumer group due to other consumer instances coming online and offline.
	InitializePartition(topic string, partition int32) (int64, error)

	// MarkAsProcessed tells the offset manager than a certain message has been successfully
	// processed by the consumer, and should be committed. The implementation does not have
	// to store this offset right away, but should return true if it intends to do this at
	// some point.
	//
	// Offsets should generally be increasing if the consumer
	// processes events serially, but this cannot be guaranteed if the consumer does any
	// asynchronous processing. This can be handled in various ways, e.g. by only accepting
	// offsets that are higehr than the offsets seen before for the same partition.
	MarkAsProcessed(topic string, partition int32, offset int64) (bool, error)

	// FinalizePartition is called when the consumergroup is done consuming a
	// partition. In this method, the offset manager can flush any remaining offsets to its
	// backend store. It should return an error if it was not able to commit the offset.
	// Note: it's possible that the consumergroup instance will start to consume the same
	// partition again after this function is called.
	FinalizePartition(topic string, partition int32, lastOffset int64, timeout time.Duration) error

	// Close is called when the consumergroup is shutting down. In normal circumstances, all
	// offsets are committed because FinalizePartition is called for all the running partition
	// consumers. You may want to check for this to be true, and try to commit any outstanding
	// offsets. If this doesn't succeed, it should return an error.
	Close() error
}

var (
	UncleanClose    = errors.New("Not all offsets were committed before shutdown was completed")
	PartitionClosed = errors.New("Tried to mark offset as processed on closed partition.")
)

// OffsetManagerConfig holds configuration setting son how the offset manager should behave.
type OffsetManagerConfig struct {
	CommitInterval time.Duration // Interval between offset flushes to the backend store.
	VerboseLogging bool          // Whether to enable verbose logging.
}

// NewOffsetManagerConfig returns a new OffsetManagerConfig with sane defaults.
func NewOffsetManagerConfig() *OffsetManagerConfig {
	return &OffsetManagerConfig{
		CommitInterval: 10 * time.Second,
	}
}

type (
	topicOffsets    map[int32]*partitionOffsetTracker
	offsetsMap      map[string]topicOffsets
	offsetCommitter func(int64) error
)

type partitionOffsetTracker struct {
	l                      sync.Mutex
	waitingForOffset       int64
	highestProcessedOffset int64
	lastCommittedOffset    int64
	done                   chan struct{}
}

type zookeeperOffsetManager struct {
	config  *OffsetManagerConfig
	l       sync.RWMutex
	offsets offsetsMap
	cg      *ConsumerGroup

	closing, closed chan struct{}
}

// NewZookeeperOffsetManager returns an offset manager that uses Zookeeper
// to store offsets.
func NewZookeeperOffsetManager(cg *ConsumerGroup, config *OffsetManagerConfig) OffsetManager {
	if config == nil {
		config = NewOffsetManagerConfig()
	}

	zom := &zookeeperOffsetManager{
		config:  config,
		cg:      cg,
		offsets: make(offsetsMap),
		closing: make(chan struct{}),
		closed:  make(chan struct{}),
	}

	go zom.offsetCommitter()

	return zom
}

func (zom *zookeeperOffsetManager) InitializePartition(topic string, partition int32) (int64, error) {
	zom.l.Lock()
	defer zom.l.Unlock()

	if zom.offsets[topic] == nil {
		zom.offsets[topic] = make(topicOffsets)
	}

	nextOffset, err := zom.cg.group.FetchOffset(topic, partition)
	if err != nil {
		return 0, err
	}

	zom.offsets[topic][partition] = &partitionOffsetTracker{
		highestProcessedOffset: nextOffset - 1,
		lastCommittedOffset:    nextOffset - 1,
		done:                   make(chan struct{}),
	}

	return nextOffset, nil
}

func (zom *zookeeperOffsetManager) FinalizePartition(topic string, partition int32, lastOffset int64, timeout time.Duration) error {
	zom.l.RLock()
	tracker := zom.offsets[topic][partition]
	zom.l.RUnlock()

	if lastOffset >= 0 {
		if lastOffset-tracker.highestProcessedOffset > 0 {
			zom.cg.Logf("%s/%d :: Last processed offset: %d. Waiting up to %ds for another %d messages to process...", topic, partition, tracker.highestProcessedOffset, timeout/time.Second, lastOffset-tracker.highestProcessedOffset)
			if !tracker.waitForOffset(lastOffset, timeout) {
				return fmt.Errorf("TIMEOUT waiting for offset %d. Last committed offset: %d", lastOffset, tracker.lastCommittedOffset)
			}
		}

		if err := zom.commitOffset(topic, partition, tracker); err != nil {
			return fmt.Errorf("FAILED to commit offset %d to Zookeeper. Last committed offset: %d", tracker.highestProcessedOffset, tracker.lastCommittedOffset)
		}
	}

	zom.l.Lock()
	delete(zom.offsets[topic], partition)
	zom.l.Unlock()

	return nil
}

func (zom *zookeeperOffsetManager) MarkAsProcessed(topic string, partition int32, offset int64) (bool, error) {
	zom.l.RLock()
	defer zom.l.RUnlock()
	if p, ok := zom.offsets[topic][partition]; ok {
		result, err := p.markAsProcessed(offset)
		if err != nil {
			zom.cg.Logf("FAILED to mark offset %d for %s/%d as processed!", offset, topic, partition)
		}
		return result, err
	} else {
		return false, fmt.Errorf("no partition %d for topic %v", partition, topic)
	}
}

func (zom *zookeeperOffsetManager) Close() error {
	close(zom.closing)
	<-zom.closed

	zom.l.Lock()
	defer zom.l.Unlock()

	var closeError error
	for _, partitionOffsets := range zom.offsets {
		if len(partitionOffsets) > 0 {
			closeError = UncleanClose
		}
	}

	return closeError
}

func (zom *zookeeperOffsetManager) offsetCommitter() {
	commitTicker := time.NewTicker(zom.config.CommitInterval)
	defer commitTicker.Stop()

	for {
		select {
		case <-zom.closing:
			close(zom.closed)
			return
		case <-commitTicker.C:
			if err := zom.commitOffsets(); err != nil {
				zom.cg.errors <- &sarama.ConsumerError{
					Topic:     "",
					Partition: -1,
					Err:       err,
				}
			}
		}
	}
}

func (zom *zookeeperOffsetManager) commitOffsets() error {
	zom.l.RLock()
	defer zom.l.RUnlock()

	var returnErr error
	for topic, partitionOffsets := range zom.offsets {
		for partition, offsetTracker := range partitionOffsets {
			err := zom.commitOffset(topic, partition, offsetTracker)
			switch err {
			case nil:
				// noop
			default:
				returnErr = err
			}
		}
	}
	return returnErr
}

func (zom *zookeeperOffsetManager) commitOffset(topic string, partition int32, tracker *partitionOffsetTracker) error {
	err := tracker.commit(func(offset int64) error {
		if offset >= 0 {
			return zom.cg.group.CommitOffset(topic, partition, offset+1)
		} else {
			return nil
		}
	})

	if err != nil {
		zom.cg.Logf("FAILED to commit offset %d for %s/%d!", tracker.highestProcessedOffset, topic, partition)
	} else if zom.config.VerboseLogging {
		zom.cg.Logf("Committed offset %d for %s/%d!", tracker.lastCommittedOffset, topic, partition)
	}

	return err
}

// MarkAsProcessed marks the provided offset as highest processed offset if
// it's higehr than any previous offset it has received.
func (pot *partitionOffsetTracker) markAsProcessed(offset int64) (bool, error) {
	if pot == nil {
		return false, PartitionClosed
	}
	pot.l.Lock()
	defer pot.l.Unlock()
	if offset > pot.highestProcessedOffset {
		pot.highestProcessedOffset = offset
		if pot.waitingForOffset == pot.highestProcessedOffset {
			close(pot.done)
		}
		return true, nil
	} else {
		return false, nil
	}
}

// Commit calls a committer function if the highest processed offset is out
// of sync with the last committed offset.
func (pot *partitionOffsetTracker) commit(committer offsetCommitter) error {
	pot.l.Lock()
	defer pot.l.Unlock()

	if pot.highestProcessedOffset > pot.lastCommittedOffset {
		if err := committer(pot.highestProcessedOffset); err != nil {
			return err
		}
		pot.lastCommittedOffset = pot.highestProcessedOffset
		return nil
	} else {
		return nil
	}
}

func (pot *partitionOffsetTracker) waitForOffset(offset int64, timeout time.Duration) bool {
	pot.l.Lock()
	if offset > pot.highestProcessedOffset {
		pot.waitingForOffset = offset
		pot.l.Unlock()
		select {
		case <-pot.done:
			return true
		case <-time.After(timeout):
			return false
		}
	} else {
		pot.l.Unlock()
		return true
	}
}
<|endoftext|>


<commit_before>// Copyright © 2016 The Things Network
// Use of this source code is governed by the MIT license that can be found in the LICENSE file.

package amqp

import (
	"fmt"
	"io"
	"time"

	"github.com/TheThingsNetwork/ttn/core/types"

	AMQP "github.com/streadway/amqp"
)

// Publisher holds an AMQP channel and can publish on uplink from devices
type Publisher interface {
	Open() error
	io.Closer
	PublishUplink(payload types.UplinkMessage) error
}

// DefaultPublisher is the default AMQP publisher
type DefaultPublisher struct {
	ctx      Logger
	conn     *AMQP.Connection
	channel  *AMQP.Channel
	exchange string
}

// Open opens an AMQP channel for publishing
func (p *DefaultPublisher) Open() error {
	channel, err := p.conn.Channel()
	if err != nil {
		return fmt.Errorf("Could not open AMQP channel (%s).", err)
	}

	if err = channel.ExchangeDeclare(p.exchange, "topic", true, false, false, false, nil); err != nil {
		channel.Close()
		return fmt.Errorf("Could not declare AMQP exchange (%s).", err)
	}

	p.channel = channel
	return nil
}

// Close closes the AMQP channel
func (p *DefaultPublisher) Close() error {
	if p.channel == nil {
		return nil
	}
	return p.channel.Close()
}

func (p *DefaultPublisher) do(action func() error) error {
	err := action()
	if err == nil {
		return nil
	}
	if err == AMQP.ErrClosed {
		err = p.Open()
		if err != nil {
			return err
		}
		err = action()
	}
	return err
}

func (p *DefaultPublisher) publish(key string, msg []byte, timestamp time.Time) error {
	return p.do(func() error {
		return p.channel.Publish(p.exchange, key, false, false, AMQP.Publishing{
			ContentType:  "application/json",
			DeliveryMode: AMQP.Persistent,
			Timestamp:    timestamp,
			Body:         msg,
		})
	})
}
<commit_msg>Code cleanup<commit_after>// Copyright © 2016 The Things Network
// Use of this source code is governed by the MIT license that can be found in the LICENSE file.

package amqp

import (
	"fmt"
	"io"
	"time"

	"github.com/TheThingsNetwork/ttn/core/types"
	AMQP "github.com/streadway/amqp"
)

// Publisher holds an AMQP channel and can publish on uplink from devices
type Publisher interface {
	Open() error
	io.Closer
	PublishUplink(payload types.UplinkMessage) error
}

// DefaultPublisher is the default AMQP publisher
type DefaultPublisher struct {
	ctx      Logger
	conn     *AMQP.Connection
	channel  *AMQP.Channel
	exchange string
}

// Open opens an AMQP channel for publishing
func (p *DefaultPublisher) Open() error {
	channel, err := p.conn.Channel()
	if err != nil {
		return fmt.Errorf("Could not open AMQP channel (%s).", err)
	}

	if err = channel.ExchangeDeclare(p.exchange, "topic", true, false, false, false, nil); err != nil {
		channel.Close()
		return fmt.Errorf("Could not declare AMQP exchange (%s).", err)
	}

	p.channel = channel
	return nil
}

// Close closes the AMQP channel
func (p *DefaultPublisher) Close() error {
	if p.channel == nil {
		return nil
	}
	return p.channel.Close()
}

func (p *DefaultPublisher) do(action func() error) error {
	err := action()
	if err == nil {
		return nil
	}
	if err == AMQP.ErrClosed {
		err = p.Open()
		if err != nil {
			return err
		}
		err = action()
	}
	return err
}

func (p *DefaultPublisher) publish(key string, msg []byte, timestamp time.Time) error {
	return p.do(func() error {
		return p.channel.Publish(p.exchange, key, false, false, AMQP.Publishing{
			ContentType:  "application/json",
			DeliveryMode: AMQP.Persistent,
			Timestamp:    timestamp,
			Body:         msg,
		})
	})
}
<|endoftext|>


<commit_before>package main

import (
	"encoding/json"
)

type Event map[string]interface{}

type CommitList []Event

// Information to translate from GitLab to GitHub format.
var eventTranslate = map[string]string{
	"total_commits_count": "size",
	"repository":          "repo",
	"after":               "head",
}

var commitTranslate = map[string]string{
	"id": "sha",
}

func (e Event) Commits() CommitList {
	generic, ok := e["commits"]
	if !ok {
		return nil
	}

	c, ok := generic.(CommitList)
	if !ok {
		return nil
	}

	return c
}

// Normalize makes GitLab events look like GitHub events.
func (e Event) normalize() {
	// Fortunately there's very little difference, at least in the push format.

	for gitlabName, githubName := range eventTranslate {
		value, ok := e[gitlabName]
		if ok {
			e[githubName] = value
		}
	}

	commits := e.Commits()
	for _, c := range commits {
		for gitlabName, githubName := range commitTranslate {
			value, ok := c[gitlabName]
			if ok {
				c[githubName] = value
			}
		}
	}

}

// Create a new event from the given JSON. If the event type is blank,
// this function will try to figure it out. Generally, GitHub events will
// present a value for this in the HTTP Request, and GitLab events place
// the event type in the JSON.
func NewEvent(jsonData []byte, eventName string) Event {
	e := Event{}
	json.Unmarshal(jsonData, e)

	if payload, ok := e["payload"].(Event); ok {
		// For GitHub events, export all payload fields into the event scope.
		for key, value := range payload {
			e[key] = value
		}
	}

	e.normalize()

	if eventName == "" {
		var gitlabType string
		gitlabType, ok := e["object_kind"].(string)
		if !ok {
			// Push events look completely different from the other events and
			// don't have an explicit type.
			gitlabType = "push"
		}
		e["type"] = gitlabType
	} else {
		e["type"] = eventName
	}
	return e
}
<commit_msg>Extra handling of gitlab non-push events<commit_after>package main

import (
	"encoding/json"
)

type Event map[string]interface{}

type CommitList []Event

// Information to translate from GitLab to GitHub format.
var eventTranslate = map[string]string{
	"total_commits_count": "size",
	"repository":          "repo",
	"after":               "head",
}

var commitTranslate = map[string]string{
	"id": "sha",
}

func (e Event) Commits() CommitList {
	generic, ok := e["commits"]
	if !ok {
		return nil
	}

	c, ok := generic.(CommitList)
	if !ok {
		return nil
	}

	return c
}

// Normalize makes GitLab events look like GitHub events.
func (e Event) normalize() {
	// Fortunately there's very little difference, at least in the push format.

	for gitlabName, githubName := range eventTranslate {
		value, ok := e[gitlabName]
		if ok {
			e[githubName] = value
		}
	}

	commits := e.Commits()
	for _, c := range commits {
		for gitlabName, githubName := range commitTranslate {
			value, ok := c[gitlabName]
			if ok {
				c[githubName] = value
			}
		}
	}

}

// Create a new event from the given JSON. If the event type is blank,
// this function will try to figure it out. Generally, GitHub events will
// present a value for this in the HTTP Request, and GitLab events place
// the event type in the JSON.
func NewEvent(jsonData []byte, eventName string) Event {
	e := Event{}
	json.Unmarshal(jsonData, e)

	if payload, ok := e["payload"].(Event); ok {
		// For GitHub events, export all payload fields into the event scope.
		for key, value := range payload {
			e[key] = value
		}
	}

	if payload, ok := e["object_attributes"].(Event); ok {
		// For GitLab events, export all object_attributes fields into the event scope.
		for key, value := range payload {
			e[key] = value
		}
	}

	e.normalize()

	if eventName == "" {
		var gitlabType string
		gitlabType, ok := e["object_kind"].(string)
		if !ok {
			// Push events look completely different from the other events and
			// don't have an explicit type.
			gitlabType = "push"
		}
		e["type"] = gitlabType
	} else {
		e["type"] = eventName
	}
	return e
}
<|endoftext|>


<commit_before>package create

import (
	"fmt"

	"github.com/giantswarm/awstpr"
	microerror "github.com/giantswarm/microkit/error"

	awsutil "github.com/giantswarm/aws-operator/client/aws"
	"github.com/giantswarm/aws-operator/resources"
	awsresources "github.com/giantswarm/aws-operator/resources/aws"
)

type securityGroupInput struct {
	Clients   awsutil.Clients
	GroupName string
	VPCID     string
}

type rulesInput struct {
	Cluster                awstpr.CustomObject
	Rules                  []awsresources.SecurityGroupRule
	MastersSecurityGroupID string
	WorkersSecurityGroupID string
	IngressSecurityGroupID string
}

const (
	calicoBGPNetworkPort = 179
	httpPort             = 80
	httpsPort            = 443
	sshPort              = 22

	defaultCIDR = "0.0.0.0/0"
)

func (s *Service) createSecurityGroup(input securityGroupInput) (*awsresources.SecurityGroup, error) {
	securityGroup := &awsresources.SecurityGroup{
		Description: input.GroupName,
		GroupName:   input.GroupName,
		VpcID:       input.VPCID,
		AWSEntity:   awsresources.AWSEntity{Clients: input.Clients},
	}
	securityGroupCreated, err := securityGroup.CreateIfNotExists()
	if err != nil {
		return nil, microerror.MaskAny(err)
	}
	if securityGroupCreated {
		s.logger.Log("info", fmt.Sprintf("created security group '%s'", input.GroupName))
	} else {
		s.logger.Log("info", fmt.Sprintf("security group '%s' already exists, reusing", input.GroupName))
	}

	return securityGroup, nil
}

func (s *Service) deleteSecurityGroup(input securityGroupInput) error {
	var securityGroup resources.ResourceWithID
	securityGroup = &awsresources.SecurityGroup{
		Description: input.GroupName,
		GroupName:   input.GroupName,
		AWSEntity:   awsresources.AWSEntity{Clients: input.Clients},
	}
	if err := securityGroup.Delete(); err != nil {
		return microerror.MaskAny(err)
	} else {
		s.logger.Log("info", fmt.Sprintf("deleted security group '%s'", input.GroupName))
	}

	return nil
}

// masterRules returns the rules for the masters security group.
func (ri rulesInput) masterRules() []awsresources.SecurityGroupRule {
	return []awsresources.SecurityGroupRule{
		{
			Port:       ri.Cluster.Spec.Cluster.Kubernetes.API.SecurePort,
			SourceCIDR: defaultCIDR,
		},
		{
			Port:       ri.Cluster.Spec.Cluster.Etcd.Port,
			SourceCIDR: defaultCIDR,
		},
		{
			Port:       sshPort,
			SourceCIDR: defaultCIDR,
		},
		{
			Port:       calicoBGPNetworkPort,
			SourceCIDR: defaultCIDR,
		},
	}
}

// workerRules returns the rules for the workers security group.
func (ri rulesInput) workerRules() []awsresources.SecurityGroupRule {
	return []awsresources.SecurityGroupRule{
		{
			Port:       ri.Cluster.Spec.Cluster.Kubernetes.IngressController.SecurePort,
			SourceCIDR: defaultCIDR,
		},
		{
			Port:       ri.Cluster.Spec.Cluster.Kubernetes.IngressController.InsecurePort,
			SourceCIDR: defaultCIDR,
		},
		{
			Port:       ri.Cluster.Spec.Cluster.Kubernetes.Kubelet.Port,
			SourceCIDR: defaultCIDR,
		},
		{
			Port:       sshPort,
			SourceCIDR: defaultCIDR,
		},
		{
			Port:       calicoBGPNetworkPort,
			SourceCIDR: defaultCIDR,
		},
	}
}

// ingressRules returns the rules for the ingress ELB security group.
func (ri rulesInput) ingressRules() []awsresources.SecurityGroupRule {
	return []awsresources.SecurityGroupRule{
		{
			Port:       httpPort,
			SourceCIDR: defaultCIDR,
		},
		{
			Port:       httpsPort,
			SourceCIDR: defaultCIDR,
		},
	}
}

func securityGroupName(clusterName string, groupName string) string {
	return fmt.Sprintf("%s-%s", clusterName, groupName)
}
<commit_msg>open read only kubelet port for heapster<commit_after>package create

import (
	"fmt"

	"github.com/giantswarm/awstpr"
	microerror "github.com/giantswarm/microkit/error"

	awsutil "github.com/giantswarm/aws-operator/client/aws"
	"github.com/giantswarm/aws-operator/resources"
	awsresources "github.com/giantswarm/aws-operator/resources/aws"
)

type securityGroupInput struct {
	Clients   awsutil.Clients
	GroupName string
	VPCID     string
}

type rulesInput struct {
	Cluster                awstpr.CustomObject
	Rules                  []awsresources.SecurityGroupRule
	MastersSecurityGroupID string
	WorkersSecurityGroupID string
	IngressSecurityGroupID string
}

const (
	calicoBGPNetworkPort = 179
	// This port is required in our current kubernetes/heapster setup, but will become unnecessary
	// once we upgrade to kubernetes 1.6 and heapster 1.3 with apiserver deployment.
	readOnlyKubeletPort = 10255
	httpPort            = 80
	httpsPort           = 443
	sshPort             = 22

	defaultCIDR = "0.0.0.0/0"
)

func (s *Service) createSecurityGroup(input securityGroupInput) (*awsresources.SecurityGroup, error) {
	securityGroup := &awsresources.SecurityGroup{
		Description: input.GroupName,
		GroupName:   input.GroupName,
		VpcID:       input.VPCID,
		AWSEntity:   awsresources.AWSEntity{Clients: input.Clients},
	}
	securityGroupCreated, err := securityGroup.CreateIfNotExists()
	if err != nil {
		return nil, microerror.MaskAny(err)
	}
	if securityGroupCreated {
		s.logger.Log("info", fmt.Sprintf("created security group '%s'", input.GroupName))
	} else {
		s.logger.Log("info", fmt.Sprintf("security group '%s' already exists, reusing", input.GroupName))
	}

	return securityGroup, nil
}

func (s *Service) deleteSecurityGroup(input securityGroupInput) error {
	var securityGroup resources.ResourceWithID
	securityGroup = &awsresources.SecurityGroup{
		Description: input.GroupName,
		GroupName:   input.GroupName,
		AWSEntity:   awsresources.AWSEntity{Clients: input.Clients},
	}
	if err := securityGroup.Delete(); err != nil {
		return microerror.MaskAny(err)
	} else {
		s.logger.Log("info", fmt.Sprintf("deleted security group '%s'", input.GroupName))
	}

	return nil
}

// masterRules returns the rules for the masters security group.
func (ri rulesInput) masterRules() []awsresources.SecurityGroupRule {
	return []awsresources.SecurityGroupRule{
		{
			Port:       ri.Cluster.Spec.Cluster.Kubernetes.API.SecurePort,
			SourceCIDR: defaultCIDR,
		},
		{
			Port:       ri.Cluster.Spec.Cluster.Etcd.Port,
			SourceCIDR: defaultCIDR,
		},
		{
			Port:       sshPort,
			SourceCIDR: defaultCIDR,
		},
		{
			Port:       calicoBGPNetworkPort,
			SourceCIDR: defaultCIDR,
		},
	}
}

// workerRules returns the rules for the workers security group.
func (ri rulesInput) workerRules() []awsresources.SecurityGroupRule {
	return []awsresources.SecurityGroupRule{
		{
			Port:       ri.Cluster.Spec.Cluster.Kubernetes.IngressController.SecurePort,
			SourceCIDR: defaultCIDR,
		},
		{
			Port:       ri.Cluster.Spec.Cluster.Kubernetes.IngressController.InsecurePort,
			SourceCIDR: defaultCIDR,
		},
		{
			Port:       ri.Cluster.Spec.Cluster.Kubernetes.Kubelet.Port,
			SourceCIDR: defaultCIDR,
		},
		{
			Port:       readOnlyKubeletPort,
			SourceCIDR: defaultCIDR,
		},
		{
			Port:       sshPort,
			SourceCIDR: defaultCIDR,
		},
		{
			Port:       calicoBGPNetworkPort,
			SourceCIDR: defaultCIDR,
		},
	}
}

// ingressRules returns the rules for the ingress ELB security group.
func (ri rulesInput) ingressRules() []awsresources.SecurityGroupRule {
	return []awsresources.SecurityGroupRule{
		{
			Port:       httpPort,
			SourceCIDR: defaultCIDR,
		},
		{
			Port:       httpsPort,
			SourceCIDR: defaultCIDR,
		},
	}
}

func securityGroupName(clusterName string, groupName string) string {
	return fmt.Sprintf("%s-%s", clusterName, groupName)
}
<|endoftext|>


<commit_before>package s3manager

import (
	"fmt"
	"io"
	"net/http"
	"strconv"
	"strings"
	"sync"

	"github.com/aws/aws-sdk-go/aws/awserr"
	"github.com/aws/aws-sdk-go/aws/awsutil"
	"github.com/aws/aws-sdk-go/aws/client"
	"github.com/aws/aws-sdk-go/aws/request"
	"github.com/aws/aws-sdk-go/service/s3"
	"github.com/aws/aws-sdk-go/service/s3/s3iface"
)

// DefaultDownloadPartSize is the default range of bytes to get at a time when
// using Download().
const DefaultDownloadPartSize = 1024 * 1024 * 5

// DefaultDownloadConcurrency is the default number of goroutines to spin up
// when using Download().
const DefaultDownloadConcurrency = 5

// The Downloader structure that calls Download(). It is safe to call Download()
// on this structure for multiple objects and across concurrent goroutines.
// Mutating the Downloader's properties is not safe to be done concurrently.
type Downloader struct {
	// The buffer size (in bytes) to use when buffering data into chunks and
	// sending them as parts to S3. The minimum allowed part size is 5MB, and
	// if this value is set to zero, the DefaultPartSize value will be used.
	PartSize int64

	// The number of goroutines to spin up in parallel when sending parts.
	// If this is set to zero, the DefaultConcurrency value will be used.
	Concurrency int

	// An S3 client to use when performing downloads.
	S3 s3iface.S3API
}

// NewDownloader creates a new Downloader instance to downloads objects from
// S3 in concurrent chunks. Pass in additional functional options  to customize
// the downloader behavior. Requires a client.ConfigProvider in order to create
// a S3 service client. The session.Session satisfies the client.ConfigProvider
// interface.
//
// Example:
//     // The session the S3 Downloader will use
//     sess := session.New()
//
//     // Create a downloader with the session and default options
//     downloader := s3manager.NewDownloader(sess)
//
//     // Create a downloader with the session and custom options
//     downloader := s3manager.NewDownloader(sess, func(d *s3manager.Uploader) {
//          d.PartSize = 64 * 1024 * 1024 // 64MB per part
//     })
func NewDownloader(c client.ConfigProvider, options ...func(*Downloader)) *Downloader {
	d := &Downloader{
		S3:          s3.New(c),
		PartSize:    DefaultDownloadPartSize,
		Concurrency: DefaultDownloadConcurrency,
	}
	for _, option := range options {
		option(d)
	}

	return d
}

// NewDownloaderWithClient creates a new Downloader instance to downloads
// objects from S3 in concurrent chunks. Pass in additional functional
// options to customize the downloader behavior. Requires a S3 service client
// to make S3 API calls.
//
// Example:
//     // The S3 client the S3 Downloader will use
//     s3Svc := s3.new(session.New())
//
//     // Create a downloader with the s3 client and default options
//     downloader := s3manager.NewDownloaderWithClient(s3Svc)
//
//     // Create a downloader with the s3 client and custom options
//     downloader := s3manager.NewDownloaderWithClient(s3Svc, func(d *s3manager.Uploader) {
//          d.PartSize = 64 * 1024 * 1024 // 64MB per part
//     })
func NewDownloaderWithClient(svc s3iface.S3API, options ...func(*Downloader)) *Downloader {
	d := &Downloader{
		S3:          svc,
		PartSize:    DefaultDownloadPartSize,
		Concurrency: DefaultDownloadConcurrency,
	}
	for _, option := range options {
		option(d)
	}

	return d
}

// Download downloads an object in S3 and writes the payload into w using
// concurrent GET requests.
//
// Additional functional options can be provided to configure the individual
// upload. These options are copies of the Uploader instance Upload is called from.
// Modifying the options will not impact the original Uploader instance.
//
// It is safe to call this method concurrently across goroutines.
//
// The w io.WriterAt can be satisfied by an os.File to do multipart concurrent
// downloads, or in memory []byte wrapper using aws.WriteAtBuffer.
func (d Downloader) Download(w io.WriterAt, input *s3.GetObjectInput, options ...func(*Downloader)) (n int64, err error) {
	impl := downloader{w: w, in: input, ctx: d}

	for _, option := range options {
		option(&impl.ctx)
	}

	return impl.download()
}

// downloader is the implementation structure used internally by Downloader.
type downloader struct {
	ctx Downloader

	in *s3.GetObjectInput
	w  io.WriterAt

	wg sync.WaitGroup
	m  sync.Mutex

	pos        int64
	totalBytes int64
	written    int64
	err        error
}

// init initializes the downloader with default options.
func (d *downloader) init() {
	d.totalBytes = -1

	if d.ctx.Concurrency == 0 {
		d.ctx.Concurrency = DefaultDownloadConcurrency
	}

	if d.ctx.PartSize == 0 {
		d.ctx.PartSize = DefaultDownloadPartSize
	}
}

// download performs the implementation of the object download across ranged
// GETs.
func (d *downloader) download() (n int64, err error) {
	d.init()

	// Spin off first worker to check additional header information
	d.notThreadSafeGetChunk()

	if total := d.getTotalBytes(); total >= 0 {
		// Spin up workers
		ch := make(chan dlchunk, d.ctx.Concurrency)

		for i := 0; i < d.ctx.Concurrency; i++ {
			d.wg.Add(1)
			go d.downloadPart(ch)
		}

		// Assign work
		for d.getErr() == nil {
			if d.pos >= total {
				break // We're finished queueing chunks
			}

			// Queue the next range of bytes to read.
			ch <- dlchunk{w: d.w, start: d.pos, size: d.ctx.PartSize}
			d.pos += d.ctx.PartSize
		}

		// Wait for completion
		close(ch)
		d.wg.Wait()
	} else {
		// Checking if we read anything new
		for d.err == nil {
			d.notThreadSafeGetChunk()
		}

		// We expect a 416 error letting us know we are done downloading the
		// total bytes. Since we do not know the content's length, this will
		// keep grabbing chunks of data until the range of bytes specified in
		// the request is out of range of the content. Once, this happens, a
		// 416 should occur.
		e, ok := d.err.(awserr.RequestFailure)
		if ok && e.StatusCode() == http.StatusRequestedRangeNotSatisfiable {
			d.err = nil
		}
	}

	// Return error
	return d.written, d.err
}

// downloadPart is an individual goroutine worker reading from the ch channel
// and performing a GetObject request on the data with a given byte range.
//
// If this is the first worker, this operation also resolves the total number
// of bytes to be read so that the worker manager knows when it is finished.
func (d *downloader) downloadPart(ch chan dlchunk) {
	defer d.wg.Done()
	for {
		chunk, ok := <-ch
		if !ok {
			break
		}
		d.downloadChunk(chunk)
	}
}

// notThreadSafeGetChunk grabs a chunk of data from the body.
// Not thread safe. Should only used when grabbing data on a single thread.
func (d *downloader) notThreadSafeGetChunk() {
	chunk := dlchunk{w: d.w, start: d.pos, size: d.ctx.PartSize}
	d.pos += d.ctx.PartSize
	d.downloadChunk(chunk)
}

// downloadChunk downloads the chunk froom s3
func (d *downloader) downloadChunk(chunk dlchunk) {
	if d.getErr() != nil {
		return
	}
	// Get the next byte range of data
	in := &s3.GetObjectInput{}
	awsutil.Copy(in, d.in)
	rng := fmt.Sprintf("bytes=%d-%d",
		chunk.start, chunk.start+chunk.size-1)
	in.Range = &rng

	req, resp := d.ctx.S3.GetObjectRequest(in)
	req.Handlers.Build.PushBack(request.MakeAddToUserAgentFreeFormHandler("S3Manager"))
	err := req.Send()

	if err != nil {
		d.setErr(err)
	} else {
		d.setTotalBytes(resp) // Set total if not yet set.

		n, err := io.Copy(&chunk, resp.Body)
		resp.Body.Close()

		if err != nil {
			d.setErr(err)
		}
		d.incrWritten(n)
	}
}

// getTotalBytes is a thread-safe getter for retrieving the total byte status.
func (d *downloader) getTotalBytes() int64 {
	d.m.Lock()
	defer d.m.Unlock()

	return d.totalBytes
}

// setTotalBytes is a thread-safe setter for setting the total byte status.
// Will extract the object's total bytes from the Content-Range if the file
// will be chunked, or Content-Length. Content-Length is used when the response
// does not include a Content-Range. Meaning the object was not chunked. This
// occurs when the full file fits within the PartSize directive.
func (d *downloader) setTotalBytes(resp *s3.GetObjectOutput) {
	d.m.Lock()
	defer d.m.Unlock()

	if d.totalBytes >= 0 {
		return
	}

	if resp.ContentRange == nil {
		// ContentRange is nil when the full file contents is provied, and
		// is not chunked. Use ContentLength instead.
		if resp.ContentLength != nil {
			d.totalBytes = *resp.ContentLength
			return
		}
	} else {
		parts := strings.Split(*resp.ContentRange, "/")

		total := int64(-1)
		var err error
		// Checking for whether or not a numbered total exists
		// If one does not exist, we will assume the total to be -1, undefined,
		// and sequentially download each chunk until hitting a 416 error
		totalStr := parts[len(parts)-1]
		if totalStr != "*" {
			total, err = strconv.ParseInt(totalStr, 10, 64)
			if err != nil {
				d.err = err
				return
			}
		}

		d.totalBytes = total
	}
}

func (d *downloader) incrWritten(n int64) {
	d.m.Lock()
	defer d.m.Unlock()

	d.written += n
}

// getErr is a thread-safe getter for the error object
func (d *downloader) getErr() error {
	d.m.Lock()
	defer d.m.Unlock()

	return d.err
}

// setErr is a thread-safe setter for the error object
func (d *downloader) setErr(e error) {
	d.m.Lock()
	defer d.m.Unlock()

	d.err = e
}

// dlchunk represents a single chunk of data to write by the worker routine.
// This structure also implements an io.SectionReader style interface for
// io.WriterAt, effectively making it an io.SectionWriter (which does not
// exist).
type dlchunk struct {
	w     io.WriterAt
	start int64
	size  int64
	cur   int64
}

// Write wraps io.WriterAt for the dlchunk, writing from the dlchunk's start
// position to its end (or EOF).
func (c *dlchunk) Write(p []byte) (n int, err error) {
	if c.cur >= c.size {
		return 0, io.EOF
	}

	n, err = c.w.WriteAt(p, c.start+c.cur)
	c.cur += int64(n)

	return
}
<commit_msg>Changing back the name<commit_after>package s3manager

import (
	"fmt"
	"io"
	"net/http"
	"strconv"
	"strings"
	"sync"

	"github.com/aws/aws-sdk-go/aws/awserr"
	"github.com/aws/aws-sdk-go/aws/awsutil"
	"github.com/aws/aws-sdk-go/aws/client"
	"github.com/aws/aws-sdk-go/aws/request"
	"github.com/aws/aws-sdk-go/service/s3"
	"github.com/aws/aws-sdk-go/service/s3/s3iface"
)

// DefaultDownloadPartSize is the default range of bytes to get at a time when
// using Download().
const DefaultDownloadPartSize = 1024 * 1024 * 5

// DefaultDownloadConcurrency is the default number of goroutines to spin up
// when using Download().
const DefaultDownloadConcurrency = 5

// The Downloader structure that calls Download(). It is safe to call Download()
// on this structure for multiple objects and across concurrent goroutines.
// Mutating the Downloader's properties is not safe to be done concurrently.
type Downloader struct {
	// The buffer size (in bytes) to use when buffering data into chunks and
	// sending them as parts to S3. The minimum allowed part size is 5MB, and
	// if this value is set to zero, the DefaultPartSize value will be used.
	PartSize int64

	// The number of goroutines to spin up in parallel when sending parts.
	// If this is set to zero, the DefaultConcurrency value will be used.
	Concurrency int

	// An S3 client to use when performing downloads.
	S3 s3iface.S3API
}

// NewDownloader creates a new Downloader instance to downloads objects from
// S3 in concurrent chunks. Pass in additional functional options  to customize
// the downloader behavior. Requires a client.ConfigProvider in order to create
// a S3 service client. The session.Session satisfies the client.ConfigProvider
// interface.
//
// Example:
//     // The session the S3 Downloader will use
//     sess := session.New()
//
//     // Create a downloader with the session and default options
//     downloader := s3manager.NewDownloader(sess)
//
//     // Create a downloader with the session and custom options
//     downloader := s3manager.NewDownloader(sess, func(d *s3manager.Uploader) {
//          d.PartSize = 64 * 1024 * 1024 // 64MB per part
//     })
func NewDownloader(c client.ConfigProvider, options ...func(*Downloader)) *Downloader {
	d := &Downloader{
		S3:          s3.New(c),
		PartSize:    DefaultDownloadPartSize,
		Concurrency: DefaultDownloadConcurrency,
	}
	for _, option := range options {
		option(d)
	}

	return d
}

// NewDownloaderWithClient creates a new Downloader instance to downloads
// objects from S3 in concurrent chunks. Pass in additional functional
// options to customize the downloader behavior. Requires a S3 service client
// to make S3 API calls.
//
// Example:
//     // The S3 client the S3 Downloader will use
//     s3Svc := s3.new(session.New())
//
//     // Create a downloader with the s3 client and default options
//     downloader := s3manager.NewDownloaderWithClient(s3Svc)
//
//     // Create a downloader with the s3 client and custom options
//     downloader := s3manager.NewDownloaderWithClient(s3Svc, func(d *s3manager.Uploader) {
//          d.PartSize = 64 * 1024 * 1024 // 64MB per part
//     })
func NewDownloaderWithClient(svc s3iface.S3API, options ...func(*Downloader)) *Downloader {
	d := &Downloader{
		S3:          svc,
		PartSize:    DefaultDownloadPartSize,
		Concurrency: DefaultDownloadConcurrency,
	}
	for _, option := range options {
		option(d)
	}

	return d
}

// Download downloads an object in S3 and writes the payload into w using
// concurrent GET requests.
//
// Additional functional options can be provided to configure the individual
// upload. These options are copies of the Uploader instance Upload is called from.
// Modifying the options will not impact the original Uploader instance.
//
// It is safe to call this method concurrently across goroutines.
//
// The w io.WriterAt can be satisfied by an os.File to do multipart concurrent
// downloads, or in memory []byte wrapper using aws.WriteAtBuffer.
func (d Downloader) Download(w io.WriterAt, input *s3.GetObjectInput, options ...func(*Downloader)) (n int64, err error) {
	impl := downloader{w: w, in: input, ctx: d}

	for _, option := range options {
		option(&impl.ctx)
	}

	return impl.download()
}

// downloader is the implementation structure used internally by Downloader.
type downloader struct {
	ctx Downloader

	in *s3.GetObjectInput
	w  io.WriterAt

	wg sync.WaitGroup
	m  sync.Mutex

	pos        int64
	totalBytes int64
	written    int64
	err        error
}

// init initializes the downloader with default options.
func (d *downloader) init() {
	d.totalBytes = -1

	if d.ctx.Concurrency == 0 {
		d.ctx.Concurrency = DefaultDownloadConcurrency
	}

	if d.ctx.PartSize == 0 {
		d.ctx.PartSize = DefaultDownloadPartSize
	}
}

// download performs the implementation of the object download across ranged
// GETs.
func (d *downloader) download() (n int64, err error) {
	d.init()

	// Spin off first worker to check additional header information
	d.getChunk()

	if total := d.getTotalBytes(); total >= 0 {
		// Spin up workers
		ch := make(chan dlchunk, d.ctx.Concurrency)

		for i := 0; i < d.ctx.Concurrency; i++ {
			d.wg.Add(1)
			go d.downloadPart(ch)
		}

		// Assign work
		for d.getErr() == nil {
			if d.pos >= total {
				break // We're finished queueing chunks
			}

			// Queue the next range of bytes to read.
			ch <- dlchunk{w: d.w, start: d.pos, size: d.ctx.PartSize}
			d.pos += d.ctx.PartSize
		}

		// Wait for completion
		close(ch)
		d.wg.Wait()
	} else {
		// Checking if we read anything new
		for d.err == nil {
			d.getChunk()
		}

		// We expect a 416 error letting us know we are done downloading the
		// total bytes. Since we do not know the content's length, this will
		// keep grabbing chunks of data until the range of bytes specified in
		// the request is out of range of the content. Once, this happens, a
		// 416 should occur.
		e, ok := d.err.(awserr.RequestFailure)
		if ok && e.StatusCode() == http.StatusRequestedRangeNotSatisfiable {
			d.err = nil
		}
	}

	// Return error
	return d.written, d.err
}

// downloadPart is an individual goroutine worker reading from the ch channel
// and performing a GetObject request on the data with a given byte range.
//
// If this is the first worker, this operation also resolves the total number
// of bytes to be read so that the worker manager knows when it is finished.
func (d *downloader) downloadPart(ch chan dlchunk) {
	defer d.wg.Done()
	for {
		chunk, ok := <-ch
		if !ok {
			break
		}
		d.downloadChunk(chunk)
	}
}

// getChunk grabs a chunk of data from the body.
// Not thread safe. Should only used when grabbing data on a single thread.
func (d *downloader) getChunk() {
	chunk := dlchunk{w: d.w, start: d.pos, size: d.ctx.PartSize}
	d.pos += d.ctx.PartSize
	d.downloadChunk(chunk)
}

// downloadChunk downloads the chunk froom s3
func (d *downloader) downloadChunk(chunk dlchunk) {
	if d.getErr() != nil {
		return
	}
	// Get the next byte range of data
	in := &s3.GetObjectInput{}
	awsutil.Copy(in, d.in)
	rng := fmt.Sprintf("bytes=%d-%d",
		chunk.start, chunk.start+chunk.size-1)
	in.Range = &rng

	req, resp := d.ctx.S3.GetObjectRequest(in)
	req.Handlers.Build.PushBack(request.MakeAddToUserAgentFreeFormHandler("S3Manager"))
	err := req.Send()

	if err != nil {
		d.setErr(err)
	} else {
		d.setTotalBytes(resp) // Set total if not yet set.

		n, err := io.Copy(&chunk, resp.Body)
		resp.Body.Close()

		if err != nil {
			d.setErr(err)
		}
		d.incrWritten(n)
	}
}

// getTotalBytes is a thread-safe getter for retrieving the total byte status.
func (d *downloader) getTotalBytes() int64 {
	d.m.Lock()
	defer d.m.Unlock()

	return d.totalBytes
}

// setTotalBytes is a thread-safe setter for setting the total byte status.
// Will extract the object's total bytes from the Content-Range if the file
// will be chunked, or Content-Length. Content-Length is used when the response
// does not include a Content-Range. Meaning the object was not chunked. This
// occurs when the full file fits within the PartSize directive.
func (d *downloader) setTotalBytes(resp *s3.GetObjectOutput) {
	d.m.Lock()
	defer d.m.Unlock()

	if d.totalBytes >= 0 {
		return
	}

	if resp.ContentRange == nil {
		// ContentRange is nil when the full file contents is provied, and
		// is not chunked. Use ContentLength instead.
		if resp.ContentLength != nil {
			d.totalBytes = *resp.ContentLength
			return
		}
	} else {
		parts := strings.Split(*resp.ContentRange, "/")

		total := int64(-1)
		var err error
		// Checking for whether or not a numbered total exists
		// If one does not exist, we will assume the total to be -1, undefined,
		// and sequentially download each chunk until hitting a 416 error
		totalStr := parts[len(parts)-1]
		if totalStr != "*" {
			total, err = strconv.ParseInt(totalStr, 10, 64)
			if err != nil {
				d.err = err
				return
			}
		}

		d.totalBytes = total
	}
}

func (d *downloader) incrWritten(n int64) {
	d.m.Lock()
	defer d.m.Unlock()

	d.written += n
}

// getErr is a thread-safe getter for the error object
func (d *downloader) getErr() error {
	d.m.Lock()
	defer d.m.Unlock()

	return d.err
}

// setErr is a thread-safe setter for the error object
func (d *downloader) setErr(e error) {
	d.m.Lock()
	defer d.m.Unlock()

	d.err = e
}

// dlchunk represents a single chunk of data to write by the worker routine.
// This structure also implements an io.SectionReader style interface for
// io.WriterAt, effectively making it an io.SectionWriter (which does not
// exist).
type dlchunk struct {
	w     io.WriterAt
	start int64
	size  int64
	cur   int64
}

// Write wraps io.WriterAt for the dlchunk, writing from the dlchunk's start
// position to its end (or EOF).
func (c *dlchunk) Write(p []byte) (n int, err error) {
	if c.cur >= c.size {
		return 0, io.EOF
	}

	n, err = c.w.WriteAt(p, c.start+c.cur)
	c.cur += int64(n)

	return
}
<|endoftext|>


<commit_before>package mysql

import (
	"database/sql"
	"fmt"
	"time"

	"github.com/jmoiron/sqlx"
	"github.com/solefaucet/sole-server/errors"
	"github.com/solefaucet/sole-server/models"
)

// GetRewardIncomes get user's reward incomes
func (s Storage) GetRewardIncomes(userID int64, limit, offset int64) ([]models.Income, error) {
	rawSQL := "SELECT * FROM incomes WHERE `user_id` = ? AND `type` = ? ORDER BY `id` DESC LIMIT ? OFFSET ?"
	args := []interface{}{userID, models.IncomeTypeReward, limit, offset}
	incomes := []models.Income{}
	err := s.selects(&incomes, rawSQL, args...)
	return incomes, err
}

// GetNumberOfRewardIncomes gets number of user's reward incomes
func (s Storage) GetNumberOfRewardIncomes(userID int64) (int64, error) {
	var count int64
	err := s.db.QueryRowx("SELECT COUNT(*) FROM incomes WHERE `user_id` = ? AND `type` = ?", userID, models.IncomeTypeReward).Scan(&count)
	return count, err
}

// CreateRewardIncome creates a new reward type income
func (s Storage) CreateRewardIncome(income models.Income, now time.Time) error {
	tx := s.db.MustBegin()

	if err := createRewardIncomeWithTx(tx, income, now); err != nil {
		tx.Rollback()
		return err
	}

	// commit
	if err := tx.Commit(); err != nil {
		return fmt.Errorf("create reward income commit transaction error: %v", err)
	}

	return nil
}

func createRewardIncomeWithTx(tx *sqlx.Tx, income models.Income, now time.Time) error {
	totalReward := income.Income

	// insert income into incomes table
	if _, err := addIncome(tx, income); err != nil {
		return err
	}

	// update user balance, total_income, referer_total_income
	if err := incrementUserBalance(tx, income.UserID, income.Income, income.RefererIncome); err != nil {
		return err
	}

	// update user rewarded_at
	if _, err := tx.Exec("UPDATE users SET `rewarded_at` = ? WHERE `id` = ?", now, income.UserID); err != nil {
		return err
	}

	// update referer balance
	if rowAffected, err := incrementRefererBalance(tx, income.RefererID, income.RefererIncome); err != nil {
		return err
	} else if rowAffected == 1 {
		totalReward += income.RefererIncome
	}

	// update total reward
	if err := incrementTotalReward(tx, totalReward, now); err != nil {
		return err
	}

	return nil
}

// insert reward income into incomes table
func addIncome(tx *sqlx.Tx, income models.Income) (sql.Result, error) {
	result, err := tx.NamedExec("INSERT INTO incomes (`user_id`, `referer_id`, `type`, `income`, `referer_income`) VALUES (:user_id, :referer_id, :type, :income, :referer_income)", income)
	if err != nil {
		return nil, fmt.Errorf("add income error: %v", err)
	}

	return result, nil
}

// increment user balance, total_income, referer_total_income
func incrementUserBalance(tx *sqlx.Tx, userID int64, delta, refererDelta float64) error {
	rawSQL := "UPDATE users SET `balance` = `balance` + ?, `total_income` = `total_income` + ?, `referer_total_income` = `referer_total_income` + ? WHERE id = ?"
	args := []interface{}{delta, delta, refererDelta, userID}
	if result, err := tx.Exec(rawSQL, args...); err != nil {
		return fmt.Errorf("increment user balance error: %v", err)
	} else if rowAffected, _ := result.RowsAffected(); rowAffected != 1 {
		return fmt.Errorf("increment user balance affected %v rows", rowAffected)
	}

	return nil
}

// increment referer balance
func incrementRefererBalance(tx *sqlx.Tx, refererID int64, delta float64) (int64, error) {
	result, err := tx.NamedExec("UPDATE users SET `balance` = `balance` + :delta, `total_income_from_referees` = `total_income_from_referees` + :delta WHERE id = :id", map[string]interface{}{
		"id":    refererID,
		"delta": delta,
	})
	if err != nil {
		return 0, fmt.Errorf("increment referer balance error: %v", err)
	}

	rowAffected, _ := result.RowsAffected()
	return rowAffected, nil
}

// increment total reward
func incrementTotalReward(tx *sqlx.Tx, totalReward float64, now time.Time) error {
	sql := "INSERT INTO total_rewards (`total`, `created_at`) VALUES (:delta, :created_at) ON DUPLICATE KEY UPDATE `total` = `total` + :delta"
	args := map[string]interface{}{
		"delta":      totalReward,
		"created_at": now,
	}

	if _, err := tx.NamedExec(sql, args); err != nil {
		return fmt.Errorf("increment total reward error: %v", err)
	}

	return nil
}

// GetOfferwowEventByID finds offerwow event by event id
func (s Storage) GetOfferwowEventByID(eventID string) (models.OfferwowEvent, error) {
	event := models.OfferwowEvent{}
	err := s.db.Get(&event, "SELECT * FROM `offerwow` WHERE `event_id` = ?", eventID)

	if err != nil {
		if err == sql.ErrNoRows {
			return event, errors.ErrNotFound
		}

		return event, fmt.Errorf("query offerwow event by id error: %v", err)
	}

	return event, nil
}

// CreateOfferwowIncome creates a new offerwow type income
func (s Storage) CreateOfferwowIncome(income models.Income, eventID string) error {
	tx := s.db.MustBegin()

	if err := createOfferwowIncomeWithTx(tx, income, eventID); err != nil {
		tx.Rollback()
		return err
	}

	// commit
	if err := tx.Commit(); err != nil {
		return fmt.Errorf("create offerwow income commit transaction error: %v", err)
	}

	return nil
}

func createOfferwowIncomeWithTx(tx *sqlx.Tx, income models.Income, eventID string) error {
	// insert offerwow income into incomes table
	result, err := addIncome(tx, income)
	if err != nil {
		return err
	}
	id, err := result.LastInsertId()
	if err != nil {
		return err
	}

	// update user balance, total_income, referer_total_income
	if err := incrementUserBalance(tx, income.UserID, income.Income, income.RefererIncome); err != nil {
		return err
	}

	// update referer balance
	if _, err := incrementRefererBalance(tx, income.RefererID, income.RefererIncome); err != nil {
		return err
	}

	// insert offerwow event
	offerwowEvent := models.OfferwowEvent{
		EventID:  eventID,
		IncomeID: id,
		Amount:   income.Income,
	}
	_, err = tx.NamedExec("INSERT INTO `offerwow` (`event_id`, `income_id`, `amount`) VALUE (:event_id, :income_id, :amount)", offerwowEvent)
	return err
}

// GetSuperrewardsOfferByID finds superrewards offer by transaction Id and user id
func (s Storage) GetSuperrewardsOfferByID(transactionID string, userID int64) (models.SuperrewardsOffer, error) {
	offer := models.SuperrewardsOffer{}
	err := s.db.Get(&offer, "SELECT * FROM `superrewards` WHERE `transaction_id` = ? AND `user_id` = ?", transactionID, userID)

	if err != nil {
		if err == sql.ErrNoRows {
			return offer, errors.ErrNotFound
		}

		return offer, fmt.Errorf("query superrewards offer by id error: %v", err)
	}

	return offer, nil
}

// CreateSuperrewardsIncome creates a new superrewards type income
func (s Storage) CreateSuperrewardsIncome(income models.Income, transactionID, offerID string) error {
	tx := s.db.MustBegin()

	if err := createSuperrewardsIncomeWithTx(tx, income, transactionID, offerID); err != nil {
		tx.Rollback()
		return err
	}

	// commit
	if err := tx.Commit(); err != nil {
		return fmt.Errorf("create superrewards income commit transaction error: %v", err)
	}

	return nil
}

func createSuperrewardsIncomeWithTx(tx *sqlx.Tx, income models.Income, transactionID, offerID string) error {
	// insert superrewards income into incomes table
	result, err := addIncome(tx, income)
	if err != nil {
		return err
	}
	id, err := result.LastInsertId()
	if err != nil {
		return err
	}

	// update user balance, total_income, referer_total_income
	if err := incrementUserBalance(tx, income.UserID, income.Income, income.RefererIncome); err != nil {
		return err
	}

	// update referer balance
	if _, err := incrementRefererBalance(tx, income.RefererID, income.RefererIncome); err != nil {
		return err
	}

	// insert superrewards offer
	offer := models.SuperrewardsOffer{
		IncomeID:      id,
		UserID:        income.UserID,
		TransactionID: transactionID,
		OfferID:       offerID,
		Amount:        income.Income,
	}
	_, err = tx.NamedExec("INSERT INTO `superrewards` (`income_id`, `user_id`, `transaction_id`, `offer_id`, `amount`) VALUE (:income_id, :user_id, :transaction_id, :offer_id, :amount)", offer)
	return err
}
<commit_msg>Fix vetshadow issue<commit_after>package mysql

import (
	"database/sql"
	"fmt"
	"time"

	"github.com/jmoiron/sqlx"
	"github.com/solefaucet/sole-server/errors"
	"github.com/solefaucet/sole-server/models"
)

// GetRewardIncomes get user's reward incomes
func (s Storage) GetRewardIncomes(userID int64, limit, offset int64) ([]models.Income, error) {
	rawSQL := "SELECT * FROM incomes WHERE `user_id` = ? AND `type` = ? ORDER BY `id` DESC LIMIT ? OFFSET ?"
	args := []interface{}{userID, models.IncomeTypeReward, limit, offset}
	incomes := []models.Income{}
	err := s.selects(&incomes, rawSQL, args...)
	return incomes, err
}

// GetNumberOfRewardIncomes gets number of user's reward incomes
func (s Storage) GetNumberOfRewardIncomes(userID int64) (int64, error) {
	var count int64
	err := s.db.QueryRowx("SELECT COUNT(*) FROM incomes WHERE `user_id` = ? AND `type` = ?", userID, models.IncomeTypeReward).Scan(&count)
	return count, err
}

// CreateRewardIncome creates a new reward type income
func (s Storage) CreateRewardIncome(income models.Income, now time.Time) error {
	tx := s.db.MustBegin()

	if err := createRewardIncomeWithTx(tx, income, now); err != nil {
		tx.Rollback()
		return err
	}

	// commit
	if err := tx.Commit(); err != nil {
		return fmt.Errorf("create reward income commit transaction error: %v", err)
	}

	return nil
}

func createRewardIncomeWithTx(tx *sqlx.Tx, income models.Income, now time.Time) error {
	totalReward := income.Income

	// insert income into incomes table
	if _, err := addIncome(tx, income); err != nil {
		return err
	}

	// update user balance, total_income, referer_total_income
	if err := incrementUserBalance(tx, income.UserID, income.Income, income.RefererIncome); err != nil {
		return err
	}

	// update user rewarded_at
	if _, err := tx.Exec("UPDATE users SET `rewarded_at` = ? WHERE `id` = ?", now, income.UserID); err != nil {
		return err
	}

	// update referer balance
	if rowAffected, err := incrementRefererBalance(tx, income.RefererID, income.RefererIncome); err != nil {
		return err
	} else if rowAffected == 1 {
		totalReward += income.RefererIncome
	}

	// update total reward
	if err := incrementTotalReward(tx, totalReward, now); err != nil {
		return err
	}

	return nil
}

// insert reward income into incomes table
func addIncome(tx *sqlx.Tx, income models.Income) (sql.Result, error) {
	result, err := tx.NamedExec("INSERT INTO incomes (`user_id`, `referer_id`, `type`, `income`, `referer_income`) VALUES (:user_id, :referer_id, :type, :income, :referer_income)", income)
	if err != nil {
		return nil, fmt.Errorf("add income error: %v", err)
	}

	return result, nil
}

// increment user balance, total_income, referer_total_income
func incrementUserBalance(tx *sqlx.Tx, userID int64, delta, refererDelta float64) error {
	rawSQL := "UPDATE users SET `balance` = `balance` + ?, `total_income` = `total_income` + ?, `referer_total_income` = `referer_total_income` + ? WHERE id = ?"
	args := []interface{}{delta, delta, refererDelta, userID}
	if result, err := tx.Exec(rawSQL, args...); err != nil {
		return fmt.Errorf("increment user balance error: %v", err)
	} else if rowAffected, _ := result.RowsAffected(); rowAffected != 1 {
		return fmt.Errorf("increment user balance affected %v rows", rowAffected)
	}

	return nil
}

// increment referer balance
func incrementRefererBalance(tx *sqlx.Tx, refererID int64, delta float64) (int64, error) {
	result, err := tx.NamedExec("UPDATE users SET `balance` = `balance` + :delta, `total_income_from_referees` = `total_income_from_referees` + :delta WHERE id = :id", map[string]interface{}{
		"id":    refererID,
		"delta": delta,
	})
	if err != nil {
		return 0, fmt.Errorf("increment referer balance error: %v", err)
	}

	rowAffected, _ := result.RowsAffected()
	return rowAffected, nil
}

// increment total reward
func incrementTotalReward(tx *sqlx.Tx, totalReward float64, now time.Time) error {
	sql := "INSERT INTO total_rewards (`total`, `created_at`) VALUES (:delta, :created_at) ON DUPLICATE KEY UPDATE `total` = `total` + :delta"
	args := map[string]interface{}{
		"delta":      totalReward,
		"created_at": now,
	}

	if _, err := tx.NamedExec(sql, args); err != nil {
		return fmt.Errorf("increment total reward error: %v", err)
	}

	return nil
}

// GetOfferwowEventByID finds offerwow event by event id
func (s Storage) GetOfferwowEventByID(eventID string) (models.OfferwowEvent, error) {
	event := models.OfferwowEvent{}
	err := s.db.Get(&event, "SELECT * FROM `offerwow` WHERE `event_id` = ?", eventID)

	if err != nil {
		if err == sql.ErrNoRows {
			return event, errors.ErrNotFound
		}

		return event, fmt.Errorf("query offerwow event by id error: %v", err)
	}

	return event, nil
}

// CreateOfferwowIncome creates a new offerwow type income
func (s Storage) CreateOfferwowIncome(income models.Income, eventID string) error {
	tx := s.db.MustBegin()

	if err := createOfferwowIncomeWithTx(tx, income, eventID); err != nil {
		tx.Rollback()
		return err
	}

	// commit
	if err := tx.Commit(); err != nil {
		return fmt.Errorf("create offerwow income commit transaction error: %v", err)
	}

	return nil
}

func createOfferwowIncomeWithTx(tx *sqlx.Tx, income models.Income, eventID string) error {
	// insert offerwow income into incomes table
	result, err := addIncome(tx, income)
	if err != nil {
		return err
	}
	id, err := result.LastInsertId()
	if err != nil {
		return err
	}

	// update user balance, total_income, referer_total_income
	if err = incrementUserBalance(tx, income.UserID, income.Income, income.RefererIncome); err != nil {
		return err
	}

	// update referer balance
	if _, err = incrementRefererBalance(tx, income.RefererID, income.RefererIncome); err != nil {
		return err
	}

	// insert offerwow event
	offerwowEvent := models.OfferwowEvent{
		EventID:  eventID,
		IncomeID: id,
		Amount:   income.Income,
	}
	_, err = tx.NamedExec("INSERT INTO `offerwow` (`event_id`, `income_id`, `amount`) VALUE (:event_id, :income_id, :amount)", offerwowEvent)
	return err
}

// GetSuperrewardsOfferByID finds superrewards offer by transaction Id and user id
func (s Storage) GetSuperrewardsOfferByID(transactionID string, userID int64) (models.SuperrewardsOffer, error) {
	offer := models.SuperrewardsOffer{}
	err := s.db.Get(&offer, "SELECT * FROM `superrewards` WHERE `transaction_id` = ? AND `user_id` = ?", transactionID, userID)

	if err != nil {
		if err == sql.ErrNoRows {
			return offer, errors.ErrNotFound
		}

		return offer, fmt.Errorf("query superrewards offer by id error: %v", err)
	}

	return offer, nil
}

// CreateSuperrewardsIncome creates a new superrewards type income
func (s Storage) CreateSuperrewardsIncome(income models.Income, transactionID, offerID string) error {
	tx := s.db.MustBegin()

	if err := createSuperrewardsIncomeWithTx(tx, income, transactionID, offerID); err != nil {
		tx.Rollback()
		return err
	}

	// commit
	if err := tx.Commit(); err != nil {
		return fmt.Errorf("create superrewards income commit transaction error: %v", err)
	}

	return nil
}

func createSuperrewardsIncomeWithTx(tx *sqlx.Tx, income models.Income, transactionID, offerID string) error {
	// insert superrewards income into incomes table
	result, err := addIncome(tx, income)
	if err != nil {
		return err
	}
	id, err := result.LastInsertId()
	if err != nil {
		return err
	}

	// update user balance, total_income, referer_total_income
	if err = incrementUserBalance(tx, income.UserID, income.Income, income.RefererIncome); err != nil {
		return err
	}

	// update referer balance
	if _, err = incrementRefererBalance(tx, income.RefererID, income.RefererIncome); err != nil {
		return err
	}

	// insert superrewards offer
	offer := models.SuperrewardsOffer{
		IncomeID:      id,
		UserID:        income.UserID,
		TransactionID: transactionID,
		OfferID:       offerID,
		Amount:        income.Income,
	}
	_, err = tx.NamedExec("INSERT INTO `superrewards` (`income_id`, `user_id`, `transaction_id`, `offer_id`, `amount`) VALUE (:income_id, :user_id, :transaction_id, :offer_id, :amount)", offer)
	return err
}
<|endoftext|>