<commit_before>package net

import (
	"context"
	"errors"
	"fmt"
	"strings"
	"sync"
	"time"

	"github.com/gobwas/glob"
	"go.uber.org/zap"

	"nimona.io/go/log"
	"nimona.io/go/peers"
	"nimona.io/go/primitives"
	"nimona.io/go/storage"
	"nimona.io/go/utils"
)

var (
	// ErrInvalidRequest when received an invalid request block
	ErrInvalidRequest = errors.New("Invalid request")
)

var (
	// ErrAllAddressesFailed for when a peer cannot be dialed
	ErrAllAddressesFailed = errors.New("all addresses failed to dial")
	// ErrNoAddresses for when a peer has no addresses
	ErrNoAddresses = errors.New("no addresses")
	// ErrNotForUs block is not meant for us
	ErrNotForUs = errors.New("block not for us")
)

// Exchange interface for mocking exchange
type Exchange interface {
	Get(ctx context.Context, id string) (interface{}, error)
	Handle(contentType string, h func(o *primitives.Block) error) (func(), error)
	Send(ctx context.Context, o *primitives.Block, address string, opts ...primitives.SendOption) error
	RegisterDiscoverer(discovery Discoverer)
}

type exchange struct {
	network     Networker
	addressBook *peers.AddressBook
	manager     *ConnectionManager
	discovery   Discoverer

	outgoingPayloads chan *outBlock
	incomingPayloads chan *incBlock

	handlers   sync.Map
	logger     *zap.Logger
	streamLock utils.Kmutex

	store         storage.Storage
	getRequests   sync.Map
	subscriptions sync.Map
}

type outBlock struct {
	context   context.Context
	recipient string
	opts      []primitives.SendOption
	block     *primitives.Block
	err       chan error
}

type incBlock struct {
	conn  *Connection
	typed *primitives.Block
}

type handler struct {
	contentType glob.Glob
	handler     func(o *primitives.Block) error
}

// NewExchange creates a exchange on a given network
func NewExchange(addressBook *peers.AddressBook, store storage.Storage, address string) (Exchange, error) {
	ctx := context.Background()

	network, err := NewNetwork(addressBook)
	if err != nil {
		return nil, err
	}

	w := &exchange{
		network:     network,
		addressBook: addressBook,
		manager:     &ConnectionManager{},

		outgoingPayloads: make(chan *outBlock, 10),
		incomingPayloads: make(chan *incBlock, 10),

		handlers:   sync.Map{},
		logger:     log.Logger(ctx).Named("exchange"),
		streamLock: utils.NewKmutex(),

		store:       store,
		getRequests: sync.Map{},
	}

	incomingConnections, err := w.network.Listen(ctx, address)
	if err != nil {
		return nil, err
	}

	go func() {
		for {
			conn := <-incomingConnections
			go func(conn *Connection) {
				w.manager.Add("peer:"+conn.RemoteID, conn)
				if err := w.HandleConnection(conn); err != nil {
					w.logger.Warn("failed to handle block", zap.Error(err))
				}
			}(conn)
		}
	}()

	go func() {
		for {
			block := <-w.incomingPayloads
			go func(block *incBlock) {
				defer func() {
					if r := recover(); r != nil {
						w.logger.Error("Recovered while processing", zap.Any("r", r))
					}
				}()
				if err := w.process(block.typed, block.conn); err != nil {
					w.logger.Error("processing block", zap.Error(err))
				}
			}(block)
		}
	}()

	go func() {
		for {
			block := <-w.outgoingPayloads
			if block.recipient == "" {
				w.logger.Info("missing recipient")
				block.err <- errors.New("missing recipient")
				continue
			}

			logger := w.logger.With(zap.String("recipient", block.recipient))

			// try to send the block directly to the recipient
			logger.Debug("getting conn to write block")
			conn, err := w.GetOrDial(block.context, block.recipient)
			if err != nil {
				logger.Debug("could not get conn to recipient", zap.Error(err))
				block.err <- err
				continue
			}

			if err := Write(block.block, conn); err != nil {
				w.manager.Close(block.recipient)
				logger.Debug("could not write to recipient", zap.Error(err))
				block.err <- err
				continue
			}

			// update peer status
			// TODO(geoah) fix status -- not that we are using this
			// w.addressBook.PutPeerStatus(recipientThumbPrint, peers.StatusConnected)
			block.err <- nil
			continue
		}
	}()

	return w, nil
}

func (w *exchange) RegisterDiscoverer(discovery Discoverer) {
	w.discovery = discovery
	ctx := context.Background()
	go func() {
		for {
			blocks, err := w.store.List()
			if err != nil {
				time.Sleep(time.Second * 10)
				continue
			}

			for _, block := range blocks {
				if err := w.discovery.PutProviders(ctx, block); err != nil {
					w.logger.Warn("could not announce provider for block", zap.String("id", block))
				}
			}

			time.Sleep(time.Second * 30)
		}
	}()
}

func (w *exchange) Handle(typePatern string, h func(o *primitives.Block) error) (func(), error) {
	g, err := glob.Compile(typePatern, '.', '/', '#')
	if err != nil {
		return nil, err
	}
	hID := RandStringBytesMaskImprSrc(8)
	w.handlers.Store(hID, &handler{
		contentType: g,
		handler:     h,
	})
	r := func() {
		w.handlers.Delete(hID)
	}
	return r, nil
}

func (w *exchange) HandleConnection(conn *Connection) error {
	w.logger.Debug("handling new connection", zap.String("remote", conn.RemoteID))
	for {
		// TODO use decoder
		typed, err := Read(conn)
		if err != nil {
			return err
		}

		w.incomingPayloads <- &incBlock{
			conn:  conn,
			typed: typed,
		}
	}
}

// Process incoming block
func (w *exchange) process(block *primitives.Block, conn *Connection) error {
	// TODO verify signature

	// TODO convert these into proper handlers
	switch block.Type {
	case "nimona.io/block.forward.request":
		req := &ForwardRequest{}
		req.FromBlock(block)
		w.logger.Info("got forwarded message", zap.String("recipient", req.Recipient))
		cerr := make(chan error, 1)
		ctx, cf := context.WithTimeout(context.Background(), time.Second)
		defer cf()
		w.outgoingPayloads <- &outBlock{
			context:   ctx,
			recipient: req.Recipient,
			block:     req.FwBlock,
			err:       cerr,
		}
		return <-cerr

	case "nimona.io/block.request":
		req := &BlockRequest{}
		req.FromBlock(block)
		if err := w.handleBlockRequest(req); err != nil {
			w.logger.Warn("could not handle request block", zap.Error(err))
		}

	case "nimona.io/block.response":
		res := &BlockResponse{}
		res.FromBlock(block)
		if err := w.handleBlockResponse(res); err != nil {
			w.logger.Warn("could not handle request block", zap.Error(err))
		}
	}

	hp := 0
	w.handlers.Range(func(_, v interface{}) bool {
		hp++
		h := v.(*handler)
		if h.contentType.Match(block.Type) {
			go func(h *handler, block *primitives.Block) {
				defer func() {
					if r := recover(); r != nil {
						w.logger.Error("Recovered while handling", zap.Any("r", r))
					}
				}()
				if err := h.handler(block); err != nil {
					w.logger.Info(
						"Could not handle event",
						zap.String("contentType", block.Type),
						zap.Error(err),
					)
				}
			}(h, block)
		}
		return true
	})

	if hp == 0 {
		fmt.Println("++++++++++++++++++++++++++++++++++++++++++++++++++++++++")
		fmt.Println("+++++ NO HANDLERS ++++++++++++++++++++++++++++++++++++++")
		fmt.Println("++++++++++++++++++++++++++++++++++++++++++++++++++++++++")
		fmt.Println("+++++", block.Type)
		fmt.Println("+++++", block.Payload)
		fmt.Println("++++++++++++++++++++++++++++++++++++++++++++++++++++++++")
	} else {
		if shouldPersist(block.Type) {
			bytes, _ := primitives.Marshal(block)
			w.store.Store(block.ID(), bytes)
		}
	}

	return nil
}

func (w *exchange) handleBlockResponse(payload *BlockResponse) error {
	// Check if nonce exists in local addressBook
	value, ok := w.getRequests.Load(payload.RequestID)
	if !ok {
		return nil
	}

	req, ok := value.(*BlockRequest)
	if !ok {
		return ErrInvalidRequest
	}

	req.response <- payload.Block
	return nil
}

func (w *exchange) handleBlockRequest(payload *BlockRequest) error {
	blockBytes, err := w.store.Get(payload.ID)
	if err != nil {
		return err
	}

	// TODO check if policy allows requested to retrieve the block

	block, err := primitives.Unmarshal(blockBytes)
	if err != nil {
		return err
	}

	resp := &BlockResponse{
		RequestID:      payload.RequestID,
		RequestedBlock: block,
	}

	signer := w.addressBook.GetLocalPeerKey()
	addr := "peer:" + payload.Sender.Thumbprint()
	if err := w.Send(context.Background(), resp.Block(), addr, primitives.SignWith(signer)); err != nil {
		w.logger.Warn("blx.handleBlockRequest could not send block", zap.Error(err))
		return err
	}

	return nil
}

func (w *exchange) Get(ctx context.Context, id string) (interface{}, error) {
	// Check local storage for block
	if blockBytes, err := w.store.Get(id); err == nil {
		block, err := primitives.Unmarshal(blockBytes)
		if err != nil {
			return nil, err
		}
		return block, nil
	}

	req := &BlockRequest{
		RequestID: RandStringBytesMaskImprSrc(8),
		ID:        id,
		response:  make(chan interface{}, 10),
	}

	defer func() {
		w.getRequests.Delete(req.RequestID)
		close(req.response)
	}()

	w.getRequests.Store(req.RequestID, req)
	// signer := w.addressBook.GetLocalPeerKey()

	// go func() {
	// 	providers, err := w.discovery.GetProviders(ctx, id)
	// 	if err != nil {
	// 		// TODO log err
	// 		return
	// 	}

	// 	for provider := range providers {
	// 		addr := "peer:" + provider.Thumbprint()
	// 		if err := w.Send(ctx, req.Block(), addr, primitives.SignWith(signer)); err != nil {
	// 			w.logger.Warn("blx.Get could not send req block", zap.Error(err))
	// 		}
	// 	}
	// }()

	for {
		select {
		case payload := <-req.response:
			return payload, nil

		case <-ctx.Done():
			return nil, storage.ErrNotFound
		}
	}
}

func (w *exchange) Send(ctx context.Context, block *primitives.Block, address string,
	opts ...primitives.SendOption) error {

	cfg := primitives.ParseSendOptions(opts...)

	if cfg.Sign && cfg.Key != nil && block.Signature == nil {
		if err := primitives.Sign(block, cfg.Key); err != nil {
			return err
		}
	}

	if shouldPersist(block.Type) {
		bytes, _ := primitives.Marshal(block)
		w.store.Store(block.ID(), bytes)
	}

	cerr := make(chan error, 1)
	w.outgoingPayloads <- &outBlock{
		context:   ctx,
		recipient: address,
		block:     block,
		err:       cerr,
		opts:      opts,
	}
	if err := <-cerr; err == nil {
		return nil
	}

	recipient := ""
	if strings.HasPrefix(address, "peer:") {
		recipient = strings.Replace(address, "peer:", "", 1)
	} else {
		return ErrAllAddressesFailed
	}

	// try to send message via their relay addresses
	peer, err := w.addressBook.GetPeerInfo(recipient)
	if err != nil {
		return err
	}

	// we unpack this because of the Send opts that might have signed this
	// fwTyped, err := primitives.UnpackDecode(bytes)
	// if err != nil {
	// 	return err
	// }

	// TODO(geoah) make sure fw block is signed
	fw := &ForwardRequest{
		Recipient: address,
		FwBlock:   block,
	}

	for _, address := range peer.Addresses {
		if strings.HasPrefix(address, "relay:") {
			relayPeerID := strings.Replace(address, "relay:", "", 1)
			if w.addressBook.GetLocalPeerInfo().Thumbprint() == relayPeerID {
				continue
			}
			relayPeer, err := w.addressBook.GetPeerInfo(relayPeerID)
			if err != nil {
				continue
			}
			cerr := make(chan error, 1)
			w.outgoingPayloads <- &outBlock{
				context:   ctx,
				recipient: "peer:" + relayPeer.Signature.Key.Thumbprint(),
				block:     fw.Block(),
				err:       cerr,
				opts:      opts,
			}
			if err := <-cerr; err == nil {
				return nil
			}
		}
	}

	return ErrAllAddressesFailed
}

func (w *exchange) GetOrDial(ctx context.Context, address string) (*Connection, error) {
	w.logger.Debug("looking for existing connection", zap.String("address", address))
	if address == "" {
		return nil, errors.New("missing address")
	}

	existingConn, err := w.manager.Get(address)
	if err == nil {
		w.logger.Debug("found existing connection", zap.String("address", address))
		return existingConn, nil
	}

	conn, err := w.network.Dial(ctx, address)
	if err != nil {
		// w.manager.Close(address)
		return nil, err
	}

	go func() {
		if err := w.HandleConnection(conn); err != nil {
			w.logger.Warn("failed to handle block", zap.Error(err))
		}
	}()

	w.manager.Add(address, conn)

	return conn, nil
}

