explicitly handle SIGINT and SIGTERM

Author: joelanford

cmd/opm/main.go

@@ -1,7 +1,10 @@
 package main
 
 import (
+	"context"
 	"os"
+	"os/signal"
+	"syscall"
 
 	utilerrors "k8s.io/apimachinery/pkg/util/errors"
 
@@ -12,7 +15,11 @@ import (
 func main() {
 	showAlphaHelp := os.Getenv("HELP_ALPHA") == "true"
 	cmd := root.NewCmd(showAlphaHelp)
-	if err := cmd.Execute(); err != nil {
+
+	ctx, cancel := signal.NotifyContext(context.Background(), syscall.SIGINT, syscall.SIGTERM)
+	defer cancel()
+
+	if err := cmd.ExecuteContext(ctx); err != nil {
 		agg, ok := err.(utilerrors.Aggregate)
 		if !ok {
 			os.Exit(1)

cmd/opm/registry/serve.go

@@ -17,7 +17,6 @@ import (
 
 	"github.com/operator-framework/operator-registry/pkg/api"
 	"github.com/operator-framework/operator-registry/pkg/lib/dns"
-	"github.com/operator-framework/operator-registry/pkg/lib/graceful"
 	"github.com/operator-framework/operator-registry/pkg/lib/log"
 	"github.com/operator-framework/operator-registry/pkg/lib/tmp"
 	"github.com/operator-framework/operator-registry/pkg/server"
@@ -54,6 +53,9 @@ func newRegistryServeCmd() *cobra.Command {
 }
 
 func serveFunc(cmd *cobra.Command, _ []string) error {
+	ctx, cancel := context.WithCancel(cmd.Context())
+	defer cancel()
+
 	// Immediately set up termination log
 	terminationLogPath, err := cmd.Flags().GetString("termination-log")
 	if err != nil {
@@ -93,19 +95,23 @@ func serveFunc(cmd *cobra.Command, _ []string) error {
 		return err
 	}
 
-	if _, err := db.ExecContext(context.TODO(), `PRAGMA soft_heap_limit=1`); err != nil {
+	if _, err := db.ExecContext(ctx, `PRAGMA soft_heap_limit=1`); err != nil {
 		logger.WithError(err).Warnf("error setting soft heap limit for sqlite")
 	}
 
 	// migrate to the latest version
-	if err := migrate(cmd, db); err != nil {
+	shouldSkipMigrate, err := cmd.Flags().GetBool("skip-migrate")
+	if err != nil {
+		return err
+	}
+	if err := migrate(ctx, shouldSkipMigrate, db); err != nil {
 		logger.WithError(err).Warnf("couldn't migrate db")
 	}
 
 	store := sqlite.NewSQLLiteQuerierFromDb(db, sqlite.OmitManifests(true))
 
 	// sanity check that the db is available
-	tables, err := store.ListTables(context.TODO())
+	tables, err := store.ListTables(ctx)
 	if err != nil {
 		logger.WithError(err).Warnf("couldn't list tables in db")
 	}
@@ -142,19 +148,18 @@ func serveFunc(cmd *cobra.Command, _ []string) error {
 	api.RegisterRegistryServer(s, server.NewRegistryServer(store))
 	health.RegisterHealthServer(s, server.NewHealthServer())
 	reflection.Register(s)
-	logger.Info("serving registry")
-	return graceful.Shutdown(logger, func() error {
-		return s.Serve(lis)
-	}, func() {
+
+	go func() {
+		<-ctx.Done()
+		logger.Info("shutting down server")
 		s.GracefulStop()
-	})
+	}()
+
+	logger.Info("serving registry")
+	return s.Serve(lis)
 }
 
-func migrate(cmd *cobra.Command, db *sql.DB) error {
-	shouldSkipMigrate, err := cmd.Flags().GetBool("skip-migrate")
-	if err != nil {
-		return err
-	}
+func migrate(ctx context.Context, shouldSkipMigrate bool, db *sql.DB) error {
 	if shouldSkipMigrate {
 		return nil
 	}
@@ -167,5 +172,5 @@ func migrate(cmd *cobra.Command, db *sql.DB) error {
 		return fmt.Errorf("failed to load migrator")
 	}
 
-	return migrator.Migrate(context.TODO())
+	return migrator.Migrate(ctx)
 }

cmd/opm/serve/serve.go

@@ -23,7 +23,6 @@ import (
 	"github.com/operator-framework/operator-registry/pkg/api"
 	"github.com/operator-framework/operator-registry/pkg/cache"
 	"github.com/operator-framework/operator-registry/pkg/lib/dns"
-	"github.com/operator-framework/operator-registry/pkg/lib/graceful"
 	"github.com/operator-framework/operator-registry/pkg/lib/log"
 	"github.com/operator-framework/operator-registry/pkg/server"
 )
@@ -91,6 +90,9 @@ will not be reflected in the served content.
 }
 
 func (s *serve) run(ctx context.Context) error {
+	ctx, cancel := context.WithCancel(ctx)
+	defer cancel()
+
 	mainLogger := s.logger.Dup()
 	p := newProfilerInterface(s.pprofAddr, mainLogger)
 	if err := p.startEndpoint(); err != nil {
@@ -169,15 +171,16 @@ func (s *serve) run(ctx context.Context) error {
 	mainLogger.Info("serving registry")
 	p.stopCpuProfileCache()
 
-	return graceful.Shutdown(s.logger, func() error {
-		return grpcServer.Serve(lis)
-	}, func() {
+	go func() {
+		<-ctx.Done()
+		mainLogger.Info("shutting down server")
 		grpcServer.GracefulStop()
 		if err := p.stopEndpoint(ctx); err != nil {
 			mainLogger.Warnf("error shutting down pprof server: %v", err)
 		}
-	})
+	}()
 
+	return grpcServer.Serve(lis)
 }
 
 // manages an HTTP pprof endpoint served by `server`,