Add utility package for graceful shutdown context

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

Author: samiulsami

cmd/pgoutbox/main.go

@@ -17,7 +17,6 @@ limitations under the License.
 package main
 
 import (
-	"context"
 	"encoding/binary"
 	"fmt"
 	"log/slog"
@@ -33,6 +32,7 @@ import (
 	"kubeops.dev/pgoutbox/internal/telemetry"
 
 	"github.com/urfave/cli/v2"
+	"kubeops.dev/pgoutbox/internal/utils"
 )
 
 // GetVersion returns latest git hash of commit.
@@ -92,19 +92,8 @@ func main() {
 					return fmt.Errorf("initialize telemetry: %w", err)
 				}
 				defer func() {
-					shutdownCtx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
+					shutdownCtx, cancel := utils.ShutdownCtx(5 * 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)
 					}

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/utils"
 )
 
 // Logical decoding plugin.
@@ -136,19 +136,9 @@ func (l *Listener) InitHandlers(ctx context.Context) {
 
 	<-ctx.Done()
 
-	shutdownCtx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
+	shutdownCtx, cancel := utils.ShutdownCtx(5 * 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)
 	}
@@ -539,7 +529,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/utils/shutdown_ctx.go

@@ -0,0 +1,27 @@
+package utils
+
+import (
+	"context"
+	"os"
+	"os/signal"
+	"syscall"
+	"time"
+)
+
+// ShutdownCtx returns a context with timeout that is cancelled on SIGINT or SIGTERM signals.
+// With a forced shutdown upon receiving a second signal.
+// ref: https://github.com/appscode-cloud/cloud-be/blob/79fd250eeb79d47abdd2d38525a0b9b21e1920cf/common/util/signal.go
+func ShutdownCtx(timeout time.Duration) (context.Context, context.CancelFunc) {
+	shutdownCtx, cancel := context.WithTimeout(context.Background(), timeout)
+
+	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
+	}()
+
+	return shutdownCtx, cancel
+}