func shouldPersist(t string) bool {
	if strings.HasPrefix(t, "nimona.io/dht") ||
		strings.HasPrefix(t, "nimona.io/telemetry") ||
		strings.HasPrefix(t, "nimona.io/handshake") {
		return false
	}
	return true
}
<commit_msg>Remove re-pushing providers<commit_after>package net

import (
	"context"
	"errors"
	"fmt"
	"strings"
	"sync"
	"time"

	"github.com/gobwas/glob"
	"go.uber.org/zap"

	"nimona.io/go/log"
	"nimona.io/go/peers"
	"nimona.io/go/primitives"
	"nimona.io/go/storage"
	"nimona.io/go/utils"
)

var (
	// ErrInvalidRequest when received an invalid request block
	ErrInvalidRequest = errors.New("Invalid request")
)

var (
	// ErrAllAddressesFailed for when a peer cannot be dialed
	ErrAllAddressesFailed = errors.New("all addresses failed to dial")
	// ErrNoAddresses for when a peer has no addresses
	ErrNoAddresses = errors.New("no addresses")
	// ErrNotForUs block is not meant for us
	ErrNotForUs = errors.New("block not for us")
)

// Exchange interface for mocking exchange
type Exchange interface {
	Get(ctx context.Context, id string) (interface{}, error)
	Handle(contentType string, h func(o *primitives.Block) error) (func(), error)
	Send(ctx context.Context, o *primitives.Block, address string, opts ...primitives.SendOption) error
	RegisterDiscoverer(discovery Discoverer)
}

type exchange struct {
	network     Networker
	addressBook *peers.AddressBook
	manager     *ConnectionManager
	discovery   Discoverer

	outgoingPayloads chan *outBlock
	incomingPayloads chan *incBlock

	handlers   sync.Map
	logger     *zap.Logger
	streamLock utils.Kmutex

	store         storage.Storage
	getRequests   sync.Map
	subscriptions sync.Map
}

type outBlock struct {
	context   context.Context
	recipient string
	opts      []primitives.SendOption
	block     *primitives.Block
	err       chan error
}

type incBlock struct {
	conn  *Connection
	typed *primitives.Block
}

type handler struct {
	contentType glob.Glob
	handler     func(o *primitives.Block) error
}

// NewExchange creates a exchange on a given network
func NewExchange(addressBook *peers.AddressBook, store storage.Storage, address string) (Exchange, error) {
	ctx := context.Background()

	network, err := NewNetwork(addressBook)
	if err != nil {
		return nil, err
	}

	w := &exchange{
		network:     network,
		addressBook: addressBook,
		manager:     &ConnectionManager{},

		outgoingPayloads: make(chan *outBlock, 10),
		incomingPayloads: make(chan *incBlock, 10),

		handlers:   sync.Map{},
		logger:     log.Logger(ctx).Named("exchange"),
		streamLock: utils.NewKmutex(),

		store:       store,
		getRequests: sync.Map{},
	}

	incomingConnections, err := w.network.Listen(ctx, address)
	if err != nil {
		return nil, err
	}

	go func() {
		for {
			conn := <-incomingConnections
			go func(conn *Connection) {
				w.manager.Add("peer:"+conn.RemoteID, conn)
				if err := w.HandleConnection(conn); err != nil {
					w.logger.Warn("failed to handle block", zap.Error(err))
				}
			}(conn)
		}
	}()

	go func() {
		for {
			block := <-w.incomingPayloads
			go func(block *incBlock) {
				defer func() {
					if r := recover(); r != nil {
						w.logger.Error("Recovered while processing", zap.Any("r", r))
					}
				}()
				if err := w.process(block.typed, block.conn); err != nil {
					w.logger.Error("processing block", zap.Error(err))
				}
			}(block)
		}
	}()

	go func() {
		for {
			block := <-w.outgoingPayloads
			if block.recipient == "" {
				w.logger.Info("missing recipient")
				block.err <- errors.New("missing recipient")
				continue
			}

			logger := w.logger.With(zap.String("recipient", block.recipient))

			// try to send the block directly to the recipient
			logger.Debug("getting conn to write block")
			conn, err := w.GetOrDial(block.context, block.recipient)
			if err != nil {
				logger.Debug("could not get conn to recipient", zap.Error(err))
				block.err <- err
				continue
			}

			if err := Write(block.block, conn); err != nil {
				w.manager.Close(block.recipient)
				logger.Debug("could not write to recipient", zap.Error(err))
				block.err <- err
				continue
			}

			// update peer status
			// TODO(geoah) fix status -- not that we are using this
			// w.addressBook.PutPeerStatus(recipientThumbPrint, peers.StatusConnected)
			block.err <- nil
			continue
		}
	}()

	return w, nil
}

func (w *exchange) RegisterDiscoverer(discovery Discoverer) {
	// TODO race?
	w.discovery = discovery
}

func (w *exchange) Handle(typePatern string, h func(o *primitives.Block) error) (func(), error) {
	g, err := glob.Compile(typePatern, '.', '/', '#')
	if err != nil {
		return nil, err
	}
	hID := RandStringBytesMaskImprSrc(8)
	w.handlers.Store(hID, &handler{
		contentType: g,
		handler:     h,
	})
	r := func() {
		w.handlers.Delete(hID)
	}
	return r, nil
}

func (w *exchange) HandleConnection(conn *Connection) error {
	w.logger.Debug("handling new connection", zap.String("remote", conn.RemoteID))
	for {
		// TODO use decoder
		typed, err := Read(conn)
		if err != nil {
			return err
		}

		w.incomingPayloads <- &incBlock{
			conn:  conn,
			typed: typed,
		}
	}
}

// Process incoming block
func (w *exchange) process(block *primitives.Block, conn *Connection) error {
	// TODO verify signature

	// TODO convert these into proper handlers
	switch block.Type {
	case "nimona.io/block.forward.request":
		req := &ForwardRequest{}
		req.FromBlock(block)
		w.logger.Info("got forwarded message", zap.String("recipient", req.Recipient))
		cerr := make(chan error, 1)
		ctx, cf := context.WithTimeout(context.Background(), time.Second)
		defer cf()
		w.outgoingPayloads <- &outBlock{
			context:   ctx,
			recipient: req.Recipient,
			block:     req.FwBlock,
			err:       cerr,
		}
		return <-cerr

	case "nimona.io/block.request":
		req := &BlockRequest{}
		req.FromBlock(block)
		if err := w.handleBlockRequest(req); err != nil {
			w.logger.Warn("could not handle request block", zap.Error(err))
		}

	case "nimona.io/block.response":
		res := &BlockResponse{}
		res.FromBlock(block)
		if err := w.handleBlockResponse(res); err != nil {
			w.logger.Warn("could not handle request block", zap.Error(err))
		}
	}

	hp := 0
	w.handlers.Range(func(_, v interface{}) bool {
		hp++
		h := v.(*handler)
		if h.contentType.Match(block.Type) {
			go func(h *handler, block *primitives.Block) {
				defer func() {
					if r := recover(); r != nil {
						w.logger.Error("Recovered while handling", zap.Any("r", r))
					}
				}()
				if err := h.handler(block); err != nil {
					w.logger.Info(
						"Could not handle event",
						zap.String("contentType", block.Type),
						zap.Error(err),
					)
				}
			}(h, block)
		}
		return true
	})

	if hp == 0 {
		fmt.Println("++++++++++++++++++++++++++++++++++++++++++++++++++++++++")
		fmt.Println("+++++ NO HANDLERS ++++++++++++++++++++++++++++++++++++++")
		fmt.Println("++++++++++++++++++++++++++++++++++++++++++++++++++++++++")
		fmt.Println("+++++", block.Type)
		fmt.Println("+++++", block.Payload)
		fmt.Println("++++++++++++++++++++++++++++++++++++++++++++++++++++++++")
	} else {
		if shouldPersist(block.Type) {
			bytes, _ := primitives.Marshal(block)
			w.store.Store(block.ID(), bytes)
		}
	}

	return nil
}

func (w *exchange) handleBlockResponse(payload *BlockResponse) error {
	// Check if nonce exists in local addressBook
	value, ok := w.getRequests.Load(payload.RequestID)
	if !ok {
		return nil
	}

	req, ok := value.(*BlockRequest)
	if !ok {
		return ErrInvalidRequest
	}

	req.response <- payload.Block
	return nil
}

func (w *exchange) handleBlockRequest(payload *BlockRequest) error {
	blockBytes, err := w.store.Get(payload.ID)
	if err != nil {
		return err
	}

	// TODO check if policy allows requested to retrieve the block

	block, err := primitives.Unmarshal(blockBytes)
	if err != nil {
		return err
	}

	resp := &BlockResponse{
		RequestID:      payload.RequestID,
		RequestedBlock: block,
	}

	signer := w.addressBook.GetLocalPeerKey()
	addr := "peer:" + payload.Sender.Thumbprint()
	if err := w.Send(context.Background(), resp.Block(), addr, primitives.SignWith(signer)); err != nil {
		w.logger.Warn("blx.handleBlockRequest could not send block", zap.Error(err))
		return err
	}

	return nil
}

func (w *exchange) Get(ctx context.Context, id string) (interface{}, error) {
	// Check local storage for block
	if blockBytes, err := w.store.Get(id); err == nil {
		block, err := primitives.Unmarshal(blockBytes)
		if err != nil {
			return nil, err
		}
		return block, nil
	}

	req := &BlockRequest{
		RequestID: RandStringBytesMaskImprSrc(8),
		ID:        id,
		response:  make(chan interface{}, 10),
	}

	defer func() {
		w.getRequests.Delete(req.RequestID)
		close(req.response)
	}()

	w.getRequests.Store(req.RequestID, req)
	// signer := w.addressBook.GetLocalPeerKey()

	// go func() {
	// 	providers, err := w.discovery.GetProviders(ctx, id)
	// 	if err != nil {
	// 		// TODO log err
	// 		return
	// 	}

	// 	for provider := range providers {
	// 		addr := "peer:" + provider.Thumbprint()
	// 		if err := w.Send(ctx, req.Block(), addr, primitives.SignWith(signer)); err != nil {
	// 			w.logger.Warn("blx.Get could not send req block", zap.Error(err))
	// 		}
	// 	}
	// }()

	for {
		select {
		case payload := <-req.response:
			return payload, nil

		case <-ctx.Done():
			return nil, storage.ErrNotFound
		}
	}
}

func (w *exchange) Send(ctx context.Context, block *primitives.Block, address string,
	opts ...primitives.SendOption) error {

	cfg := primitives.ParseSendOptions(opts...)

	if cfg.Sign && cfg.Key != nil && block.Signature == nil {
		if err := primitives.Sign(block, cfg.Key); err != nil {
			return err
		}
	}

	if shouldPersist(block.Type) {
		bytes, _ := primitives.Marshal(block)
		w.store.Store(block.ID(), bytes)
	}

	cerr := make(chan error, 1)
	w.outgoingPayloads <- &outBlock{
		context:   ctx,
		recipient: address,
		block:     block,
		err:       cerr,
		opts:      opts,
	}
	if err := <-cerr; err == nil {
		return nil
	}

	recipient := ""
	if strings.HasPrefix(address, "peer:") {
		recipient = strings.Replace(address, "peer:", "", 1)
	} else {
		return ErrAllAddressesFailed
	}

	// try to send message via their relay addresses
	peer, err := w.addressBook.GetPeerInfo(recipient)
	if err != nil {
		return err
	}

	// we unpack this because of the Send opts that might have signed this
	// fwTyped, err := primitives.UnpackDecode(bytes)
	// if err != nil {
	// 	return err
	// }

	// TODO(geoah) make sure fw block is signed
	fw := &ForwardRequest{
		Recipient: address,
		FwBlock:   block,
	}

	for _, address := range peer.Addresses {
		if strings.HasPrefix(address, "relay:") {
			relayPeerID := strings.Replace(address, "relay:", "", 1)
			if w.addressBook.GetLocalPeerInfo().Thumbprint() == relayPeerID {
				continue
			}
			relayPeer, err := w.addressBook.GetPeerInfo(relayPeerID)
			if err != nil {
				continue
			}
			cerr := make(chan error, 1)
			w.outgoingPayloads <- &outBlock{
				context:   ctx,
				recipient: "peer:" + relayPeer.Signature.Key.Thumbprint(),
				block:     fw.Block(),
				err:       cerr,
				opts:      opts,
			}
			if err := <-cerr; err == nil {
				return nil
			}
		}
	}

	return ErrAllAddressesFailed
}

func (w *exchange) GetOrDial(ctx context.Context, address string) (*Connection, error) {
	w.logger.Debug("looking for existing connection", zap.String("address", address))
	if address == "" {
		return nil, errors.New("missing address")
	}

	existingConn, err := w.manager.Get(address)
	if err == nil {
		w.logger.Debug("found existing connection", zap.String("address", address))
		return existingConn, nil
	}

	conn, err := w.network.Dial(ctx, address)
	if err != nil {
		// w.manager.Close(address)
		return nil, err
	}

	go func() {
		if err := w.HandleConnection(conn); err != nil {
			w.logger.Warn("failed to handle block", zap.Error(err))
		}
	}()

	w.manager.Add(address, conn)

	return conn, nil
}

