github.com/wfusion/gofusion@v1.1.14/common/infra/watermill/pubsub/amqp/connection.go (about)

     1  package amqp
     2  
     3  import (
     4  	"sync"
     5  
     6  	"github.com/cenkalti/backoff/v4"
     7  	"github.com/pkg/errors"
     8  	"go.uber.org/atomic"
     9  
    10  	"github.com/wfusion/gofusion/common/infra/watermill"
    11  
    12  	amqp "github.com/rabbitmq/amqp091-go"
    13  )
    14  
    15  // ConnectionWrapper manages an AMQP connection.
    16  type ConnectionWrapper struct {
    17  	config ConnectionConfig
    18  
    19  	logger watermill.LoggerAdapter
    20  
    21  	amqpConnection     *amqp.Connection
    22  	amqpConnectionLock sync.Mutex
    23  	connected          chan struct{}
    24  
    25  	closing chan struct{}
    26  	closed  atomic.Uint32
    27  
    28  	connectionWaitGroup sync.WaitGroup
    29  }
    30  
    31  // NewConnection returns a new connection wrapper.
    32  func NewConnection(
    33  	config ConnectionConfig,
    34  	logger watermill.LoggerAdapter,
    35  ) (*ConnectionWrapper, error) {
    36  	if logger == nil {
    37  		logger = watermill.NopLogger{}
    38  	}
    39  
    40  	pubSub := &ConnectionWrapper{
    41  		config:    config,
    42  		logger:    logger,
    43  		closing:   make(chan struct{}),
    44  		connected: make(chan struct{}),
    45  	}
    46  	if err := pubSub.connect(); err != nil {
    47  		return nil, err
    48  	}
    49  
    50  	go pubSub.handleConnectionClose()
    51  
    52  	return pubSub, nil
    53  }
    54  
    55  func (c *ConnectionWrapper) Close() error {
    56  	if !c.closed.CompareAndSwap(0, 1) {
    57  		// Already closed.
    58  		return nil
    59  	}
    60  
    61  	close(c.closing)
    62  
    63  	c.logger.Info("[Common] watermill closing AMQP pub/sub", nil)
    64  	defer c.logger.Info("[Common] watermill amqp closed AMQP pub/sub", nil)
    65  
    66  	c.connectionWaitGroup.Wait()
    67  
    68  	if err := c.amqpConnection.Close(); err != nil {
    69  		c.logger.Error("[Common] watermill amqp connection close error", err, nil)
    70  	}
    71  
    72  	return nil
    73  }
    74  
    75  func (c *ConnectionWrapper) connect() error {
    76  	c.amqpConnectionLock.Lock()
    77  	defer c.amqpConnectionLock.Unlock()
    78  
    79  	amqpConfig := c.config.AmqpConfig
    80  	if amqpConfig != nil && amqpConfig.TLSClientConfig != nil && c.config.TLSConfig != nil {
    81  		return errors.New("both Config.AmqpConfig.TLSClientConfig and Config.TLSConfig are set")
    82  	}
    83  
    84  	var connection *amqp.Connection
    85  	var err error
    86  
    87  	if amqpConfig != nil {
    88  		connection, err = amqp.DialConfig(c.config.AmqpURI, *c.config.AmqpConfig)
    89  	} else if c.config.TLSConfig != nil {
    90  		connection, err = amqp.DialTLS(c.config.AmqpURI, c.config.TLSConfig)
    91  	} else {
    92  		connection, err = amqp.Dial(c.config.AmqpURI)
    93  	}
    94  
    95  	if err != nil {
    96  		return errors.Wrap(err, "cannot connect to AMQP")
    97  	}
    98  	c.amqpConnection = connection
    99  	close(c.connected)
   100  
   101  	c.logger.Info("[Common] watermill connected to AMQP", nil)
   102  
   103  	return nil
   104  }
   105  
   106  func (c *ConnectionWrapper) Connection() *amqp.Connection {
   107  	return c.amqpConnection
   108  }
   109  
   110  func (c *ConnectionWrapper) Connected() chan struct{} {
   111  	return c.connected
   112  }
   113  
   114  func (c *ConnectionWrapper) IsConnected() bool {
   115  	select {
   116  	case <-c.connected:
   117  		return true
   118  	default:
   119  		return false
   120  	}
   121  }
   122  
   123  func (c *ConnectionWrapper) Closed() bool {
   124  	return c.closed.Load() == 1
   125  }
   126  
   127  func (c *ConnectionWrapper) handleConnectionClose() {
   128  	for {
   129  		c.logger.Debug("[Common] watermill amqp handle connection close is waiting for c.connected", nil)
   130  		<-c.connected
   131  		c.logger.Debug("[Common] watermill amqp handle connection close is for connection or pub/sub close", nil)
   132  
   133  		notifyCloseConnection := c.amqpConnection.NotifyClose(make(chan *amqp.Error))
   134  
   135  		select {
   136  		case <-c.closing:
   137  			c.logger.Debug("[Common] watermill amqp stopping handle connection close", nil)
   138  			c.connected = make(chan struct{})
   139  			return
   140  		case err := <-notifyCloseConnection:
   141  			c.connected = make(chan struct{})
   142  			c.logger.Error("[Common] watermill received close notification from AMQP, reconnecting", err, nil)
   143  			c.reconnect()
   144  		}
   145  	}
   146  }
   147  
   148  func (c *ConnectionWrapper) reconnect() {
   149  	reconnectConfig := c.config.Reconnect
   150  	if reconnectConfig == nil {
   151  		reconnectConfig = DefaultReconnectConfig()
   152  	}
   153  
   154  	if err := backoff.Retry(func() error {
   155  		err := c.connect()
   156  		if err == nil {
   157  			return nil
   158  		}
   159  
   160  		c.logger.Error("[Common] watermill cannot reconnect to AMQP, retrying", err, nil)
   161  
   162  		if c.Closed() {
   163  			return backoff.Permanent(errors.Wrap(err, "closing AMQP connection"))
   164  		}
   165  
   166  		return err
   167  	}, reconnectConfig.backoffConfig()); err != nil {
   168  		// should only exit, if closing Pub/Sub
   169  		c.logger.Error("[Common] watermill amqp reconnect failed", err, nil)
   170  	}
   171  }