Have you ever been in a situation when you wanted your application to always finish ongoing processing like downloading a file, sending a message to a broker, crediting amount to a bank account, sending response to an upstream service, etc. before application exits either naturally or prematurely ?

If you are like me, you probably have and if not, it will eventually happen. Because graceful shutdown mechanism is significant for applications which perform frequent background processing tasks to ensure the integrity of results in case of abrupt termination. It ensures that any ongoing operation/processing is run to its completion when application is forced to exit. Every golang developer who deals with goroutines should be aware of it.

By the end of this post, you will be able to write go applications that can handle shutdown signals gracefully, perform cleanup tasks efficiently and exit smoothly. Sounds good?

Let’s get started!

TL;DR

We need to add following two functionalities in a Go application to ensure graceful shutdown during abrupt termination -

  1. The main function must wait for all background goroutines to finish. Use done channel idiom or sync.WaitGroup struct.
  2. The application must listen for incoming os signal related to the termination e.g. SIGTERM, SIGHUP or SIGINT. See documentation.

We can trap the termination signal and close the application gracefully with following pattern -

func main() {
	ctx, cancel := signal.NotifyContext(
		context.Background(),
		syscall.SIGTERM,
		syscall.SIGHUP,
		syscall.SIGINT,
	)
	done := make(chan struct{})
	go func() {
		defer close(done)
		for { // for-select-done idiom
			select {
			case <-ctx.Done():
				return
				// ... other cases
			}
		}
	}()
	// do something else here

	<-done // wait for goroutine to finish
	cancel()
}

Simple Notification App

In this application, we will receive a message from a source and notifies about it to a destination.

Usually in real world, message can come from various sources like upstream service, message broker, user request, event, etc. and destination can be frontend, downstream service, message broker, user response, database, etc.

For our application we will simply simulate the source and destination by defining two functions -

  • MessageChannel: It returns receive-only channel which acts as a source for messages.
// MessageChannel returns a receive only channel and starts a goroutine
// which sends message to that channel after every 2 seconds.
func MessageChannel() <-chan string {
	ch := make(chan string)
	go func() {
		defer close(ch)
		for i := 0; i < 3; i++ {
			time.Sleep(2 * time.Second)
			ch <- fmt.Sprintf("message %d", i)
		}
	}()
	return ch
}

Note that for loop has three iterations so we are only generating three messages.

  • NotificationChannel: It returns send-only channel acts as a destination for notifications.
// NotificationChannel returns a send only channel and starts
// a goroutine which receives notifications from that channel.
func NotificationChannel() chan<- string {
	ch := make(chan string)
	go func() {
		for _ = range ch {
			// do nothing
		}
	}()
	return ch
}

Now, lets write a main function that starts a worker goroutine to receive a message, perform some processing and then send a notification -

func main() {
	source := MessageChannel()
	destination := NotificationChannel()

	done := make(chan struct{})
  // worker goroutine
	go func() {
		defer close(done)
		for {
			message, ok := <-source
			if !ok {
				return
			}
			fmt.Printf("message received: %s\n", message)
			// some processing
			fmt.Println("processing...")
			time.Sleep(time.Second)
			notification := fmt.Sprintf("notification - %s", message)
			destination <- notification
			fmt.Println("notification sent")
		}
	}()
	<-done
}

Our main function waits for the worker goroutine to finish using done channel idiom. The processing of message is simulated with time.Sleep for a second.

It is fine, you can also use sync.WaitGroup instead of done channel.

Let’s run our application and wait for it to complete, we see the following expected output -

message received: message 0
processing...
notification sent
message received: message 1
processing...
notification sent
message received: message 2
processing...
notification sent

The Problem

Go applications may be abruptly terminated by triggers such as container stop/restart when using docker, pressing Ctrl+C or closing the terminal.

Let’s press Ctrl+C while our application is processing the received message. It abruptly terminates the execution and the message is lost -

message received: message 0
processing...
notification sent
message received: message 1
processing...
exit status 0xc000013a

The above output shows that notification for message 0 was sent to destination but notification for message 1 is lost due to premature termination during processing.

How can we ensure that the notification for any received message is not lost and always gets delivered to the destination?

Let’s write solution to this problem.

The Solution (Graceful Shutdown)

To implement graceful shutdown in our application, we can use signal.NotifyContext function and for-select-done idiom. The NotifyContext creates a child context whose Done channel is closed when os sends specified signal. In our case, we are interested in interrupt signal os.Interrupt or syscall.SIGINT (there are also other signals SIGHUP and SIGTERM. See doc). The updated main function looks like below -

func main() {
	source := MessageChannel()
	destination := NotificationChannel()

	ctx, cancel := signal.NotifyContext(context.Background(), os.Interrupt)
	done := make(chan struct{})
	go func() {
		defer close(done)
		for {
			select {
			case <-ctx.Done():
				return
			case message, ok := <-source:
				if !ok {
					return
				}
				fmt.Printf("message received: %s\n", message)
				// some processing
				fmt.Println("processing...")
				time.Sleep(time.Second)
				notification := fmt.Sprintf("notification - %s", message)
				destination <- notification
				fmt.Println("notification sent")
			}
		}
	}()
	<-done
	cancel() // release context resources
	fmt.Printf("main exits")
}

The worker goroutine waits for either context’s Done channel to get closed or the message from the source. The case block whose operation succeeds first is executed. The last Printf statement in main function acts as a marker to track whether execution reaches end of main function when our application is terminated forcefully.

Let’s press Ctrl+C again while application is processing a message. You will notice that our application exits with a small delay -

message received: message 0
processing...
notification sent
message received: message 1
processing...
notification sent
main exits

Instead of exiting immediately during processing, the output shows that notification for message 1 was sent to destination and main function returns normally after worker goroutine is finished.

Hence, we have successfully armed our application with graceful shutdown mechanism and increased reliability.

Thanks for stopping by! I'll catch you in the next post.

Keep learning, stay curious, and keep coding!