func shouldPersist(t string) bool {
	if strings.HasPrefix(t, "nimona.io/dht") ||
		strings.HasPrefix(t, "nimona.io/telemetry") ||
		strings.HasPrefix(t, "nimona.io/handshake") {
		return false
	}
	return true
}
<|endoftext|>


<commit_before>package main

import (
	"bytes"
	"encoding/json"
	"fmt"
	"io"
	"net/http"
	"os"
	"strconv"
	"time"

	"github.com/estafette/estafette-ci-contracts"

	"github.com/rs/zerolog/log"
	"github.com/sethgrid/pester"
)

// EndOfLifeHelper has methods to shutdown the runner after a fatal or successful run
type EndOfLifeHelper interface {
	handleFatal(contracts.BuildLog, error, string)
	sendBuildFinishedEvent(buildStatus string) error
	sendBuildCleanEvent(buildStatus string) error
	sendBuildJobLogEvent(buildLog contracts.BuildLog) error
}

type endOfLifeHelperImpl struct {
	runAsJob bool
	config   contracts.BuilderConfig
}

// NewEndOfLifeHelper returns a new EndOfLifeHelper
func NewEndOfLifeHelper(runAsJob bool, config contracts.BuilderConfig) EndOfLifeHelper {
	return &endOfLifeHelperImpl{
		runAsJob: runAsJob,
		config:   config,
	}
}

func (elh *endOfLifeHelperImpl) handleFatal(buildLog contracts.BuildLog, err error, message string) {

	// add error messages as step to show in logs
	fatalStep := contracts.BuildLogStep{
		Step:     "init",
		LogLines: []contracts.BuildLogLine{},
		ExitCode: -1,
		Status:   "failed",
	}
	lineNumber := 1

	if err != nil {
		fatalStep.LogLines = append(fatalStep.LogLines, contracts.BuildLogLine{
			LineNumber: lineNumber,
			Timestamp:  time.Now().UTC(),
			StreamType: "stderr",
			Text:       err.Error(),
		})
		lineNumber++
	}
	if message != "" {
		fatalStep.LogLines = append(fatalStep.LogLines, contracts.BuildLogLine{
			LineNumber: lineNumber,
			Timestamp:  time.Now().UTC(),
			StreamType: "stderr",
			Text:       message,
		})
		lineNumber++
	}

	buildLog.Steps = append(buildLog.Steps, fatalStep)

	_ = elh.sendBuildFinishedEvent("failed")
	_ = elh.sendBuildJobLogEvent(buildLog)
	_ = elh.sendBuildCleanEvent("failed")

	if elh.runAsJob {
		log.Error().Err(err).Msg(message)
		os.Exit(0)
	} else {
		log.Fatal().Err(err).Msg(message)
	}
}

func (elh *endOfLifeHelperImpl) sendBuildJobLogEvent(buildLog contracts.BuildLog) (err error) {

	ciServerBuilderPostLogsURL := elh.config.CIServer.PostLogsURL
	ciAPIKey := elh.config.CIServer.APIKey
	jobName := *elh.config.JobName

	if ciServerBuilderPostLogsURL != "" && ciAPIKey != "" && jobName != "" {

		// convert BuildJobLogs to json
		var requestBody io.Reader

		var data []byte
		if *elh.config.Action == "release" {
			// copy buildLog to releaseLog and marshal that
			releaseLog := contracts.ReleaseLog{
				ID:         buildLog.ID,
				RepoSource: buildLog.RepoSource,
				RepoOwner:  buildLog.RepoOwner,
				RepoName:   buildLog.RepoName,
				ReleaseID:  strconv.Itoa(elh.config.ReleaseParams.ReleaseID),
				Steps:      buildLog.Steps,
				InsertedAt: buildLog.InsertedAt,
			}
			data, err = json.Marshal(releaseLog)
			if err != nil {
				log.Error().Err(err).Msgf("Failed marshalling ReleaseLog for job %v", jobName)
				return
			}
		} else {
			data, err = json.Marshal(buildLog)
			if err != nil {
				log.Error().Err(err).Msgf("Failed marshalling BuildLog for job %v", jobName)
				return
			}
		}

		requestBody = bytes.NewReader(data)

		// create client, in order to add headers
		client := pester.New()
		client.MaxRetries = 5
		client.Backoff = pester.ExponentialJitterBackoff
		client.KeepLog = true
		client.Timeout = time.Duration(120) * time.Second
		request, err := http.NewRequest("POST", ciServerBuilderPostLogsURL, requestBody)
		if err != nil {
			log.Error().Err(err).Msgf("Failed creating http client for job %v", jobName)
			return err
		}

		// add headers
		request.Header.Add("Authorization", fmt.Sprintf("Bearer %v", ciAPIKey))
		request.Header.Add("Content-Type", "application/json")

		// perform actual request
		response, err := client.Do(request)
		if err != nil {
			log.Error().Err(err).Str("logs", client.LogString()).Msgf("Failed shipping logs to %v for job %v", ciServerBuilderPostLogsURL, jobName)
			return err
		}

		defer response.Body.Close()

		log.Debug().Str("logs", client.LogString()).Msgf("Successfully shipped logs to %v for job %v", ciServerBuilderPostLogsURL, jobName)
	}

	return nil
}

func (elh *endOfLifeHelperImpl) sendBuildFinishedEvent(buildStatus string) error {
	return elh.sendBuilderEvent(buildStatus, fmt.Sprintf("builder:%v", buildStatus))
}

func (elh *endOfLifeHelperImpl) sendBuildCleanEvent(buildStatus string) error {
	return elh.sendBuilderEvent(buildStatus, "builder:clean")
}

func (elh *endOfLifeHelperImpl) sendBuilderEvent(buildStatus, event string) (err error) {

	ciServerBuilderEventsURL := elh.config.CIServer.BuilderEventsURL
	ciAPIKey := elh.config.CIServer.APIKey
	jobName := *elh.config.JobName

	if ciServerBuilderEventsURL != "" && ciAPIKey != "" && jobName != "" {
		// convert EstafetteCiBuilderEvent to json
		var requestBody io.Reader

		buildID := ""
		if elh.config.BuildParams != nil {
			buildID = strconv.Itoa(elh.config.BuildParams.BuildID)
		}

		releaseID := ""
		if elh.config.ReleaseParams != nil {
			releaseID = strconv.Itoa(elh.config.ReleaseParams.ReleaseID)
		}

		ciBuilderEvent := EstafetteCiBuilderEvent{
			JobName:      jobName,
			RepoSource:   elh.config.Git.RepoSource,
			RepoOwner:    elh.config.Git.RepoOwner,
			RepoName:     elh.config.Git.RepoName,
			RepoBranch:   elh.config.Git.RepoBranch,
			RepoRevision: elh.config.Git.RepoRevision,
			ReleaseID:    releaseID,
			BuildID:      buildID,
			BuildStatus:  buildStatus,
		}
		data, err := json.Marshal(ciBuilderEvent)
		if err != nil {
			log.Error().Err(err).Msgf("Failed marshalling EstafetteCiBuilderEvent for job %v", jobName)
			return err
		}
		requestBody = bytes.NewReader(data)

		// create client, in order to add headers
		client := pester.New()
		client.MaxRetries = 5
		client.Backoff = pester.ExponentialJitterBackoff
		client.KeepLog = true
		client.Timeout = time.Duration(10) * time.Second
		request, err := http.NewRequest("POST", ciServerBuilderEventsURL, requestBody)
		if err != nil {
			log.Error().Err(err).Msgf("Failed creating http client for job %v", jobName)
			return err
		}

		// add headers
		request.Header.Add("X-Estafette-Event", event)
		request.Header.Add("X-Estafette-Event-Job-Name", jobName)
		request.Header.Add("Authorization", fmt.Sprintf("Bearer %v", ciAPIKey))

		// perform actual request
		response, err := client.Do(request)
		if err != nil {
			log.Error().Err(err).Str("logs", client.LogString()).Msgf("Failed performing http request to %v for job %v", ciServerBuilderEventsURL, jobName)
			return err
		}

		defer response.Body.Close()

		log.Debug().Str("logs", client.LogString()).Str("url", ciServerBuilderEventsURL).Msgf("Succesfully sent %v event to api", event)
	}

	return nil
}
<commit_msg>rename pester logs<commit_after>package main

import (
	"bytes"
	"encoding/json"
	"fmt"
	"io"
	"net/http"
	"os"
	"strconv"
	"time"

	"github.com/estafette/estafette-ci-contracts"

	"github.com/rs/zerolog/log"
	"github.com/sethgrid/pester"
)

// EndOfLifeHelper has methods to shutdown the runner after a fatal or successful run
type EndOfLifeHelper interface {
	handleFatal(contracts.BuildLog, error, string)
	sendBuildFinishedEvent(buildStatus string) error
	sendBuildCleanEvent(buildStatus string) error
	sendBuildJobLogEvent(buildLog contracts.BuildLog) error
}

type endOfLifeHelperImpl struct {
	runAsJob bool
	config   contracts.BuilderConfig
}

// NewEndOfLifeHelper returns a new EndOfLifeHelper
func NewEndOfLifeHelper(runAsJob bool, config contracts.BuilderConfig) EndOfLifeHelper {
	return &endOfLifeHelperImpl{
		runAsJob: runAsJob,
		config:   config,
	}
}

func (elh *endOfLifeHelperImpl) handleFatal(buildLog contracts.BuildLog, err error, message string) {

	// add error messages as step to show in logs
	fatalStep := contracts.BuildLogStep{
		Step:     "init",
		LogLines: []contracts.BuildLogLine{},
		ExitCode: -1,
		Status:   "failed",
	}
	lineNumber := 1

	if err != nil {
		fatalStep.LogLines = append(fatalStep.LogLines, contracts.BuildLogLine{
			LineNumber: lineNumber,
			Timestamp:  time.Now().UTC(),
			StreamType: "stderr",
			Text:       err.Error(),
		})
		lineNumber++
	}
	if message != "" {
		fatalStep.LogLines = append(fatalStep.LogLines, contracts.BuildLogLine{
			LineNumber: lineNumber,
			Timestamp:  time.Now().UTC(),
			StreamType: "stderr",
			Text:       message,
		})
		lineNumber++
	}

	buildLog.Steps = append(buildLog.Steps, fatalStep)

	_ = elh.sendBuildFinishedEvent("failed")
	_ = elh.sendBuildJobLogEvent(buildLog)
	_ = elh.sendBuildCleanEvent("failed")

	if elh.runAsJob {
		log.Error().Err(err).Msg(message)
		os.Exit(0)
	} else {
		log.Fatal().Err(err).Msg(message)
	}
}

func (elh *endOfLifeHelperImpl) sendBuildJobLogEvent(buildLog contracts.BuildLog) (err error) {

	ciServerBuilderPostLogsURL := elh.config.CIServer.PostLogsURL
	ciAPIKey := elh.config.CIServer.APIKey
	jobName := *elh.config.JobName

	if ciServerBuilderPostLogsURL != "" && ciAPIKey != "" && jobName != "" {

		// convert BuildJobLogs to json
		var requestBody io.Reader

		var data []byte
		if *elh.config.Action == "release" {
			// copy buildLog to releaseLog and marshal that
			releaseLog := contracts.ReleaseLog{
				ID:         buildLog.ID,
				RepoSource: buildLog.RepoSource,
				RepoOwner:  buildLog.RepoOwner,
				RepoName:   buildLog.RepoName,
				ReleaseID:  strconv.Itoa(elh.config.ReleaseParams.ReleaseID),
				Steps:      buildLog.Steps,
				InsertedAt: buildLog.InsertedAt,
			}
			data, err = json.Marshal(releaseLog)
			if err != nil {
				log.Error().Err(err).Msgf("Failed marshalling ReleaseLog for job %v", jobName)
				return
			}
		} else {
			data, err = json.Marshal(buildLog)
			if err != nil {
				log.Error().Err(err).Msgf("Failed marshalling BuildLog for job %v", jobName)
				return
			}
		}

		requestBody = bytes.NewReader(data)

		// create client, in order to add headers
		client := pester.New()
		client.MaxRetries = 5
		client.Backoff = pester.ExponentialJitterBackoff
		client.KeepLog = true
		client.Timeout = time.Duration(120) * time.Second
		request, err := http.NewRequest("POST", ciServerBuilderPostLogsURL, requestBody)
		if err != nil {
			log.Error().Err(err).Msgf("Failed creating http client for job %v", jobName)
			return err
		}

		// add headers
		request.Header.Add("Authorization", fmt.Sprintf("Bearer %v", ciAPIKey))
		request.Header.Add("Content-Type", "application/json")

		// perform actual request
		response, err := client.Do(request)
		if err != nil {
			log.Error().Err(err).Str("logs", client.LogString()).Msgf("Failed shipping logs to %v for job %v", ciServerBuilderPostLogsURL, jobName)
			return err
		}

		defer response.Body.Close()

		log.Debug().Str("logs", client.LogString()).Msgf("Successfully shipped logs to %v for job %v", ciServerBuilderPostLogsURL, jobName)
	}

	return nil
}

