xframe/component/queue/rabbitmq/sender.go

package rabbitmq

import (
	"context"
	"time"

	"git.wishpal.cn/wishpal_ironfan/xframe/base/proc"
	"git.wishpal.cn/wishpal_ironfan/xframe/base/rescue"
	"git.wishpal.cn/wishpal_ironfan/xframe/base/threading"
	"git.wishpal.cn/wishpal_ironfan/xframe/component/logger"
	"github.com/pkg/errors"
	amqp "github.com/rabbitmq/amqp091-go"
)

type RabbitSender struct {
	conf      RabbitSenderConf
	conn      *amqp.Connection
	connError chan *amqp.Error

	channelPool chan *amqp.Channel
	closedPool  bool //used for putChannel, avoid triggering panic when writing data to closed channels

	stopCtx       context.Context // control the lifecycle of get_channel and reconnect
	stopCtxCancel context.CancelFunc
}

func NewSender(conf RabbitSenderConf) (*RabbitSender, error) {
	ctx, cancel := context.WithCancel(context.Background())
	sender := &RabbitSender{
		conf:          conf,
		stopCtx:       ctx,
		stopCtxCancel: cancel,
	}

	// start sender
	if err := sender.start(); err != nil {
		return nil, err
	}

	// watch reconnect
	threading.GoSafeVoid(sender.watchReconnect)

	// watch channel pool size4
	threading.GoSafeVoid(func() {
		for {
			logger.Infof("rabbitmq sender channel pool(%d).", len(sender.channelPool))
			time.Sleep(1 * time.Second)
		}
	})

	// quit
	proc.AddShutdownListener(func() {
		cancel()
		sender.cleanup()
	})

	return sender, nil
}

func (q *RabbitSender) start() error {
	// connect to rabbitmq
	conn, err := amqp.Dial(q.conf.URL)
	if err != nil {
		return err
	}
	defer func() {
		if err != nil && conn != nil {
			conn.Close()
		}
	}()

	// declares an exchange on the server
	channel, err := conn.Channel()
	if err != nil {
		return err
	}

	err = channel.ExchangeDeclare(
		q.conf.Exchange.Name,
		q.conf.Exchange.Type,
		q.conf.Exchange.Durable,
		q.conf.Exchange.AutoDelete,
		q.conf.Exchange.Internal,
		q.conf.Exchange.NoWait,
		nil,
	)
	if err != nil {
		return err
	}

	// start multiple publisher channels, process message publishing
	pool := make(chan *amqp.Channel, q.conf.Concurrency)
	for i := 0; i < q.conf.Concurrency; i++ {
		channel, err = conn.Channel()
		if err != nil {
			return err
		}

		pool <- channel
	}

	q.conn = conn
	q.channelPool = pool
	q.connError = make(chan *amqp.Error, 1)
	q.conn.NotifyClose(q.connError)
	q.closedPool = false

	return nil
}

func (q *RabbitSender) reconnect() {
	var (
		retry    = 0
		maxRetry = 8
	)
	for {
		backoffDuration := time.Duration(1<<retry) * time.Second
		if retry > maxRetry {
			backoffDuration = time.Duration(1<<maxRetry) * time.Second
		}

		select {
		case <-q.stopCtx.Done():
			logger.Warnf("rabbitmq has stopped.")
			return
		case <-time.After(backoffDuration):
			if err := q.start(); err != nil {
				logger.Warnf("rabbitmq URL:%s, exchange:%s-%s, attempting to reconnect in %v-%d, err:%v",
					q.conf.URL, q.conf.Exchange.Name, q.conf.Exchange.Type, backoffDuration, retry, err)
				retry += 1
				continue
			}

			logger.Warnf("rabbitmq URL:%s, exchange:%s-%s, retry:%d, success reconnect",
				q.conf.URL, q.conf.Exchange.Name, q.conf.Exchange.Type, retry)
			return
		}
	}
}

func (q *RabbitSender) watchReconnect() {
	for {
		select {
		case <-q.stopCtx.Done():
			logger.Warnf("rabbitmq has stopped.")
			return
		case err := <-q.connError:
			logger.Errorf("rabbitmq URL:%s, exchange:%s-%s, conn err:%v",
				q.conf.URL, q.conf.Exchange.Name, q.conf.Exchange.Type, err)

			// resource cleanup
			q.cleanup()

			// reconnect
			q.reconnect()
		}
	}
}

func (q *RabbitSender) getChannel() (*amqp.Channel, error) {
	select {
	case <-q.stopCtx.Done():
		return nil, errors.New("rabbitmq has stopped.")
	case ch, ok := <-q.channelPool:
		if !ok {
			return nil, errors.New("channelPool closed. possibly attempting to reconnect.")
		}
		return ch, nil
	case <-time.After(time.Duration(10) * time.Millisecond):
		return nil, errors.New("timed out waiting for a channel")
	}
}

func (q *RabbitSender) putChannel(channel *amqp.Channel) {
	defer rescue.Recover()

	select {
	case <-q.stopCtx.Done():
		logger.Warnf("rabbitmq has stopped.")
	default:
		if q.closedPool {
			if !channel.IsClosed() {
				channel.Close()
			}
			return
		}
		q.channelPool <- channel
	}
}

func (q *RabbitSender) cleanup() {
	if q.conn != nil {
		q.conn.Close()
	}

	close(q.channelPool)
	q.closedPool = true
	for ch := range q.channelPool {
		if ch != nil {
			ch.Close()
		}
	}
}

func (q *RabbitSender) Send(ctx context.Context, msg Message) error {
	channel, err := q.getChannel()
	if err != nil {
		return err
	}
	defer q.putChannel(channel)

	err = channel.PublishWithContext(
		ctx,
		q.conf.Exchange.Name,
		msg.RouteKey,
		false,
		false,
		amqp.Publishing{
			ContentType:  q.conf.ContentType,
			Body:         msg.Body,
			DeliveryMode: amqp.Persistent,
		},
	)
	if err != nil {
		return err
	}

	return nil
}

func (q *RabbitSender) Stop() {
	q.stopCtxCancel()
	q.cleanup()
}