rabbitmq-amqp-go-client/pkg/rabbitmqamqp/amqp_environment.go

package rabbitmqamqp

import (
	"context"
	"fmt"
	"github.com/Azure/go-amqp"
	"sync"
	"sync/atomic"
)

type TEndPointStrategy int

const (
	StrategyRandom     TEndPointStrategy = iota
	StrategySequential TEndPointStrategy = iota
)

type Endpoint struct {
	Address string
	Options *AmqpConnOptions
}

func DefaultEndpoints() []Endpoint {
	ep := Endpoint{
		Address: "amqp://",
		Options: &AmqpConnOptions{
			SASLType: amqp.SASLTypeAnonymous(),
		},
	}

	return []Endpoint{ep}
}

type Environment struct {
	connections      sync.Map
	endPoints        []Endpoint
	EndPointStrategy TEndPointStrategy
	nextConnectionId int32
}

func NewEnvironment(address string, options *AmqpConnOptions) *Environment {
	return NewClusterEnvironmentWithStrategy([]Endpoint{{Address: address, Options: options}}, StrategyRandom)
}

func NewClusterEnvironment(endPoints []Endpoint) *Environment {
	return NewClusterEnvironmentWithStrategy(endPoints, StrategyRandom)
}

func NewClusterEnvironmentWithStrategy(endPoints []Endpoint, strategy TEndPointStrategy) *Environment {
	return &Environment{
		connections:      sync.Map{},
		endPoints:        endPoints,
		EndPointStrategy: strategy,
		nextConnectionId: 0,
	}
}

// NewConnection get a new connection from the environment.
// It picks an endpoint from the list of endpoints, based on EndPointStrategy, and tries to open a connection.
// It fails if all the endpoints are not reachable.
// The Environment will keep track of the connection and close it when the environment is closed.
func (e *Environment) NewConnection(ctx context.Context) (*AmqpConnection, error) {

	tmp := make([]Endpoint, len(e.endPoints))
	copy(tmp, e.endPoints)
	lastError := error(nil)
	for len(tmp) > 0 {
		idx := 0

		switch e.EndPointStrategy {
		case StrategyRandom:
			idx = random(len(tmp))
		case StrategySequential:
			idx = 0
		}

		addr := tmp[idx]
		// remove the index from the tmp list
		tmp = append(tmp[:idx], tmp[idx+1:]...)
		var cloned *AmqpConnOptions
		if addr.Options != nil {
			cloned = addr.Options.Clone()
		}
		connection, err := Dial(ctx, addr.Address, cloned)
		if err != nil {
			Error("Failed to open connection", ExtractWithoutPassword(addr.Address), err)
			lastError = err
			continue
		}

		// here we use it to make each connection unique
		atomic.AddInt32(&e.nextConnectionId, 1)
		connection.amqpConnOptions.Id = fmt.Sprintf("%s_%d", connection.amqpConnOptions.Id, e.nextConnectionId)
		e.connections.Store(connection.Id(), connection)

		connection.refMap = &e.connections
		return connection, nil
	}
	return nil, fmt.Errorf("fail to open connection. Last error: %w", lastError)
}

// Connections gets the active connections in the environment
func (e *Environment) Connections() []*AmqpConnection {
	connections := make([]*AmqpConnection, 0)
	e.connections.Range(func(key, value interface{}) bool {
		connections = append(connections, value.(*AmqpConnection))
		return true
	})
	return connections
}

// CloseConnections closes all the connections in the environment with all the publishers and consumers.
func (e *Environment) CloseConnections(ctx context.Context) error {
	var err error
	e.connections.Range(func(key, value any) bool {
		connection := value.(*AmqpConnection)
		if cerr := connection.Close(ctx); cerr != nil {
			err = cerr
		}
		return true
	})
	return err
}