Add utility package for graceful shutdown context

Signed-off-by: sami <sami@appscode.com>

Author: samiulsami

cmd/pgoutbox/main.go

@@ -17,14 +17,11 @@ limitations under the License.
 package main
 
 import (
-	"context"
 	"encoding/binary"
 	"fmt"
 	"log/slog"
 	"os"
-	"os/signal"
 	"runtime/debug"
-	"syscall"
 	"time"
 
 	"kubeops.dev/pgoutbox/apis"
@@ -33,6 +30,7 @@ import (
 	"kubeops.dev/pgoutbox/internal/telemetry"
 
 	"github.com/urfave/cli/v2"
+	"kubeops.dev/pgoutbox/internal/util"
 )
 
 // GetVersion returns latest git hash of commit.
@@ -73,7 +71,7 @@ func main() {
 			},
 		},
 		Action: func(c *cli.Context) error {
-			ctx, cancel := signal.NotifyContext(c.Context, syscall.SIGINT, syscall.SIGTERM)
+			ctx, cancel := util.SetupSignalContext()
 			defer cancel()
 
 			cfg, err := apis.InitConfig(c.String("config"))
@@ -91,24 +89,6 @@ func main() {
 				if err = telemetry.InitMetrics(ctx, version); err != nil {
 					return fmt.Errorf("initialize telemetry: %w", err)
 				}
-				defer func() {
-					shutdownCtx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
-					defer cancel()
-
-					// ref: https://github.com/appscode-cloud/cloud-be/blob/79fd250eeb79d47abdd2d38525a0b9b21e1920cf/common/util/signal.go
-					shutdownHandler := make(chan os.Signal, 2)
-					signal.Notify(shutdownHandler, os.Interrupt, syscall.SIGTERM)
-					go func() {
-						<-shutdownHandler
-						cancel()
-						<-shutdownHandler
-						os.Exit(1) // force exit upon receiving second signal
-					}()
-
-					if err := telemetry.Shutdown(shutdownCtx); err != nil {
-						slog.Error("telemetry shutdown failed", "err", err)
-					}
-				}()
 			}
 
 			pgxConn, pgConn, err := initPgxConnections(cfg.Database, logger, time.Minute*10)

internal/listener/listener.go

@@ -11,7 +11,6 @@ import (
 	"strconv"
 	"sync"
 	"sync/atomic"
-	"syscall"
 	"time"
 
 	"kubeops.dev/pgoutbox/apis"
@@ -20,6 +19,7 @@ import (
 	"github.com/jackc/pglogrepl"
 	"github.com/prometheus/client_golang/prometheus/promhttp"
 	"golang.org/x/sync/errgroup"
+	"kubeops.dev/pgoutbox/internal/telemetry"
 )
 
 // Logical decoding plugin.
@@ -136,22 +136,16 @@ func (l *Listener) InitHandlers(ctx context.Context) {
 
 	<-ctx.Done()
 
-	shutdownCtx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
+	shutdownCtx, cancel := context.WithTimeout(context.Background(), 30*time.Second)
 	defer cancel()
 
-	// ref: https://github.com/appscode-cloud/cloud-be/blob/79fd250eeb79d47abdd2d38525a0b9b21e1920cf/common/util/signal.go
-	shutdownHandler := make(chan os.Signal, 2)
-	signal.Notify(shutdownHandler, os.Interrupt, syscall.SIGTERM)
-	go func() {
-		<-shutdownHandler
-		cancel()
-		<-shutdownHandler
-		os.Exit(1) // force exit upon receiving second signal
-	}()
-
 	if err := srv.Shutdown(shutdownCtx); err != nil {
 		l.log.Error("http server shutdown error", "err", err)
 	}
+
+	if err := telemetry.Shutdown(shutdownCtx); err != nil {
+		l.log.Error("telemetry shutdown error", "err", err)
+	}
 }
 
 const contentTypeTextPlain = "text/plain"
@@ -406,7 +400,6 @@ func (l *Listener) processRawMessage(ctx context.Context, rawMsg []byte, txWAL *
 					return fmt.Errorf("publish: %w", err)
 				}
 				l.monitor.RecordPublishDuration(ctx, time.Since(publishStart).Seconds(), subjectName)
-
 				l.monitor.IncPublishedEvents(ctx, subjectName, event.Table)
 
 				l.log.Info(
@@ -539,7 +532,7 @@ func (l *Listener) readLSN() pglogrepl.LSN {
 func (l *Listener) setLSN(ctx context.Context, lsn pglogrepl.LSN) {
 	l.mu.Lock()
 	defer l.mu.Unlock()
-	defer l.monitor.RecordLSN(ctx, int64(l.lsn))
 
 	l.lsn = lsn
+	l.monitor.RecordLSN(ctx, int64(l.lsn))
 }

internal/listener/listener_test.go

@@ -539,7 +539,7 @@ func TestListener_Stream(t *testing.T) {
 			tt.setup()
 
 			ctx, cancel := context.WithTimeout(context.Background(), tt.args.timeout)
-			defer cancel()
+			_ = cancel
 
 			w := &Listener{
 				log:        logger,
@@ -643,7 +643,7 @@ func TestListener_Process(t *testing.T) {
 			setup: func() {
 				var cancel context.CancelFunc
 				ctx, cancel = context.WithTimeout(ctx, time.Millisecond*200)
-				defer cancel()
+				_ = cancel
 
 				setIsReplicationActive("slot1", false, nil)
 
@@ -681,7 +681,8 @@ func TestListener_Process(t *testing.T) {
 			setup: func() {
 				var cancel context.CancelFunc
 				ctx, cancel = context.WithTimeout(ctx, time.Millisecond*20)
-				defer cancel()
+				_ = cancel
+
 				setCreatePublication("pgoutbox", errors.New("some err"))
 				setGetSlotLSN("slot1", "100/200", nil)
 				setStartReplication(
@@ -716,7 +717,8 @@ func TestListener_Process(t *testing.T) {
 			setup: func() {
 				var cancel context.CancelFunc
 				ctx, cancel = context.WithTimeout(ctx, time.Millisecond*20)
-				defer cancel()
+				_ = cancel
+
 				setCreatePublication("pgoutbox", nil)
 				setGetSlotLSN("slot1", "100/200", errors.New("some err"))
 			},
@@ -739,7 +741,8 @@ func TestListener_Process(t *testing.T) {
 			setup: func() {
 				var cancel context.CancelFunc
 				ctx, cancel = context.WithTimeout(ctx, time.Millisecond*20)
-				defer cancel()
+				_ = cancel
+
 				setCreatePublication("pgoutbox", nil)
 				setGetSlotLSN("slot1", "", nil)
 				setCreateReplicationSlot(

internal/util/signal.go

@@ -0,0 +1,60 @@
+// ref: https://github.com/appscode-cloud/cloud-be/blob/79fd250eeb79d47abdd2d38525a0b9b21e1920cf/common/util/signal.go
+
+package util
+
+import (
+	"context"
+	"os"
+	"os/signal"
+	"syscall"
+)
+
+var (
+	onlyOneSignalHandler = make(chan struct{})
+	shutdownHandler      chan os.Signal
+	shutdownSignals      = []os.Signal{os.Interrupt, syscall.SIGTERM}
+)
+
+// SetupSignalHandler registered for SIGTERM and SIGINT. A stop channel is returned
+// which is closed on one of these signals. If a second signal is caught, the program
+// is terminated with exit code 1.
+// Only one of SetupSignalContext and SetupSignalHandler should be called, and only can
+// be called once.
+func SetupSignalHandler() <-chan struct{} {
+	c, _ := SetupSignalContext()
+	return c.Done()
+}
+
+// SetupSignalContext is the same as SetupSignalHandler, but a context.Context is returned.
+// Only one of SetupSignalContext and SetupSignalHandler should be called, and only can
+// be called once.
+func SetupSignalContext() (context.Context, context.CancelFunc) {
+	close(onlyOneSignalHandler) // panics when called twice
+
+	shutdownHandler = make(chan os.Signal, 2)
+
+	ctx, cancel := context.WithCancel(context.Background())
+	signal.Notify(shutdownHandler, shutdownSignals...)
+	go func() {
+		<-shutdownHandler
+		cancel()
+		<-shutdownHandler
+		os.Exit(1) // second signal. Exit directly.
+	}()
+
+	return ctx, cancel
+}
+
+// RequestShutdown emulates a received event that is considered as shutdown signal (SIGTERM/SIGINT)
+// This returns whether a handler was notified
+func RequestShutdown() bool {
+	if shutdownHandler != nil {
+		select {
+		case shutdownHandler <- shutdownSignals[0]:
+			return true
+		default:
+		}
+	}
+
+	return false
+}