github.com/authzed/spicedb@v1.32.1-0.20240520085336-ebda56537386/pkg/cmd/signals.go (about)

     1  package cmd
     2  
     3  import (
     4  	"context"
     5  	"os"
     6  	"os/signal"
     7  	"syscall"
     8  	"time"
     9  
    10  	log "github.com/authzed/spicedb/internal/logging"
    11  )
    12  
    13  // SignalContextWithGracePeriod creates a new context that will be cancelled
    14  // when an interrupt/SIGTERM signal is received and the provided grace period
    15  // subsequently finishes.
    16  func SignalContextWithGracePeriod(ctx context.Context, gracePeriod time.Duration) context.Context {
    17  	newCtx, cancelfn := context.WithCancel(ctx)
    18  	go func() {
    19  		signalctx, _ := signal.NotifyContext(newCtx, os.Interrupt, syscall.SIGTERM)
    20  		<-signalctx.Done()
    21  		log.Ctx(ctx).Info().Msg("received interrupt")
    22  
    23  		if gracePeriod > 0 {
    24  			interruptGrace, _ := signal.NotifyContext(context.Background(), os.Interrupt)
    25  			graceTimer := time.NewTimer(gracePeriod)
    26  
    27  			log.Ctx(ctx).Info().Stringer("timeout", gracePeriod).Msg("starting shutdown grace period")
    28  
    29  			select {
    30  			case <-graceTimer.C:
    31  			case <-interruptGrace.Done():
    32  				log.Ctx(ctx).Warn().Msg("interrupted shutdown grace period")
    33  			}
    34  		}
    35  		log.Ctx(ctx).Info().Msg("shutting down")
    36  		cancelfn()
    37  	}()
    38  
    39  	return newCtx
    40  }