func (elh *endOfLifeHelperImpl) sendBuildFinishedEvent(buildStatus string) error {
	return elh.sendBuilderEvent(buildStatus, fmt.Sprintf("builder:%v", buildStatus))
}

func (elh *endOfLifeHelperImpl) sendBuildCleanEvent(buildStatus string) error {
	return elh.sendBuilderEvent(buildStatus, "builder:clean")
}

func (elh *endOfLifeHelperImpl) sendBuilderEvent(buildStatus, event string) (err error) {

	ciServerBuilderEventsURL := elh.config.CIServer.BuilderEventsURL
	ciAPIKey := elh.config.CIServer.APIKey
	jobName := *elh.config.JobName

	if ciServerBuilderEventsURL != "" && ciAPIKey != "" && jobName != "" {
		// convert EstafetteCiBuilderEvent to json
		var requestBody io.Reader

		buildID := ""
		if elh.config.BuildParams != nil {
			buildID = strconv.Itoa(elh.config.BuildParams.BuildID)
		}

		releaseID := ""
		if elh.config.ReleaseParams != nil {
			releaseID = strconv.Itoa(elh.config.ReleaseParams.ReleaseID)
		}

		ciBuilderEvent := EstafetteCiBuilderEvent{
			JobName:      jobName,
			RepoSource:   elh.config.Git.RepoSource,
			RepoOwner:    elh.config.Git.RepoOwner,
			RepoName:     elh.config.Git.RepoName,
			RepoBranch:   elh.config.Git.RepoBranch,
			RepoRevision: elh.config.Git.RepoRevision,
			ReleaseID:    releaseID,
			BuildID:      buildID,
			BuildStatus:  buildStatus,
		}
		data, err := json.Marshal(ciBuilderEvent)
		if err != nil {
			log.Error().Err(err).Msgf("Failed marshalling EstafetteCiBuilderEvent for job %v", jobName)
			return err
		}
		requestBody = bytes.NewReader(data)

		// create client, in order to add headers
		client := pester.New()
		client.MaxRetries = 5
		client.Backoff = pester.ExponentialJitterBackoff
		client.KeepLog = true
		client.Timeout = time.Duration(10) * time.Second
		request, err := http.NewRequest("POST", ciServerBuilderEventsURL, requestBody)
		if err != nil {
			log.Error().Err(err).Msgf("Failed creating http client for job %v", jobName)
			return err
		}

		// add headers
		request.Header.Add("X-Estafette-Event", event)
		request.Header.Add("X-Estafette-Event-Job-Name", jobName)
		request.Header.Add("Authorization", fmt.Sprintf("Bearer %v", ciAPIKey))

		// perform actual request
		response, err := client.Do(request)
		if err != nil {
			log.Error().Err(err).Str("pesterLogs", client.LogString()).Msgf("Failed performing http request to %v for job %v", ciServerBuilderEventsURL, jobName)
			return err
		}

		defer response.Body.Close()

		log.Debug().Str("pesterLogs", client.LogString()).Str("url", ciServerBuilderEventsURL).Msgf("Succesfully sent %v event to api", event)
	}

	return nil
}
<|endoftext|>


<commit_before>package gophercloud

import "errors"

var (
	// ErrServiceNotFound is returned when no service in a service catalog matches
	// the provided EndpointOpts. This is generally returned by provider service
	// factory methods like "NewComputeV2()" and can mean that a service is not
	// enabled for your account.
	ErrServiceNotFound = errors.New("No suitable service could be found in the service catalog.")

	// ErrEndpointNotFound is returned when no available endpoints match the
	// provided EndpointOpts. This is also generally returned by provider service
	// factory methods, and usually indicates that a region was specified
	// incorrectly.
	ErrEndpointNotFound = errors.New("No suitable endpoint could be found in the service catalog.")
)

// Availability indicates to whom a specific service endpoint is accessible:
// the internet at large, internal networks only, or only to administrators.
// Different identity services use different terminology for these. Identity v2
// lists them as different kinds of URLs within the service catalog ("adminURL",
// "internalURL", and "publicURL"), while v3 lists them as "Interfaces" in an
// endpoints response.
type Availability string

const (
	// AvailabilityAdmin indicates that an endpoint only available to administrators.
	AvailabilityAdmin Availability = "admin"

	// AvailabilityPublic indicates that an endpoint available to everyone on the
	// internet.
	AvailabilityPublic Availability = "public"

	// AvailabilityInternal indicates that an endpoint is only available within
	// the cluster's internal network.
	AvailabilityInternal Availability = "internal"
)

// EndpointOpts specifies search criteria used to by queries against an
// OpenStack service catalog. The options must contain enough information to
// unambiguously identify one, and only one, endpoint within the catalog.
//
// Usually, these are passed to service client factory functions in a provider
// package, like "rackspace.NewComputeV2()".
type EndpointOpts struct {
	// Type [required] is the service type for the client (e.g., "compute",
	// "object-store"). Generally, this will be supplied by the service client
	// function, but a user-given value will be honored if provided.
	Type string

	// Name [optional] is the service name for the client (e.g., "nova") as it
	// appears in the service catalog. Services can have the same Type but a
	// different Name, which is why both Type and Name are sometimes needed.
	Name string

	// Region [required] is the geographic region in which the endpoint resides,
	// generally specifying which datacenter should house your resources.
	// Required only for services that span multiple regions.
	Region string

	// Availability [optional] is the visibility of the endpoint to be returned.
	// Valid types include the constants AvailabilityPublic, AvailabilityInternal,
	// or AvailabilityAdmin from this package.
	//
	// Availability is not required, and defaults to AvailabilityPublic. Not all
	// providers or services offer all Availability options.
	Availability Availability
}

// EndpointLocator is a function that describes how to locate a single endpoint
// from a service catalog for a specific ProviderClient. It should be set
// during ProviderClient authentication and used to discover related
// ServiceClients.
type EndpointLocator func(EndpointOpts) (string, error)

// ApplyDefaults sets EndpointOpts fields if not already set. Currently,
// EndpointOpts.Availability defaults to the public endpoint.
func (eo *EndpointOpts) ApplyDefaults(t string) {
	if eo.Type == "" {
		eo.Type = t
	}
	if eo.Availability == "" {
		eo.Availability = AvailabilityPublic
	}
}
<commit_msg>Explicitly name a few more things as internal.<commit_after>package gophercloud

import "errors"

var (
	// ErrServiceNotFound is returned when no service in a service catalog matches
	// the provided EndpointOpts. This is generally returned by provider service
	// factory methods like "NewComputeV2()" and can mean that a service is not
	// enabled for your account.
	ErrServiceNotFound = errors.New("No suitable service could be found in the service catalog.")

	// ErrEndpointNotFound is returned when no available endpoints match the
	// provided EndpointOpts. This is also generally returned by provider service
	// factory methods, and usually indicates that a region was specified
	// incorrectly.
	ErrEndpointNotFound = errors.New("No suitable endpoint could be found in the service catalog.")
)

// Availability indicates to whom a specific service endpoint is accessible:
// the internet at large, internal networks only, or only to administrators.
// Different identity services use different terminology for these. Identity v2
// lists them as different kinds of URLs within the service catalog ("adminURL",
// "internalURL", and "publicURL"), while v3 lists them as "Interfaces" in an
// endpoints response.
type Availability string

const (
	// AvailabilityAdmin indicates that an endpoint only available to administrators.
	AvailabilityAdmin Availability = "admin"

	// AvailabilityPublic indicates that an endpoint available to everyone on the
	// internet.
	AvailabilityPublic Availability = "public"

	// AvailabilityInternal indicates that an endpoint is only available within
	// the cluster's internal network.
	AvailabilityInternal Availability = "internal"
)

// EndpointOpts specifies search criteria used to by queries against an
// OpenStack service catalog. The options must contain enough information to
// unambiguously identify one, and only one, endpoint within the catalog.
//
// Usually, these are passed to service client factory functions in a provider
// package, like "rackspace.NewComputeV2()".
type EndpointOpts struct {
	// Type [required] is the service type for the client (e.g., "compute",
	// "object-store"). Generally, this will be supplied by the service client
	// function, but a user-given value will be honored if provided.
	Type string

	// Name [optional] is the service name for the client (e.g., "nova") as it
	// appears in the service catalog. Services can have the same Type but a
	// different Name, which is why both Type and Name are sometimes needed.
	Name string

	// Region [required] is the geographic region in which the endpoint resides,
	// generally specifying which datacenter should house your resources.
	// Required only for services that span multiple regions.
	Region string

	// Availability [optional] is the visibility of the endpoint to be returned.
	// Valid types include the constants AvailabilityPublic, AvailabilityInternal,
	// or AvailabilityAdmin from this package.
	//
	// Availability is not required, and defaults to AvailabilityPublic. Not all
	// providers or services offer all Availability options.
	Availability Availability
}

/*
EndpointLocator is an internal function to be used by provider implementations.

It provides an implementation that locates a single endpoint from a service
catalog for a specific ProviderClient based on user-provided EndpointOpts. The
provider then uses it to discover related ServiceClients.
*/
type EndpointLocator func(EndpointOpts) (string, error)

// ApplyDefaults is an internal method to be used by provider implementations.
//
// It sets EndpointOpts fields if not already set, including a default type.
// Currently, EndpointOpts.Availability defaults to the public endpoint.
func (eo *EndpointOpts) ApplyDefaults(t string) {
	if eo.Type == "" {
		eo.Type = t
	}
	if eo.Availability == "" {
		eo.Availability = AvailabilityPublic
	}
}
<|endoftext|>


<commit_before>package models

import (
	"time"

	"gopkg.in/mgo.v2"
	"gopkg.in/mgo.v2/bson"
)

var (
	Todo *_Todo

	todoCollection = "todo"
	todoIndexes    = []mgo.Index{}
)

type _Todo struct{}

type TodoModel struct {
	Id        bson.ObjectId `bson:"_id" json:"-"`
	Uid       string        `bson:"uid" json:"uid"`
	Title     string        `bson:"title" json:"title"`
	Content   string        `bson:"content" json:"content"`
	Due       time.Time     `bson:"due" json:"due"`
	Finished  bool          `bson:"finished" json:"finished"`
	CreatedAt time.Time     `bson:"created_at" json:"created_at"`
	UpdatedAt time.Time     `bson:"updated_at" json:"updated_at"`

	isNewRecord bool `bson:"-"`
}

func (_ *_Todo) NewModel(uid, title, content string, due time.Time) *TodoModel {
	return &TodoModel{
		Id:          bson.NewObjectId(),
		Uid:         uid,
		Title:       title,
		Content:     content,
		Due:         due,
		Finished:    false,
		isNewRecord: true,
	}
}

func (todo *TodoModel) Save() (err error) {
	if !todo.Id.Valid() {
		return ErrInvalidId
	}

	Todo.Query(func(c *mgo.Collection) {
		t := time.Now()

		if todo.isNewRecord {
			todo.CreatedAt = t
			todo.UpdatedAt = t

			if err = c.Insert(todo); err == nil {
				todo.isNewRecord = false
			}
		} else {
			migration := bson.M{
				"$set": bson.M{
					"title":      todo.Title,
					"content":    todo.Content,
					"due":        todo.Due,
					"finished":   todo.Finished,
					"updated_at": t,
				},
			}

			err = c.UpdateId(todo.Id, migration)
		}
	})

	return
}

type ListTodoParams struct {
	Page       int
	Limit      int
	From       time.Time
	To         time.Time
	Uid        string
	IsAll      bool
	IsFinished bool
}

func (_ *_Todo) List(params ListTodoParams) (total int, todos []*TodoModel, err error) {
	offset := (params.Page - 1) * params.Limit
	if offset < 0 {
		offset = 0
	}

	query := bson.M{}
	if params.From.After(params.To) {
		return 0, nil, ErrInvalidParams
	}

	dueRange := bson.M{}
	if !params.From.IsZero() {
		dueRange["$gte"] = params.From
	}
	if !params.To.IsZero() {
		dueRange["$lte"] = params.To
	}
	if len(dueRange) > 0 {
		query["due"] = dueRange
	}

	if params.Uid != "" {
		query["uid"] = params.Uid
	}

	if !params.IsAll {
		query["finished"] = params.IsFinished
	}

	Todo.Query(func(c *mgo.Collection) {
		err = c.Find(query).Skip(offset).Limit(params.Limit).Sort("due").All(&todos)
		if err == nil {
			total, err = c.Find(query).Count()
		}
	})

	return
}

func (_ *_Todo) Delete(id string) (err error) {
	if !bson.IsObjectIdHex(id) {
		return ErrInvalidId
	}

	Todo.Query(func(c *mgo.Collection) {
		err = c.RemoveId(bson.ObjectIdHex(id))
	})

	return
}

func (_ *_Todo) Query(query func(c *mgo.Collection)) {
	Mongo().Query(tomatoCollection, tomatoIndexes, query)
}
<commit_msg>fix test fail<commit_after>package models

import (
	"time"

	"gopkg.in/mgo.v2"
	"gopkg.in/mgo.v2/bson"
)

