Skip to content

Instantly share code, notes, and snippets.

Embed
What would you like to do?
A robust rabbitmq client for Go
const (
// When reconnecting to the server after connection failure
reconnectDelay = 5 * time.Second
)
// Client holds necessery information for rabbitMQ
type Client struct {
pushQueue string
logger zerolog.Logger
connection *amqp.Connection
channel *amqp.Channel
done chan os.Signal
notifyClose chan *amqp.Error
notifyConfirm chan amqp.Confirmation
isConnected bool
alive bool
threads int
wg *sync.WaitGroup
}
// New is a constructor that takes address, push and listen queue names, logger, and a channel that will notify rabbitmq client on server shutdown. We calculate the number of threads, create the client, and start the connection process. Connect method connects to the rabbitmq server and creates push/listen channels if they don't exist.
func New(listenQueue, pushQueue, addr string, l zerolog.Logger, done chan os.Signal) *Client {
threads := runtime.GOMAXPROCS(0)
if numCPU := runtime.NumCPU(); numCPU > threads {
threads = numCPU
}
client := Client{
logger: l,
threads: threads,
pushQueue: pushQueue,
done: done,
alive: true,
wg: &sync.WaitGroup{},
}
client.wg.Add(threads)
go client.handleReconnect(addr)
return &client
}
// handleReconnect will wait for a connection error on
// notifyClose, and then continuously attempt to reconnect.
func (c *Client) handleReconnect(listenQueue, addr string) {
for c.alive {
c.isConnected = false
fmt.Printf("Attempting to connect to rabbitMQ: %s\n", addr)
var retryCount int
for !c.connect(listenQueue, addr) {
c.logger.Printf("disconnected from rabbitMQ and failed to connect")
time.Sleep(reconnectDelay + time.Duration(retryCount)*time.Second)
retryCount++
}
select {
case <-c.done:
return
case <-c.notifyClose:
}
}
}
// connect will make a single attempt to connect to
// RabbitMq. It returns the success of the attempt.
func (c *Client) connect(listenQueue, addr string) bool {
conn, err := amqp.Dial(addr)
if err != nil {
c.logger.Printf("failed to dial rabbitMQ server: %v", err)
return false
}
ch, err := conn.Channel()
if err != nil {
c.logger.Printf("failed connecting to channel: %v", err)
return false
}
ch.Confirm(false)
_, err = ch.QueueDeclare(
listenQueue,
true, // Durable
false, // Delete when unused
false, // Exclusive
false, // No-wait
nil, // Arguments
)
if err != nil {
c.logger.Printf("failed to declare listen queue: %v", err)
return false
}
_, err = ch.QueueDeclare(
c.pushQueue,
true, // Durable
false, // Delete when unused
false, // Exclusive
false, // No-wait
nil, // Arguments
)
if err != nil {
c.logger.Printf("failed to declare push queue: %v", err)
return false
}
c.changeConnection(conn, ch)
c.isConnected = true
return true
}
// changeConnection takes a new connection to the queue,
// and updates the channel listeners to reflect this.
func (c *Client) changeConnection(connection *amqp.Connection, channel *amqp.Channel) {
c.connection = connection
c.channel = channel
c.notifyClose = make(chan *amqp.Error)
c.notifyConfirm = make(chan amqp.Confirmation)
c.channel.NotifyClose(c.notifyClose)
c.channel.NotifyPublish(c.notifyConfirm)
}
// Push will push data onto the queue, and wait for a confirmation.
// If no confirms are received until within the resendTimeout,
// it continuously resends messages until a confirmation is received.
// This will block until the server sends a confirm.
func (c *Client) Push(data []byte) error {
if !c.isConnected {
return errors.New("failed to push push: not connected")
}
for {
err := c.UnsafePush(data)
if err != nil {
continue
}
select {
case confirm := <-c.notifyConfirm:
if confirm.Ack {
return nil
}
case <-time.After(resendDelay):
}
}
}
func (c *Client) UnsafePush(data []byte) error {
if !c.isConnected {
return errNotConnected
}
return c.channel.Publish(
"", // Exchange
c.name, // Routing key
false, // Mandatory
false, // Immediate
amqp.Publishing{
ContentType: "text/plain",
Body: data,
},
)
}
func (c *Client) Stream(cancelCtx context.Context) error {
for {
if c.isConnected {
break
}
time.Sleep(1 * time.Second)
}
err := c.channel.Qos(1, 0, false)
if err != nil {
return err
}
for i := 1; i <= c.threads; i++ {
msgs, err := c.channel.Consume(
c.streamQueue,
consumerName(i), // Consumer
false, // Auto-Ack
false, // Exclusive
false, // No-local
false, // No-Wait
nil, // Args
)
if err != nil {
return err
}
go func() {
defer c.wg.Done()
for {
select {
case <-cancelCtx.Done():
return
case msg := <-msgs:
c.parseEvent(msg)
}
}
}()
}
return nil
}
type event struct{
Job string `json:"job"`
Data string `json:"data"`
}
func (c *Client) parseEvent(msg amqp.Delivery) {
l := c.logger.Log().Timestamp()
startTime := time.Now()
var evt event
err := json.Unmarshal(msg.Body, &evt)
if err != nil {
logAndNack(msg, l, startTime, "unmarshalling body: %s - %s", string(msg.Body), err.Error())
return
}
if evt.Data == "" {
logAndNack(msg, l, startTime, "received event without data")
return
}
defer func(ctx context.Context, e event, m amqp.Delivery, logger *zerolog.Event) {
if err := recover(); err != nil {
stack := make([]byte, 8096)
stack = stack[:runtime.Stack(stack, false)]
l.Bytes("stack", stack).Str("level", "fatal").Interface("error", err).Msg("panic recovery for rabbitMQ message")
msg.Nack(false, false)
}
}(ctx, evt, msg, l)
switch evt.Job {
case "job1":
// Call an actual function
err = func()
case "job1":
err = func()
default:
msg.Reject(false)
return
}
if err != nil {
logAndNack(msg, l, startTime, err.Error())
return
}
l.Str("level", "info").Int64("took-ms", time.Since(startTime).Milliseconds()).Msgf("%s succeeded" evt.Job)
msg.Ack(false)
}
func logAndNack(msg amqp.Delivery, l *zerolog.Event, t time.Time, err string, args ...interface{}) {
msg.Nack(false, false)
l.Int64("took-ms", time.Since(t).Milliseconds()).Str("level", "error").Msg(fmt.Sprintf(err, args...))
}
func (c *Client) Close() error {
if !c.isConnected {
return errAlreadyClosed
}
c.alive = false
fmt.Println("Waiting for current messages to be processed...")
c.wg.Wait()
for i := 1; i <= c.threads; i++ {
fmt.Println("Closing consumer: ", i)
err := c.channel.Cancel(consumerName(i), false)
if err != nil {
return fmt.Errorf("error canceling consumer %s: %v", consumerName(i), err)
}
}
err := c.channel.Close()
if err != nil {
return err
}
err = c.connection.Close()
if err != nil {
return err
}
c.isConnected = false
fmt.Println("gracefully stopped rabbitMQ connection")
return nil
}
func consumerName(i int) string {
return fmt.Sprintf("go-consumer-%v", i)
}
@kamal-github

This comment has been minimized.

Copy link

@kamal-github kamal-github commented Sep 26, 2020

I believe, For Push() L122 you are returning on connection error and if the connection failure happens right after L122, then we are waiting until the message gets published and confirmed. so Sometime it will throw an error, and other times it will reconnect and publish successfully eventually.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
You can’t perform that action at this time.