124 lines
3.2 KiB
Go
124 lines
3.2 KiB
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
|
|
}
|