var (
	Todo *_Todo

	todoCollection = "todo"
	todoIndexes    = []mgo.Index{}
)

type _Todo struct{}

type TodoModel struct {
	Id        bson.ObjectId `bson:"_id" json:"-"`
	Uid       string        `bson:"uid" json:"uid"`
	Title     string        `bson:"title" json:"title"`
	Content   string        `bson:"content" json:"content"`
	Due       time.Time     `bson:"due" json:"due"`
	Finished  bool          `bson:"finished" json:"finished"`
	CreatedAt time.Time     `bson:"created_at" json:"created_at"`
	UpdatedAt time.Time     `bson:"updated_at" json:"updated_at"`

	isNewRecord bool `bson:"-"`
}

func (_ *_Todo) NewModel(uid, title, content string, due time.Time) *TodoModel {
	return &TodoModel{
		Id:          bson.NewObjectId(),
		Uid:         uid,
		Title:       title,
		Content:     content,
		Due:         due,
		Finished:    false,
		isNewRecord: true,
	}
}

func (todo *TodoModel) Save() (err error) {
	if !todo.Id.Valid() {
		return ErrInvalidId
	}

	Todo.Query(func(c *mgo.Collection) {
		t := time.Now()

		if todo.isNewRecord {
			todo.CreatedAt = t
			todo.UpdatedAt = t

			if err = c.Insert(todo); err == nil {
				todo.isNewRecord = false
			}
		} else {
			migration := bson.M{
				"$set": bson.M{
					"title":      todo.Title,
					"content":    todo.Content,
					"due":        todo.Due,
					"finished":   todo.Finished,
					"updated_at": t,
				},
			}

			err = c.UpdateId(todo.Id, migration)
		}
	})

	return
}

type ListTodoParams struct {
	Page       int
	Limit      int
	From       time.Time
	To         time.Time
	Uid        string
	IsAll      bool
	IsFinished bool
}

func (_ *_Todo) List(params ListTodoParams) (total int, todos []*TodoModel, err error) {
	offset := (params.Page - 1) * params.Limit
	if offset < 0 {
		offset = 0
	}

	query := bson.M{}
	if params.From.After(params.To) {
		return 0, nil, ErrInvalidParams
	}

	dueRange := bson.M{}
	if !params.From.IsZero() {
		dueRange["$gte"] = params.From
	}
	if !params.To.IsZero() {
		dueRange["$lte"] = params.To
	}
	if len(dueRange) > 0 {
		query["due"] = dueRange
	}

	if params.Uid != "" {
		query["uid"] = params.Uid
	}

	if !params.IsAll {
		query["finished"] = params.IsFinished
	}

	Todo.Query(func(c *mgo.Collection) {
		err = c.Find(query).Skip(offset).Limit(params.Limit).Sort("due").All(&todos)
		if err == nil {
			total, err = c.Find(query).Count()
		}
	})

	return
}

func (_ *_Todo) Delete(id string) (err error) {
	if !bson.IsObjectIdHex(id) {
		return ErrInvalidId
	}

	Todo.Query(func(c *mgo.Collection) {
		err = c.RemoveId(bson.ObjectIdHex(id))
	})

	return
}

func (_ *_Todo) Query(query func(c *mgo.Collection)) {
	Mongo().Query(todoCollection, todoIndexes, query)
}
<|endoftext|>


<commit_before>//  Copyright (c) 2013 Couchbase, Inc.
//  Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file
//  except in compliance with the License. You may obtain a copy of the License at
//    http://www.apache.org/licenses/LICENSE-2.0
//  Unless required by applicable law or agreed to in writing, software distributed under the
//  License is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND,
//  either express or implied. See the License for the specific language governing permissions
//  and limitations under the License.

package base

import (
	"fmt"
	"regexp"

	"github.com/couchbaselabs/go-couchbase"
	"github.com/couchbaselabs/walrus"
	"github.com/dustin/gomemcached"
	"github.com/dustin/gomemcached/client"
)

func init() {
	// Increase max memcached request size to 20M bytes, to support large docs (attachments!)
	// arriving in a tap feed. (see issues #210, #333.)
	gomemcached.MaxBodyLen = int(20.0e6)
}

type Bucket walrus.Bucket
type TapArguments walrus.TapArguments
type TapFeed walrus.TapFeed
type AuthHandler couchbase.AuthHandler

// Full specification of how to connect to a bucket
type BucketSpec struct {
	Server, PoolName, BucketName string
	Auth                         AuthHandler
}

// Implementation of walrus.Bucket that talks to a Couchbase server
type couchbaseBucket struct {
	*couchbase.Bucket
}

type couchbaseFeedImpl struct {
	*couchbase.TapFeed
	events <-chan walrus.TapEvent
}

func (feed *couchbaseFeedImpl) Events() <-chan walrus.TapEvent {
	return feed.events
}

func (bucket couchbaseBucket) GetName() string {
	return bucket.Name
}

func (bucket couchbaseBucket) Write(k string, flags int, exp int, v interface{}, opt walrus.WriteOptions) (err error) {
	return bucket.Bucket.Write(k, flags, exp, v, couchbase.WriteOptions(opt))
}

func (bucket couchbaseBucket) Update(k string, exp int, callback walrus.UpdateFunc) error {
	return bucket.Bucket.Update(k, exp, couchbase.UpdateFunc(callback))
}

func (bucket couchbaseBucket) WriteUpdate(k string, exp int, callback walrus.WriteUpdateFunc) error {
	cbCallback := func(current []byte) (updated []byte, opt couchbase.WriteOptions, err error) {
		updated, walrusOpt, err := callback(current)
		opt = couchbase.WriteOptions(walrusOpt)
		return
	}
	return bucket.Bucket.WriteUpdate(k, exp, cbCallback)
}

func (bucket couchbaseBucket) View(ddoc, name string, params map[string]interface{}) (walrus.ViewResult, error) {
	vres := walrus.ViewResult{}
	return vres, bucket.Bucket.ViewCustom(ddoc, name, params, &vres)
}

func (bucket couchbaseBucket) StartTapFeed(args walrus.TapArguments) (walrus.TapFeed, error) {
	cbArgs := memcached.TapArguments{
		Backfill: args.Backfill,
		Dump:     args.Dump,
		KeysOnly: args.KeysOnly,
	}
	cbFeed, err := bucket.Bucket.StartTapFeed(&cbArgs)
	if err != nil {
		return nil, err
	}

	// Create a bridge from the Couchbase tap feed to a Walrus tap feed:
	events := make(chan walrus.TapEvent)
	impl := couchbaseFeedImpl{cbFeed, events}
	go func() {
		for cbEvent := range cbFeed.C {
			events <- walrus.TapEvent{
				Opcode:   walrus.TapOpcode(cbEvent.Opcode),
				Expiry:   cbEvent.Expiry,
				Flags:    cbEvent.Flags,
				Key:      cbEvent.Key,
				Value:    cbEvent.Value,
				Sequence: cbEvent.Cas,
			}
		}
	}()
	return &impl, nil
}

func (bucket couchbaseBucket) Dump() {
	Warn("Dump not implemented for couchbaseBucket")
}

// Creates a Bucket that talks to a real live Couchbase server.
func GetCouchbaseBucket(spec BucketSpec) (bucket Bucket, err error) {
        cbbucket, err := couchbase.GetBucket(spec.Server, spec.PoolName, spec.BucketName)
        if err == nil {
                bucket = couchbaseBucket{cbbucket}
        }
        return
}
func GetBucket(spec BucketSpec) (bucket Bucket, err error) {
	if isWalrus, _ := regexp.MatchString(`^(walrus:|file:|/|\.)`, spec.Server); isWalrus {
		Log("Opening Walrus database %s on <%s>", spec.BucketName, spec.Server)
		walrus.Logging = LogKeys["Walrus"]
		bucket, err = walrus.GetBucket(spec.Server, spec.PoolName, spec.BucketName)
	} else {
		suffix := ""
		if spec.Auth != nil {
			username, _ := spec.Auth.GetCredentials()
			suffix = fmt.Sprintf(" as user %q", username)
		}
		Log("Opening Couchbase database %s on <%s>%s", spec.BucketName, spec.Server, suffix)
		bucket, err = GetCouchbaseBucket(spec)
	}

	if LogKeys["Bucket"] {
		bucket = &LoggingBucket{bucket: bucket}
	}
	return
}
<commit_msg>Ran file through Fmt<commit_after>//  Copyright (c) 2013 Couchbase, Inc.
//  Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file
//  except in compliance with the License. You may obtain a copy of the License at
//    http://www.apache.org/licenses/LICENSE-2.0
//  Unless required by applicable law or agreed to in writing, software distributed under the
//  License is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND,
//  either express or implied. See the License for the specific language governing permissions
//  and limitations under the License.

package base

import (
	"fmt"
	"regexp"

	"github.com/couchbaselabs/go-couchbase"
	"github.com/couchbaselabs/walrus"
	"github.com/dustin/gomemcached"
	"github.com/dustin/gomemcached/client"
)

func init() {
	// Increase max memcached request size to 20M bytes, to support large docs (attachments!)
	// arriving in a tap feed. (see issues #210, #333.)
	gomemcached.MaxBodyLen = int(20.0e6)
}

type Bucket walrus.Bucket
type TapArguments walrus.TapArguments
type TapFeed walrus.TapFeed
type AuthHandler couchbase.AuthHandler

// Full specification of how to connect to a bucket
type BucketSpec struct {
	Server, PoolName, BucketName string
	Auth                         AuthHandler
}

// Implementation of walrus.Bucket that talks to a Couchbase server
type couchbaseBucket struct {
	*couchbase.Bucket
}

type couchbaseFeedImpl struct {
	*couchbase.TapFeed
	events <-chan walrus.TapEvent
}

func (feed *couchbaseFeedImpl) Events() <-chan walrus.TapEvent {
	return feed.events
}

func (bucket couchbaseBucket) GetName() string {
	return bucket.Name
}

func (bucket couchbaseBucket) Write(k string, flags int, exp int, v interface{}, opt walrus.WriteOptions) (err error) {
	return bucket.Bucket.Write(k, flags, exp, v, couchbase.WriteOptions(opt))
}

func (bucket couchbaseBucket) Update(k string, exp int, callback walrus.UpdateFunc) error {
	return bucket.Bucket.Update(k, exp, couchbase.UpdateFunc(callback))
}

func (bucket couchbaseBucket) WriteUpdate(k string, exp int, callback walrus.WriteUpdateFunc) error {
	cbCallback := func(current []byte) (updated []byte, opt couchbase.WriteOptions, err error) {
		updated, walrusOpt, err := callback(current)
		opt = couchbase.WriteOptions(walrusOpt)
		return
	}
	return bucket.Bucket.WriteUpdate(k, exp, cbCallback)
}

func (bucket couchbaseBucket) View(ddoc, name string, params map[string]interface{}) (walrus.ViewResult, error) {
	vres := walrus.ViewResult{}
	return vres, bucket.Bucket.ViewCustom(ddoc, name, params, &vres)
}

func (bucket couchbaseBucket) StartTapFeed(args walrus.TapArguments) (walrus.TapFeed, error) {
	cbArgs := memcached.TapArguments{
		Backfill: args.Backfill,
		Dump:     args.Dump,
		KeysOnly: args.KeysOnly,
	}
	cbFeed, err := bucket.Bucket.StartTapFeed(&cbArgs)
	if err != nil {
		return nil, err
	}

	// Create a bridge from the Couchbase tap feed to a Walrus tap feed:
	events := make(chan walrus.TapEvent)
	impl := couchbaseFeedImpl{cbFeed, events}
	go func() {
		for cbEvent := range cbFeed.C {
			events <- walrus.TapEvent{
				Opcode:   walrus.TapOpcode(cbEvent.Opcode),
				Expiry:   cbEvent.Expiry,
				Flags:    cbEvent.Flags,
				Key:      cbEvent.Key,
				Value:    cbEvent.Value,
				Sequence: cbEvent.Cas,
			}
		}
	}()
	return &impl, nil
}

func (bucket couchbaseBucket) Dump() {
	Warn("Dump not implemented for couchbaseBucket")
}

// Creates a Bucket that talks to a real live Couchbase server.
func GetCouchbaseBucket(spec BucketSpec) (bucket Bucket, err error) {
	cbbucket, err := couchbase.GetBucket(spec.Server, spec.PoolName, spec.BucketName)
	if err == nil {
		bucket = couchbaseBucket{cbbucket}
	}
	return
}

func GetBucket(spec BucketSpec) (bucket Bucket, err error) {
	if isWalrus, _ := regexp.MatchString(`^(walrus:|file:|/|\.)`, spec.Server); isWalrus {
		Log("Opening Walrus database %s on <%s>", spec.BucketName, spec.Server)
		walrus.Logging = LogKeys["Walrus"]
		bucket, err = walrus.GetBucket(spec.Server, spec.PoolName, spec.BucketName)
	} else {
		suffix := ""
		if spec.Auth != nil {
			username, _ := spec.Auth.GetCredentials()
			suffix = fmt.Sprintf(" as user %q", username)
		}
		Log("Opening Couchbase database %s on <%s>%s", spec.BucketName, spec.Server, suffix)
		bucket, err = GetCouchbaseBucket(spec)
	}

	if LogKeys["Bucket"] {
		bucket = &LoggingBucket{bucket: bucket}
	}
	return
}
<|endoftext|>


