Graceful Shutdown Go App

graceful.go

package graceful

import (
	"context"
	"log"
	"os"
	"os/signal"
	"sync"
	"syscall"
	"time"
)

// operation is a clean up function on shutting down
type Operation func(ctx context.Context)

// gracefulShutdown waits for termination syscalls and doing clean up operations after received it
func Shutdown(ctx context.Context, timeout time.Duration, ops []Operation) <-chan struct{} {
	wait := make(chan struct{})
	go func() {
		s := make(chan os.Signal, 1)

		// add any other syscalls that you want to be notified with.
		signal.Notify(s, os.Interrupt, syscall.SIGINT, syscall.SIGTERM)
		<-s

		// set timeout for the ops to be done to prevent system hang.
		// TODO: should to consider whether to timeout and exit directly.
		timeoutFunc := time.AfterFunc(timeout, func() {
			log.Printf("timeout %d ms has been elapsed, force exit", timeout.Milliseconds())
			os.Exit(1)
		})

		defer timeoutFunc.Stop()

		var wg sync.WaitGroup

		// Do the operations asynchronously to save time.
		for _, v := range ops {
			wg.Add(1)

			op := v
			go func() {
				defer wg.Done()
				op(ctx)
			}()
		}

		wg.Wait()

		close(wait)
	}()

	return wait
}

main.go

package main

import (
    "context"
    "fmt"
    "log"
    "net/http"
    "time"

    "git.shiyou.kingsoft.com/go/graceful"
)

func main() {
    server := &http.Server{Addr: fmt.Sprintf(":%d", 9527)}

    go func() {
        server.ListenAndServe()
    }()

    wait := graceful.Shutdown(
        context.Background(),
        2 * time.Second,
        []graceful.Operation{
            func(ctx context.Context) {
                server.Shutdown(ctx)
            },
        }
    )

    <-wait

    log.Println("Good bye ...")
}