Pub/Sub is commonly used to run tasks asynchronously via queued messages. These tasks can often be long-running and involve multi-step, stateful I/O operations. As much as I'd wish that all the pubsub events have atomic and transactional logic, the world is not ideal and sometimes we must design for chaos, especially in eventually consistent systems. In this post, we’ll explore a simple pattern for gracefully shutting down Pub/Sub consumers in Go.
Gracefully shutting down Pub/Sub consumers can be achieved with a simple approach:
A consumer can shut down for various reasons—Kubernetes upgrades, autoscaling restarts, manual throttling, redeployments, or even just Murphy's law. Abrupt shutdowns can be dangerous. I’ve encountered incidents where a consumer took a distributed lock before restarting during a redeployment, and the lock was never released. When consumers aren’t fully idempotent or transactional, these sudden shutdowns can leave asynchronous state machines stuck in an inconsistent state, sometimes for an extended period or even permanently.
To avoid this, let’s explore a recipe to make the system resilient to terminations and interruptions. While it's hard to handle cases like OOM failures where no termination signal is sent, we can definitely make consumers resilient to most shutdown scenarios.
A good way to implement graceful shutdown for Pub/Sub consumers is to think of it like a restaurant. When a restaurant closes (say at 12PM), it stops accepting new customers but continues serving those already inside until they finish. Eventually, the restaurant shuts down completely. We can apply a similar strategy to Pub/Sub consumers:
1. Listen for Termination Signals
In Go, you can listen for termination signals using channels. Channels are perfect for sending and receiving signals between goroutines, and since Google’s official Pub/Sub SDK processes messages in goroutines, they fit our use case well.
type terminationChan chan bool
// Create a channel to handle termination signals for routines
func CreateTerminationChannel() terminationChan {
// Create a buffered channel to hold the termination signal
terminationChannel := make(chan bool, 1)
// Listen for SIGTERM, SIGINT, and OS interrupt signals
sigChan := make(chan os.Signal, 1)
signal.Notify(sigChan, os.Interrupt, syscall.SIGTERM, syscall.SIGINT)
go func() {
<-sigChan
// Send termination signal through the channel
terminationChannel <- true
}()
// Return the termination channel
return terminationChannel
}
Here, we’ve created a boolean termination channel to capture termination signals. Now that we have this channel, we can pass it to our Pub/Sub client to handle shutdowns gracefully.
2. Reject New Incoming Messages
Once a termination signal is received, it’s important to stop processing any new messages. This can be done in multiple ways, but I chose to simply wrap my subscriber handler in a decorator.
func withGracefulShutdown(handler func(context.Context, *pubsub.Message), terminationChannel chan bool) func(context.Context, *pubsub.Message) {
return func(ctx context.Context, message *pubsub.Message) {
select {
case <-terminationChannel:
message.Nack() // Reject the message
default:
handler(ctx, message) // Process the message if no termination signal
}
}
}
This function takes the message handler and the termination channel, then returns a decorated handler. If a termination signal is received, new messages are immediately rejected using Nack(). This ensures no new messages are processed after receiving the shutdown signal.
3. Finish Ongoing Message Processing and Exit
When message processing starts just before a termination signal is received, it's important to allow time for those messages to complete processing to avoid leaving them in a half-processed state. We can use Go’s waitgroups to handle this. By incrementing a global waitgroup whenever a message starts processing and decrementing it once processing finishes, we can track how many messages are still being processed at any given time.
Upon receiving a termination signal, we wait for the waitgroup to reach zero before shutting down the consumer. The code would look something like this:
func startConsumer(ctx context.Context, subscription *pubsub.Subscription, terminationChannel chan bool) {
cancellableCtx, cancel := context.WithCancel(ctx)
var handlerWaitGroup sync.WaitGroup
go func() {
<-terminationChannel
logger.Info("Received termination signal")
logger.Info("Waiting for ongoing message processing")
handlerWaitGroup.Wait()
logger.Info("All ongoing messages processed; shutting down subscriber")
cancel()
logger.Info("Subscriber terminated")
}()
subscription.Receive(cancellableCtx, withGracefulShutdown(func(messageCtx context.Context, message *pubsub.Message) {
// increment the waitgroup when the processing starts and decrement it when the handler returns
handlerWaitGroup.Add(1)
defer handlerWaitGroup.Done()
// Message processing logic
}, terminationChannel))
}
This approach ensures that all messages are fully processed before the consumer shuts down. The Receive method will return once the context is cancelled, thus exiting the consumer elegantly.
This approach is effective for most use cases where messages are short-lived, typically running for 3-5 minutes or less—depending on how much delay you can tolerate for deployment or autoscaling. However, based on your specific requirements, there are additional steps you can take to further improve your graceful shutdown logic.
Rishichandra Wawhal, Engineering, Zamp