<commit_before>// +build linux

package netns

import (
	"fmt"
	"io/ioutil"
	"os"
	"path/filepath"
	"strconv"
	"strings"

	"golang.org/x/sys/unix"
)

// Deprecated: use syscall pkg instead (go >= 1.5 needed).
const (
	CLONE_NEWUTS  = 0x04000000 /* New utsname group? */
	CLONE_NEWIPC  = 0x08000000 /* New ipcs */
	CLONE_NEWUSER = 0x10000000 /* New user namespace */
	CLONE_NEWPID  = 0x20000000 /* New pid namespace */
	CLONE_NEWNET  = 0x40000000 /* New network namespace */
	CLONE_IO      = 0x80000000 /* Get io context */
)

// Setns sets namespace using syscall. Note that this should be a method
// in syscall but it has not been added.
func Setns(ns NsHandle, nstype int) (err error) {
	return unix.Setns(int(ns), nstype)
}

// Set sets the current network namespace to the namespace represented
// by NsHandle.
func Set(ns NsHandle) (err error) {
	return Setns(ns, CLONE_NEWNET)
}

// New creates a new network namespace, sets it as current and returns
// a handle to it.
func New() (ns NsHandle, err error) {
	if err := unix.Unshare(CLONE_NEWNET); err != nil {
		return -1, err
	}
	return Get()
}

// Get gets a handle to the current threads network namespace.
func Get() (NsHandle, error) {
	return GetFromThread(os.Getpid(), unix.Gettid())
}

// GetFromPath gets a handle to a network namespace
// identified by the path
func GetFromPath(path string) (NsHandle, error) {
	fd, err := unix.Open(path, unix.O_RDONLY|unix.O_CLOEXEC, 0)
	if err != nil {
		return -1, err
	}
	return NsHandle(fd), nil
}

// GetFromName gets a handle to a named network namespace such as one
// created by `ip netns add`.
func GetFromName(name string) (NsHandle, error) {
	return GetFromPath(fmt.Sprintf("/var/run/netns/%s", name))
}

// GetFromPid gets a handle to the network namespace of a given pid.
func GetFromPid(pid int) (NsHandle, error) {
	return GetFromPath(fmt.Sprintf("/proc/%d/ns/net", pid))
}

// GetFromThread gets a handle to the network namespace of a given pid and tid.
func GetFromThread(pid, tid int) (NsHandle, error) {
	return GetFromPath(fmt.Sprintf("/proc/%d/task/%d/ns/net", pid, tid))
}

// GetFromDocker gets a handle to the network namespace of a docker container.
// Id is prefixed matched against the running docker containers, so a short
// identifier can be used as long as it isn't ambiguous.
func GetFromDocker(id string) (NsHandle, error) {
	pid, err := getPidForContainer(id)
	if err != nil {
		return -1, err
	}
	return GetFromPid(pid)
}

// borrowed from docker/utils/utils.go
func findCgroupMountpoint(cgroupType string) (string, error) {
	output, err := ioutil.ReadFile("/proc/mounts")
	if err != nil {
		return "", err
	}

	// /proc/mounts has 6 fields per line, one mount per line, e.g.
	// cgroup /sys/fs/cgroup/devices cgroup rw,relatime,devices 0 0
	for _, line := range strings.Split(string(output), "\n") {
		parts := strings.Split(line, " ")
		if len(parts) == 6 && parts[2] == "cgroup" {
			for _, opt := range strings.Split(parts[3], ",") {
				if opt == cgroupType {
					return parts[1], nil
				}
			}
		}
	}

	return "", fmt.Errorf("cgroup mountpoint not found for %s", cgroupType)
}

// Returns the relative path to the cgroup docker is running in.
// borrowed from docker/utils/utils.go
// modified to get the docker pid instead of using /proc/self
func getThisCgroup(cgroupType string) (string, error) {
	dockerpid, err := ioutil.ReadFile("/var/run/docker.pid")
	if err != nil {
		return "", err
	}
	result := strings.Split(string(dockerpid), "\n")
	if len(result) == 0 || len(result[0]) == 0 {
		return "", fmt.Errorf("docker pid not found in /var/run/docker.pid")
	}
	pid, err := strconv.Atoi(result[0])
	if err != nil {
		return "", err
	}
	output, err := ioutil.ReadFile(fmt.Sprintf("/proc/%d/cgroup", pid))
	if err != nil {
		return "", err
	}
	for _, line := range strings.Split(string(output), "\n") {
		parts := strings.Split(line, ":")
		// any type used by docker should work
		if parts[1] == cgroupType {
			return parts[2], nil
		}
	}
	return "", fmt.Errorf("cgroup '%s' not found in /proc/%d/cgroup", cgroupType, pid)
}

// Returns the first pid in a container.
// borrowed from docker/utils/utils.go
// modified to only return the first pid
// modified to glob with id
// modified to search for newer docker containers
func getPidForContainer(id string) (int, error) {
	pid := 0

	// memory is chosen randomly, any cgroup used by docker works
	cgroupType := "memory"

	cgroupRoot, err := findCgroupMountpoint(cgroupType)
	if err != nil {
		return pid, err
	}

	cgroupThis, err := getThisCgroup(cgroupType)
	if err != nil {
		return pid, err
	}

	id += "*"

	attempts := []string{
		filepath.Join(cgroupRoot, cgroupThis, id, "tasks"),
		// With more recent lxc versions use, cgroup will be in lxc/
		filepath.Join(cgroupRoot, cgroupThis, "lxc", id, "tasks"),
		// With more recent docker, cgroup will be in docker/
		filepath.Join(cgroupRoot, cgroupThis, "docker", id, "tasks"),
		// Even more recent docker versions under systemd use docker-<id>.scope/
		filepath.Join(cgroupRoot, "system.slice", "docker-"+id+".scope", "tasks"),
		// Even more recent docker versions under cgroup/systemd/docker/<id>/
		filepath.Join(cgroupRoot, "..", "systemd", "docker", id, "tasks"),
		// Kubernetes with docker and CNI is even more different
		filepath.Join(cgroupRoot, "..", "systemd", "kubepods", "*", "pod*", id, "tasks"),
		// Another flavor of containers location in recent kubernetes 1.11+
		filepath.Join(cgroupRoot, cgroupThis, "kubepods.slice", "kubepods-besteffort.slice", "*", "docker-"+id+".scope", "tasks"),
		// When runs inside of a container with recent kubernetes 1.11+
		filepath.Join(cgroupRoot, "kubepods.slice", "kubepods-besteffort.slice", "*", "docker-"+id+".scope", "tasks"),
	}

	var filename string
	for _, attempt := range attempts {
		filenames, _ := filepath.Glob(attempt)
		if len(filenames) > 1 {
			return pid, fmt.Errorf("Ambiguous id supplied: %v", filenames)
		} else if len(filenames) == 1 {
			filename = filenames[0]
			break
		}
	}

	if filename == "" {
		return pid, fmt.Errorf("Unable to find container: %v", id[:len(id)-1])
	}

	output, err := ioutil.ReadFile(filename)
	if err != nil {
		return pid, err
	}

	result := strings.Split(string(output), "\n")
	if len(result) == 0 || len(result[0]) == 0 {
		return pid, fmt.Errorf("No pid found for container")
	}

	pid, err = strconv.Atoi(result[0])
	if err != nil {
		return pid, fmt.Errorf("Invalid pid '%s': %s", result[0], err)
	}

	return pid, nil
}
<commit_msg>Allows to create and delete named network namespaces<commit_after>// +build linux

package netns

import (
	"fmt"
	"io/ioutil"
	"os"
	"path"
	"path/filepath"
	"strconv"
	"strings"
	"syscall"

	"golang.org/x/sys/unix"
)

// Deprecated: use syscall pkg instead (go >= 1.5 needed).
const (
	CLONE_NEWUTS  = 0x04000000   /* New utsname group? */
	CLONE_NEWIPC  = 0x08000000   /* New ipcs */
	CLONE_NEWUSER = 0x10000000   /* New user namespace */
	CLONE_NEWPID  = 0x20000000   /* New pid namespace */
	CLONE_NEWNET  = 0x40000000   /* New network namespace */
	CLONE_IO      = 0x80000000   /* Get io context */
	bindMountPath = "/run/netns" /* Bind mount path for named netns */
)

// Setns sets namespace using syscall. Note that this should be a method
// in syscall but it has not been added.
func Setns(ns NsHandle, nstype int) (err error) {
	return unix.Setns(int(ns), nstype)
}

// Set sets the current network namespace to the namespace represented
// by NsHandle.
func Set(ns NsHandle) (err error) {
	return Setns(ns, CLONE_NEWNET)
}

// New creates a new network namespace, sets it as current and returns
// a handle to it.
func New() (ns NsHandle, err error) {
	if err := unix.Unshare(CLONE_NEWNET); err != nil {
		return -1, err
	}
	return Get()
}

// NewNamed creates a new named network namespace and returns a handle to it
func NewNamed(name string) (NsHandle, error) {
	if _, err := os.Stat(bindMountPath); os.IsNotExist(err) {
		err = os.MkdirAll(bindMountPath, 0755)
		if err != nil {
			return None(), err
		}
	}

	newNs, err := New()
	if err != nil {
		return None(), err
	}

	namedPath := path.Join(bindMountPath, name)

	f, err := os.OpenFile(namedPath, os.O_CREATE|os.O_EXCL, 0444)
	if err != nil {
		return None(), err
	}
	f.Close()

	nsPath := fmt.Sprintf("/proc/%d/task/%d/ns/net", os.Getpid(), syscall.Gettid())
	err = syscall.Mount(nsPath, namedPath, "bind", syscall.MS_BIND, "")
	if err != nil {
		return None(), err
	}

	return newNs, nil
}

// DeleteNamed deletes a named network namespace
func DeleteNamed(name string) error {
	namedPath := path.Join(bindMountPath, name)

	err := syscall.Unmount(namedPath, syscall.MNT_DETACH)
	if err != nil {
		return err
	}

	return os.Remove(namedPath)
}

// Get gets a handle to the current threads network namespace.
func Get() (NsHandle, error) {
	return GetFromThread(os.Getpid(), unix.Gettid())
}

// GetFromPath gets a handle to a network namespace
// identified by the path
func GetFromPath(path string) (NsHandle, error) {
	fd, err := unix.Open(path, unix.O_RDONLY|unix.O_CLOEXEC, 0)
	if err != nil {
		return -1, err
	}
	return NsHandle(fd), nil
}

// GetFromName gets a handle to a named network namespace such as one
// created by `ip netns add`.
func GetFromName(name string) (NsHandle, error) {
	return GetFromPath(fmt.Sprintf("/var/run/netns/%s", name))
}

// GetFromPid gets a handle to the network namespace of a given pid.
func GetFromPid(pid int) (NsHandle, error) {
	return GetFromPath(fmt.Sprintf("/proc/%d/ns/net", pid))
}

// GetFromThread gets a handle to the network namespace of a given pid and tid.
func GetFromThread(pid, tid int) (NsHandle, error) {
	return GetFromPath(fmt.Sprintf("/proc/%d/task/%d/ns/net", pid, tid))
}

// GetFromDocker gets a handle to the network namespace of a docker container.
// Id is prefixed matched against the running docker containers, so a short
// identifier can be used as long as it isn't ambiguous.
func GetFromDocker(id string) (NsHandle, error) {
	pid, err := getPidForContainer(id)
	if err != nil {
		return -1, err
	}
	return GetFromPid(pid)
}

// borrowed from docker/utils/utils.go
func findCgroupMountpoint(cgroupType string) (string, error) {
	output, err := ioutil.ReadFile("/proc/mounts")
	if err != nil {
		return "", err
	}

	// /proc/mounts has 6 fields per line, one mount per line, e.g.
	// cgroup /sys/fs/cgroup/devices cgroup rw,relatime,devices 0 0
	for _, line := range strings.Split(string(output), "\n") {
		parts := strings.Split(line, " ")
		if len(parts) == 6 && parts[2] == "cgroup" {
			for _, opt := range strings.Split(parts[3], ",") {
				if opt == cgroupType {
					return parts[1], nil
				}
			}
		}
	}

	return "", fmt.Errorf("cgroup mountpoint not found for %s", cgroupType)
}

// Returns the relative path to the cgroup docker is running in.
// borrowed from docker/utils/utils.go
// modified to get the docker pid instead of using /proc/self
func getThisCgroup(cgroupType string) (string, error) {
	dockerpid, err := ioutil.ReadFile("/var/run/docker.pid")
	if err != nil {
		return "", err
	}
	result := strings.Split(string(dockerpid), "\n")
	if len(result) == 0 || len(result[0]) == 0 {
		return "", fmt.Errorf("docker pid not found in /var/run/docker.pid")
	}
	pid, err := strconv.Atoi(result[0])
	if err != nil {
		return "", err
	}
	output, err := ioutil.ReadFile(fmt.Sprintf("/proc/%d/cgroup", pid))
	if err != nil {
		return "", err
	}
	for _, line := range strings.Split(string(output), "\n") {
		parts := strings.Split(line, ":")
		// any type used by docker should work
		if parts[1] == cgroupType {
			return parts[2], nil
		}
	}
	return "", fmt.Errorf("cgroup '%s' not found in /proc/%d/cgroup", cgroupType, pid)
}

