package amqp import ( "errors" "sync" "time" "github.com/rs/zerolog/log" "github.com/streadway/amqp" ) type AMQPSession struct { connection *amqp.Connection channel *amqp.Channel done chan bool streamsReInit []chan bool ready chan bool notifyConnClose chan *amqp.Error notifyChanClose chan *amqp.Error notifyConfirm chan amqp.Confirmation isready bool mutex sync.Mutex mutexCh sync.Mutex } const ( // When reconnecting to the server after connection failure reconnectDelay = 5 * time.Second // When setting up the channel after a channel exception reInitDelay = 2 * time.Second // When resending messages the server didn't confirm resendDelay = 5 * time.Second ) var ( errNotConnected = errors.New("not connected to a server") errAlreadyClosed = errors.New("already closed: not connected to the server") errShutdown = errors.New("session is shutting down") ) // NewAMQPSession creates a new consumer state instance, and automatically // attempts to connect to the server. func NewAMQPSession(addr string) (*AMQPSession, error) { session := AMQPSession{ ready: make(chan bool, 1), streamsReInit: make([]chan bool, 0), } go session.handleReconnect(addr) return &session, nil } func (session *AMQPSession) waitChannelReady() { session.mutex.Lock() if !session.isready { session.mutex.Unlock() select { case <-session.ready: return } } else { session.mutex.Unlock() } } func (session *AMQPSession) setReady(ready bool) { session.mutex.Lock() session.isready = ready session.mutex.Unlock() if ready { session.ready <- true } } // connect will create a new AMQP connection func (session *AMQPSession) connect(addr string) (*amqp.Connection, error) { conn, err := amqp.Dial(addr) if err != nil { return nil, err } session.changeConnection(conn) log.Info().Msg("AMQP Connected to AMQP!") return conn, nil } // handleReconnect will wait for a connection error on // notifyConnClose, and then continuously attempt to reconnect. func (session *AMQPSession) handleReconnect(addr string) { for { session.setReady(false) log.Info().Msg("AMQP Attempting to connect") conn, err := session.connect(addr) if err != nil { log.Info().Msg("AMQP Failed to connect. Retrying...") select { case <-session.done: return case <-time.After(reconnectDelay): } continue } if done := session.handleReInit(conn); done { break } } log.Fatal().Msg("AMQP stopped") } // handleReconnect will wait for a channel error // and then continuously attempt to re-initialize both channels func (session *AMQPSession) handleReInit(conn *amqp.Connection) bool { for { session.setReady(false) err := session.init(conn) if err != nil { log.Info().Msg("AMQP Failed to initialize channel. Retrying...") select { case <-session.done: log.Info().Msg("AMQP Connection closed. Done") return true case <-time.After(reInitDelay): } continue } select { case <-session.done: log.Info().Msg("AMQP Connection closed. Done") return true case <-session.notifyConnClose: log.Info().Msg("AMQP Connection closed. Reconnecting...") for _, ch := range session.streamsReInit { ch <- true } return false case <-session.notifyChanClose: log.Info().Msg("AMQP Channel closed. Re-running init...") for _, ch := range session.streamsReInit { ch <- true } } } } // init will initialize channel func (session *AMQPSession) init(conn *amqp.Connection) error { ch, err := conn.Channel() if err != nil { return err } err = ch.Confirm(false) if err != nil { return err } session.changeChannel(ch) session.setReady(true) log.Info().Msg("AMQP Setup complete!") return nil } // changeConnection takes a new connection to the queue, // and updates the close listener to reflect this. func (session *AMQPSession) changeConnection(connection *amqp.Connection) { session.connection = connection session.notifyConnClose = make(chan *amqp.Error) session.connection.NotifyClose(session.notifyConnClose) } // changeChannel takes a new channel to the queue, // and updates the channel listeners to reflect this. func (session *AMQPSession) changeChannel(channel *amqp.Channel) { session.mutexCh.Lock() session.channel = channel session.notifyChanClose = make(chan *amqp.Error) session.notifyConfirm = make(chan amqp.Confirmation, 1) session.channel.NotifyClose(session.notifyChanClose) session.channel.NotifyPublish(session.notifyConfirm) session.mutexCh.Unlock() } // Push will push data onto the queue, and wait for a confirm. // If no confirms are received until within the resendTimeout, // it continuously re-sends messages until a confirm is received. // This will block until the server sends a confirm. Errors are // only returned if the push action itself fails, see UnsafePush. func (session *AMQPSession) Push(ex, rt string, data []byte) error { for { err := session.UnsafePush(ex, rt, data) if err != nil { log.Info().Msg("AMQP Push failed. Retrying...") select { case <-session.done: return errShutdown case <-time.After(resendDelay): } continue } select { case confirm := <-session.notifyConfirm: if confirm.Ack { log.Info().Msg("AMQP Push confirmed!") return nil } case <-time.After(resendDelay): } log.Info().Msg("AMQP Push didn't confirm. Retrying...") } } // Push will push to the queue without checking for // confirmation. It returns an error if it fails to connect. // No guarantees are provided for whether the server will // recieve the message. func (session *AMQPSession) UnsafePush(ex, rt string, data []byte) error { session.waitChannelReady() session.mutexCh.Lock() defer session.mutexCh.Unlock() return session.channel.Publish( ex, // Exchange rt, // Routing key false, // Mandatory false, // Immediate amqp.Publishing{ ContentType: "text/plain", Body: data, }, ) } // Stream will continuously put queue items on the channel. // It is required to call delivery.Ack when it has been // successfully processed, or delivery.Nack when it fails. // Ignoring this will cause data to build up on the server. func (session *AMQPSession) Stream(exName, qName, rKey string, consumer func(amqp.Delivery)) error { session.mutex.Lock() reinit := make(chan bool, 1) session.streamsReInit = append(session.streamsReInit, reinit) session.mutex.Unlock() go func() { for { session.waitChannelReady() _, err := session.channel.QueueDeclare( qName, false, // Durable true, // Delete when unused true, // Exclusive false, // No-wait nil, // Arguments ) if err != nil { panic("AMQP failed to declare queue " + qName + ": " + err.Error()) } session.channel.ExchangeDeclare(exName, "topic", false, false, false, false, nil) session.channel.QueueBind(qName, rKey, exName, false, nil) ch, err := session.channel.Consume( qName, "", // Consumer true, // Auto-Ack true, // Exclusive false, // No-local false, // No-Wait nil, // Args ) if err != nil { panic("AMQP failed to consume queue " + err.Error()) } func() { for { select { case delivery := <-ch: consumer(delivery) case <-reinit: return } } }() log.Warn().Msg("AMQP Stream loop interrupted") } }() return nil } // Close will delete the queue, close the channel and the connection. func (session *AMQPSession) Close(qName string) error { log.Error().Msg("Closing connection to RabbitMQ") _, err := session.channel.QueueDelete(qName, false, false, false) if err != nil { return err } err = session.channel.Close() if err != nil { return err } err = session.connection.Close() if err != nil { return err } return nil }