// Returns the first pid in a container.
// borrowed from docker/utils/utils.go
// modified to only return the first pid
// modified to glob with id
// modified to search for newer docker containers
func getPidForContainer(id string) (int, error) {
	pid := 0

	// memory is chosen randomly, any cgroup used by docker works
	cgroupType := "memory"

	cgroupRoot, err := findCgroupMountpoint(cgroupType)
	if err != nil {
		return pid, err
	}

	cgroupThis, err := getThisCgroup(cgroupType)
	if err != nil {
		return pid, err
	}

	id += "*"

	attempts := []string{
		filepath.Join(cgroupRoot, cgroupThis, id, "tasks"),
		// With more recent lxc versions use, cgroup will be in lxc/
		filepath.Join(cgroupRoot, cgroupThis, "lxc", id, "tasks"),
		// With more recent docker, cgroup will be in docker/
		filepath.Join(cgroupRoot, cgroupThis, "docker", id, "tasks"),
		// Even more recent docker versions under systemd use docker-<id>.scope/
		filepath.Join(cgroupRoot, "system.slice", "docker-"+id+".scope", "tasks"),
		// Even more recent docker versions under cgroup/systemd/docker/<id>/
		filepath.Join(cgroupRoot, "..", "systemd", "docker", id, "tasks"),
		// Kubernetes with docker and CNI is even more different
		filepath.Join(cgroupRoot, "..", "systemd", "kubepods", "*", "pod*", id, "tasks"),
		// Another flavor of containers location in recent kubernetes 1.11+
		filepath.Join(cgroupRoot, cgroupThis, "kubepods.slice", "kubepods-besteffort.slice", "*", "docker-"+id+".scope", "tasks"),
		// When runs inside of a container with recent kubernetes 1.11+
		filepath.Join(cgroupRoot, "kubepods.slice", "kubepods-besteffort.slice", "*", "docker-"+id+".scope", "tasks"),
	}

	var filename string
	for _, attempt := range attempts {
		filenames, _ := filepath.Glob(attempt)
		if len(filenames) > 1 {
			return pid, fmt.Errorf("Ambiguous id supplied: %v", filenames)
		} else if len(filenames) == 1 {
			filename = filenames[0]
			break
		}
	}

	if filename == "" {
		return pid, fmt.Errorf("Unable to find container: %v", id[:len(id)-1])
	}

	output, err := ioutil.ReadFile(filename)
	if err != nil {
		return pid, err
	}

	result := strings.Split(string(output), "\n")
	if len(result) == 0 || len(result[0]) == 0 {
		return pid, fmt.Errorf("No pid found for container")
	}

	pid, err = strconv.Atoi(result[0])
	if err != nil {
		return pid, fmt.Errorf("Invalid pid '%s': %s", result[0], err)
	}

	return pid, nil
}
<|endoftext|>


<commit_before>package amqp

import (
	"errors"
	"fmt"
	"github.com/eventials/goevents/messaging"
	log "github.com/sirupsen/logrus"
	amqplib "github.com/streadway/amqp"
	"regexp"
	"strings"
	"sync"
	"time"
)

var (
	logger = log.WithFields(log.Fields{
		"type":     "goevents",
		"sub_type": "consumer",
	})
)

type handler struct {
	action       string
	fn           messaging.EventHandler
	re           *regexp.Regexp
	maxRetries   int32
	retryDelay   time.Duration
	delayedRetry bool
}

type consumer struct {
	config ConsumerConfig

	m  sync.Mutex
	wg sync.WaitGroup

	conn     *Connection
	autoAck  bool
	handlers []handler

	channel    *amqplib.Channel
	queue      *amqplib.Queue
	retryQueue *amqplib.Queue

	exchangeName string
	queueName    string

	closed bool
}

// ConsumerConfig to be used when creating a new producer.
type ConsumerConfig struct {
	ConsumeRetryInterval time.Duration
	PrefetchCount        int
	DurableQueue         bool
	AutoDelete           bool
	PrefixName           string
}

// NewConsumer returns a new AMQP Consumer.
// Uses a default ConsumerConfig with 2 second of consume retry interval.
func NewConsumer(c messaging.Connection, autoAck bool, exchange, queue string) (*consumer, error) {
	return NewConsumerConfig(c, autoAck, exchange, queue, ConsumerConfig{
		ConsumeRetryInterval: 2 * time.Second,
		PrefetchCount:        0,
		DurableQueue:         true,
		AutoDelete:           false,
	})
}

// NewConsumerConfig returns a new AMQP Consumer.
func NewConsumerConfig(c messaging.Connection, autoAck bool, exchange, queue string, config ConsumerConfig) (*consumer, error) {
	consumer := &consumer{
		config:       config,
		conn:         c.(*Connection),
		autoAck:      autoAck,
		handlers:     make([]handler, 0),
		exchangeName: exchange,
		queueName:    queue,
	}

	err := consumer.setupTopology()

	if err != nil {
		return nil, err
	}

	go consumer.handleReestablishedConnnection()

	return consumer, err

}

func (c *consumer) Close() {
	func() {
		c.m.Lock()
		defer c.m.Unlock()

		// Unsubscribe all handlers
		c.handlers = make([]handler, 0)

		c.closed = true
		c.channel.Close()
	}()

	// Wait all go routine finish.
	c.wg.Wait()
}

func (c *consumer) uniqueNameWithPrefix() string {
	return fmt.Sprint("%s%d", c.config.PrefixName, time.Now().UnixNano())
}

func (c *consumer) setupTopology() error {
	c.m.Lock()
	defer c.m.Unlock()

	var err error

	if c.channel != nil {
		c.channel.Close()
	}

	c.channel, err = c.conn.OpenChannel()

	if err != nil {
		return err
	}

	err = c.channel.Qos(c.config.PrefetchCount, 0, true)

	if err != nil {
		return err
	}

	err = c.channel.ExchangeDeclare(
		c.exchangeName, // name
		"topic",        // type
		true,           // durable
		false,          // auto-delete
		false,          // internal
		false,          // no-wait
		nil,            // arguments
	)

	if err != nil {
		return err
	}

	if c.config.PrefixName != "" {
		c.queueName = c.uniqueNameWithPrefix()
	}

	q, err := c.channel.QueueDeclare(
		c.queueName,           // name
		c.config.DurableQueue, // durable
		c.config.AutoDelete,   // auto-delete
		false,                 // exclusive
		false,                 // no-wait
		nil,                   // arguments
	)

	if err != nil {
		return err
	}

	c.queue = &q

	return nil
}

func (c *consumer) handleReestablishedConnnection() {
	for !c.closed {
		<-c.conn.NotifyReestablish()

		err := c.setupTopology()

		if err != nil {
			logger.WithFields(log.Fields{
				"error": err,
			}).Error("Error setting up topology after reconnection.")
		}
	}
}

func (c *consumer) dispatch(msg amqplib.Delivery) {
	if h, ok := c.getHandler(msg); ok {
		delay, isRetry := getXRetryDelayHeader(msg)

		if isRetry {
			logger.WithFields(log.Fields{
				"delay":      delay.String(),
				"message_id": msg.MessageId,
			}).Info("Delaying message.")

			time.Sleep(delay)
		}

		retryCount, _ := getXRetryCountHeader(msg)

		c.doDispatch(msg, h, retryCount, delay)
	} else {
		if !c.autoAck {
			err := msg.Nack(false, true)

			if err != nil {
				logger.WithFields(log.Fields{
					"error":      err,
					"message_id": msg.MessageId,
				}).Error("Failed to nack message.")
			}
		}
	}
}

func (c *consumer) callAndHandlePanic(msg amqplib.Delivery, h *handler) (err error) {
	defer func() {
		if r := recover(); r != nil {
			switch x := r.(type) {
			case string:
				err = errors.New(x)
			case error:
				err = x
			default:
				err = errors.New("Unknown panic")
			}
		}
	}()

	err = h.fn(messaging.Event{
		Id:     msg.MessageId,
		Action: h.action,
		Body:   msg.Body,
	})

	return
}

func (c *consumer) doDispatch(msg amqplib.Delivery, h *handler, retryCount int32, delay time.Duration) {
	err := c.callAndHandlePanic(msg, h)

	if err != nil {
		if h.maxRetries > 0 {
			if retryCount >= h.maxRetries {
				logger.WithFields(log.Fields{
					"max_retries": h.maxRetries,
					"message_id":  msg.MessageId,
				}).Error("Maximum retries reached. Giving up.")

				if !c.autoAck {
					msg.Ack(false)
				}
			} else {
				logger.WithFields(log.Fields{
					"error":      err,
					"message_id": msg.MessageId,
				}).Error("Failed to process event. Retrying...")

				c.retryMessage(msg, h, retryCount, delay)
			}
		} else {
			logger.WithFields(log.Fields{
				"error":      err,
				"message_id": msg.MessageId,
			}).Error("Failed to process event.")
		}
	} else if !c.autoAck {
		msg.Ack(false)
	}
}

func (c *consumer) publishMessage(msg amqplib.Publishing, queue string) error {
	channel, err := c.conn.OpenChannel()

	if err != nil {
		return err
	}

	defer channel.Close()

	if err := channel.Confirm(false); err != nil {
		return fmt.Errorf("Channel could not be put into confirm mode: %s", err)
	}

	confirms := channel.NotifyPublish(make(chan amqplib.Confirmation, 1))

	err = channel.Publish("", queue, false, false, msg)

	if err != nil {
		return err
	} else {
		if confirmed := <-confirms; !confirmed.Ack {
			return ErrNotAcked
		}
	}

	return nil
}

func (c *consumer) retryMessage(msg amqplib.Delivery, h *handler, retryCount int32, delay time.Duration) {
	if delay > 0 {
		if h.delayedRetry {
			delay *= 2
		}
	} else {
		delay = h.retryDelay
	}

	retryMsg := amqplib.Publishing{
		Headers: amqplib.Table{
			"x-retry-count": retryCount + 1,
			"x-retry-max":   h.maxRetries,
			"x-retry-delay": delay.String(),
			"x-action-key":  getAction(msg),
		},
		Timestamp:    time.Now(),
		DeliveryMode: msg.DeliveryMode,
		Body:         msg.Body,
		MessageId:    msg.MessageId,
	}

	err := c.publishMessage(retryMsg, c.queueName)

	if err != nil {
		logger.WithFields(log.Fields{
			"error": err,
		}).Error("Failed to retry.")

		if !c.autoAck {
			msg.Nack(false, true)
		}
	} else if !c.autoAck {
		msg.Ack(false)
	}
}

func (c *consumer) getHandler(msg amqplib.Delivery) (*handler, bool) {
	c.m.Lock()
	defer c.m.Unlock()

	action := getAction(msg)

	for _, h := range c.handlers {
		if h.re.MatchString(action) {
			return &h, true
		}
	}

	return nil, false
}

// Subscribe allows to subscribe an action handler.
func (c *consumer) Subscribe(action string, handlerFn messaging.EventHandler, options *messaging.SubscribeOptions) error {
	// TODO: Replace # pattern too.
	pattern := strings.Replace(action, "*", "(.*)", 0)
	re, err := regexp.Compile(pattern)

	if err != nil {
		return err
	}

	err = c.channel.QueueBind(
		c.queueName,    // queue name
		action,         // routing key
		c.exchangeName, // exchange
		false,          // no-wait
		nil,            // arguments
	)

	if err != nil {
		return err
	}

	if options == nil {
		options = &messaging.SubscribeOptions{
			RetryDelay:   time.Duration(0),
			DelayedRetry: false,
			MaxRetries:   0,
		}
	}

	c.handlers = append(c.handlers, handler{
		action:       action,
		fn:           handlerFn,
		re:           re,
		maxRetries:   options.MaxRetries,
		retryDelay:   options.RetryDelay,
		delayedRetry: options.DelayedRetry,
	})

	return nil
}

// Unsubscribe allows to unsubscribe an action handler.
func (c *consumer) Unsubscribe(action string) error {
	err := c.channel.QueueUnbind(
		c.queueName,    // queue name
		action,         // routing key
		c.exchangeName, // exchange
		nil,            // arguments
	)

	if err != nil {
		return err
	}

	idx := -1

	for i, h := range c.handlers {
		if h.action == action {
			idx = i
			break
		}
	}

	if idx != -1 {
		c.handlers = append(c.handlers[:idx], c.handlers[idx+1:]...)
	}

	return nil
}

// Listen start to listen for new messages.
func (c *consumer) Consume() {
	logger.Info("Registered handlers:")

	for _, handler := range c.handlers {
		logger.Infof("  %s", handler.action)
	}

	for !c.closed {
		logger.WithFields(log.Fields{
			"queue": c.queueName,
		}).Debug("Setting up consumer channel...")

		msgs, err := c.channel.Consume(
			c.queueName, // queue
			"",          // consumer
			c.autoAck,   // auto ack
			false,       // exclusive
			false,       // no local
			false,       // no wait
			nil,         // args
		)

		if err != nil {
			logger.WithFields(log.Fields{
				"queue": c.queueName,
				"error": err,
			}).Error("Error setting up consumer...")

			time.Sleep(c.config.ConsumeRetryInterval)

			c.setupTopology()

			continue
		}

		logger.WithFields(log.Fields{
			"queue": c.queueName,
		}).Info("Consuming messages...")

		for m := range msgs {
			logger.Info("Received from channel.")
			c.wg.Add(1)
			go func(msg amqplib.Delivery) {
				c.dispatch(msg)
				c.wg.Done()
			}(m)
		}

		logger.WithFields(log.Fields{
			"queue":  c.queueName,
			"closed": c.closed,
		}).Info("Consumption finished.")
	}
}

func getAction(msg amqplib.Delivery) string {
	if ac, ok := msg.Headers["x-action-key"]; ok {
		return ac.(string)
	} else {
		return msg.RoutingKey
	}
}

func getXRetryCountHeader(msg amqplib.Delivery) (int32, bool) {
	if c, ok := msg.Headers["x-retry-count"]; ok {
		return c.(int32), true
	}

	return 0, false
}

func getXRetryDelayHeader(msg amqplib.Delivery) (time.Duration, bool) {
	if d, ok := msg.Headers["x-retry-delay"]; ok {
		t, err := time.ParseDuration(d.(string))
		if err == nil {
			return t, true
		}
	}

	return time.Duration(0), false
}
<commit_msg>fix fmt<commit_after>package amqp

import (
	"errors"
	"fmt"
	"github.com/eventials/goevents/messaging"
	log "github.com/sirupsen/logrus"
	amqplib "github.com/streadway/amqp"
	"regexp"
	"strings"
	"sync"
	"time"
)

var (
	logger = log.WithFields(log.Fields{
		"type":     "goevents",
		"sub_type": "consumer",
	})
)

type handler struct {
	action       string
	fn           messaging.EventHandler
	re           *regexp.Regexp
	maxRetries   int32
	retryDelay   time.Duration
	delayedRetry bool
}

type consumer struct {
	config ConsumerConfig

	m  sync.Mutex
	wg sync.WaitGroup

	conn     *Connection
	autoAck  bool
	handlers []handler

	channel    *amqplib.Channel
	queue      *amqplib.Queue
	retryQueue *amqplib.Queue

	exchangeName string
	queueName    string

	closed bool
}

// ConsumerConfig to be used when creating a new producer.
type ConsumerConfig struct {
	ConsumeRetryInterval time.Duration
	PrefetchCount        int
	DurableQueue         bool
	AutoDelete           bool
	PrefixName           string
}

// NewConsumer returns a new AMQP Consumer.
// Uses a default ConsumerConfig with 2 second of consume retry interval.
func NewConsumer(c messaging.Connection, autoAck bool, exchange, queue string) (*consumer, error) {
	return NewConsumerConfig(c, autoAck, exchange, queue, ConsumerConfig{
		ConsumeRetryInterval: 2 * time.Second,
		PrefetchCount:        0,
		DurableQueue:         true,
		AutoDelete:           false,
	})
}

// NewConsumerConfig returns a new AMQP Consumer.
func NewConsumerConfig(c messaging.Connection, autoAck bool, exchange, queue string, config ConsumerConfig) (*consumer, error) {
	consumer := &consumer{
		config:       config,
		conn:         c.(*Connection),
		autoAck:      autoAck,
		handlers:     make([]handler, 0),
		exchangeName: exchange,
		queueName:    queue,
	}

	err := consumer.setupTopology()

	if err != nil {
		return nil, err
	}

	go consumer.handleReestablishedConnnection()

	return consumer, err

}

func (c *consumer) Close() {
	func() {
		c.m.Lock()
		defer c.m.Unlock()

		// Unsubscribe all handlers
		c.handlers = make([]handler, 0)

		c.closed = true
		c.channel.Close()
	}()

	// Wait all go routine finish.
	c.wg.Wait()
}

func (c *consumer) uniqueNameWithPrefix() string {
	return fmt.Sprintf("%s%d", c.config.PrefixName, time.Now().UnixNano())
}

func (c *consumer) setupTopology() error {
	c.m.Lock()
	defer c.m.Unlock()

	var err error

	if c.channel != nil {
		c.channel.Close()
	}

	c.channel, err = c.conn.OpenChannel()

	if err != nil {
		return err
	}

	err = c.channel.Qos(c.config.PrefetchCount, 0, true)

	if err != nil {
		return err
	}

	err = c.channel.ExchangeDeclare(
		c.exchangeName, // name
		"topic",        // type
		true,           // durable
		false,          // auto-delete
		false,          // internal
		false,          // no-wait
		nil,            // arguments
	)

	if err != nil {
		return err
	}

	if c.config.PrefixName != "" {
		c.queueName = c.uniqueNameWithPrefix()
	}

	q, err := c.channel.QueueDeclare(
		c.queueName,           // name
		c.config.DurableQueue, // durable
		c.config.AutoDelete,   // auto-delete
		false,                 // exclusive
		false,                 // no-wait
		nil,                   // arguments
	)

	if err != nil {
		return err
	}

	c.queue = &q

	return nil
}

func (c *consumer) handleReestablishedConnnection() {
	for !c.closed {
		<-c.conn.NotifyReestablish()

		err := c.setupTopology()

		if err != nil {
			logger.WithFields(log.Fields{
				"error": err,
			}).Error("Error setting up topology after reconnection.")
		}
	}
}

func (c *consumer) dispatch(msg amqplib.Delivery) {
	if h, ok := c.getHandler(msg); ok {
		delay, isRetry := getXRetryDelayHeader(msg)

		if isRetry {
			logger.WithFields(log.Fields{
				"delay":      delay.String(),
				"message_id": msg.MessageId,
			}).Info("Delaying message.")

			time.Sleep(delay)
		}

		retryCount, _ := getXRetryCountHeader(msg)

		c.doDispatch(msg, h, retryCount, delay)
	} else {
		if !c.autoAck {
			err := msg.Nack(false, true)

			if err != nil {
				logger.WithFields(log.Fields{
					"error":      err,
					"message_id": msg.MessageId,
				}).Error("Failed to nack message.")
			}
		}
	}
}

func (c *consumer) callAndHandlePanic(msg amqplib.Delivery, h *handler) (err error) {
	defer func() {
		if r := recover(); r != nil {
			switch x := r.(type) {
			case string:
				err = errors.New(x)
			case error:
				err = x
			default:
				err = errors.New("Unknown panic")
			}
		}
	}()

	err = h.fn(messaging.Event{
		Id:     msg.MessageId,
		Action: h.action,
		Body:   msg.Body,
	})

	return
}

func (c *consumer) doDispatch(msg amqplib.Delivery, h *handler, retryCount int32, delay time.Duration) {
	err := c.callAndHandlePanic(msg, h)

	if err != nil {
		if h.maxRetries > 0 {
			if retryCount >= h.maxRetries {
				logger.WithFields(log.Fields{
					"max_retries": h.maxRetries,
					"message_id":  msg.MessageId,
				}).Error("Maximum retries reached. Giving up.")

				if !c.autoAck {
					msg.Ack(false)
				}
			} else {
				logger.WithFields(log.Fields{
					"error":      err,
					"message_id": msg.MessageId,
				}).Error("Failed to process event. Retrying...")

				c.retryMessage(msg, h, retryCount, delay)
			}
		} else {
			logger.WithFields(log.Fields{
				"error":      err,
				"message_id": msg.MessageId,
			}).Error("Failed to process event.")
		}
	} else if !c.autoAck {
		msg.Ack(false)
	}
}

func (c *consumer) publishMessage(msg amqplib.Publishing, queue string) error {
	channel, err := c.conn.OpenChannel()

	if err != nil {
		return err
	}

	defer channel.Close()

	if err := channel.Confirm(false); err != nil {
		return fmt.Errorf("Channel could not be put into confirm mode: %s", err)
	}

	confirms := channel.NotifyPublish(make(chan amqplib.Confirmation, 1))

	err = channel.Publish("", queue, false, false, msg)

	if err != nil {
		return err
	} else {
		if confirmed := <-confirms; !confirmed.Ack {
			return ErrNotAcked
		}
	}

	return nil
}

func (c *consumer) retryMessage(msg amqplib.Delivery, h *handler, retryCount int32, delay time.Duration) {
	if delay > 0 {
		if h.delayedRetry {
			delay *= 2
		}
	} else {
		delay = h.retryDelay
	}

	retryMsg := amqplib.Publishing{
		Headers: amqplib.Table{
			"x-retry-count": retryCount + 1,
			"x-retry-max":   h.maxRetries,
			"x-retry-delay": delay.String(),
			"x-action-key":  getAction(msg),
		},
		Timestamp:    time.Now(),
		DeliveryMode: msg.DeliveryMode,
		Body:         msg.Body,
		MessageId:    msg.MessageId,
	}

	err := c.publishMessage(retryMsg, c.queueName)

	if err != nil {
		logger.WithFields(log.Fields{
			"error": err,
		}).Error("Failed to retry.")

		if !c.autoAck {
			msg.Nack(false, true)
		}
	} else if !c.autoAck {
		msg.Ack(false)
	}
}

func (c *consumer) getHandler(msg amqplib.Delivery) (*handler, bool) {
	c.m.Lock()
	defer c.m.Unlock()

	action := getAction(msg)

	for _, h := range c.handlers {
		if h.re.MatchString(action) {
			return &h, true
		}
	}

	return nil, false
}

// Subscribe allows to subscribe an action handler.
func (c *consumer) Subscribe(action string, handlerFn messaging.EventHandler, options *messaging.SubscribeOptions) error {
	// TODO: Replace # pattern too.
	pattern := strings.Replace(action, "*", "(.*)", 0)
	re, err := regexp.Compile(pattern)

	if err != nil {
		return err
	}

	err = c.channel.QueueBind(
		c.queueName,    // queue name
		action,         // routing key
		c.exchangeName, // exchange
		false,          // no-wait
		nil,            // arguments
	)

	if err != nil {
		return err
	}

	if options == nil {
		options = &messaging.SubscribeOptions{
			RetryDelay:   time.Duration(0),
			DelayedRetry: false,
			MaxRetries:   0,
		}
	}

	c.handlers = append(c.handlers, handler{
		action:       action,
		fn:           handlerFn,
		re:           re,
		maxRetries:   options.MaxRetries,
		retryDelay:   options.RetryDelay,
		delayedRetry: options.DelayedRetry,
	})

	return nil
}

// Unsubscribe allows to unsubscribe an action handler.
func (c *consumer) Unsubscribe(action string) error {
	err := c.channel.QueueUnbind(
		c.queueName,    // queue name
		action,         // routing key
		c.exchangeName, // exchange
		nil,            // arguments
	)

	if err != nil {
		return err
	}

	idx := -1

	for i, h := range c.handlers {
		if h.action == action {
			idx = i
			break
		}
	}

	if idx != -1 {
		c.handlers = append(c.handlers[:idx], c.handlers[idx+1:]...)
	}

	return nil
}

// Listen start to listen for new messages.
func (c *consumer) Consume() {
	logger.Info("Registered handlers:")

	for _, handler := range c.handlers {
		logger.Infof("  %s", handler.action)
	}

	for !c.closed {
		logger.WithFields(log.Fields{
			"queue": c.queueName,
		}).Debug("Setting up consumer channel...")

		msgs, err := c.channel.Consume(
			c.queueName, // queue
			"",          // consumer
			c.autoAck,   // auto ack
			false,       // exclusive
			false,       // no local
			false,       // no wait
			nil,         // args
		)

		if err != nil {
			logger.WithFields(log.Fields{
				"queue": c.queueName,
				"error": err,
			}).Error("Error setting up consumer...")

			time.Sleep(c.config.ConsumeRetryInterval)

			c.setupTopology()

			continue
		}

		logger.WithFields(log.Fields{
			"queue": c.queueName,
		}).Info("Consuming messages...")

		for m := range msgs {
			logger.Info("Received from channel.")
			c.wg.Add(1)
			go func(msg amqplib.Delivery) {
				c.dispatch(msg)
				c.wg.Done()
			}(m)
		}

		logger.WithFields(log.Fields{
			"queue":  c.queueName,
			"closed": c.closed,
		}).Info("Consumption finished.")
	}
}

func getAction(msg amqplib.Delivery) string {
	if ac, ok := msg.Headers["x-action-key"]; ok {
		return ac.(string)
	} else {
		return msg.RoutingKey
	}
}

func getXRetryCountHeader(msg amqplib.Delivery) (int32, bool) {
	if c, ok := msg.Headers["x-retry-count"]; ok {
		return c.(int32), true
	}

	return 0, false
}

func getXRetryDelayHeader(msg amqplib.Delivery) (time.Duration, bool) {
	if d, ok := msg.Headers["x-retry-delay"]; ok {
		t, err := time.ParseDuration(d.(string))
		if err == nil {
			return t, true
		}
	}

	return time.Duration(0), false
}
<|endoftext|>