From b153b7f83716e031f175eb4abe84948d55ba6fb2 Mon Sep 17 00:00:00 2001
From: Jan Tytgat <jan.tytgat@corelayer.eu>
Date: Wed, 23 Apr 2025 12:52:11 +0200
Subject: [PATCH] Configure signals for graceful shutdown

Signed-off-by: Jan Tytgat <jan.tytgat@corelayer.eu>
---
 application/application.go | 53 +++++++++++++++++++++++++++++++++++---
 examples/simple/main.go    | 13 ++++++----
 2 files changed, 57 insertions(+), 9 deletions(-)

diff --git a/application/application.go b/application/application.go
index b43f9b3..7e5c3a9 100644
--- a/application/application.go
+++ b/application/application.go
@@ -3,6 +3,7 @@ package application
 import (
 	"context"
 	"fmt"
+	"os/signal"
 
 	"github.com/spf13/cobra"
 )
@@ -20,19 +21,63 @@ func New(c Config) (Application, error) {
 	}
 
 	return &application{
-		cmd: cmd,
+		cmd:    cmd,
+		config: c,
 	}, nil
 }
 
 type application struct {
-	cmd *cobra.Command
+	cmd    *cobra.Command
+	config Config
 }
 
 func (a *application) Start(ctx context.Context) error {
-	return a.cmd.ExecuteContext(ctx)
+
+	sigCtx, sigStop := signal.NotifyContext(ctx, a.config.ShutdownSignals...)
+	defer sigStop() // Ensure that this gets called.
+
+	// Result channel for command
+	chExe := make(chan error)
+	go a.execute(sigCtx, chExe)
+
+	select {
+	// sigCtx.Done() returns a channel that will have a message
+	// when the context is cancelled. We wait for that signal, which means
+	// we received the signal, or our context was cancelled for some other reason.
+	case <-sigCtx.Done():
+		sigStop()
+		return a.Shutdown()
+	case err := <-chExe:
+		return err
+	}
 }
 
 func (a *application) Shutdown() error {
-	fmt.Println("Shutdown")
+	switch a.config.EnableGracefulShutdown {
+	case true:
+		return a.gracefulShutdown()
+	case false:
+		return a.shutdown()
+		// default:
+		// 	return a.shutdown()
+	}
+	return nil
+}
+
+func (a *application) execute(ctx context.Context, chErr chan error) {
+	chErr <- a.cmd.ExecuteContext(ctx)
+}
+
+func (a *application) gracefulShutdown() error {
+	fmt.Println("graceful shutdown")
+
+	shutdownCtx, shutdownCancel := context.WithTimeout(context.Background(), a.config.ShutdownTimeout)
+	defer shutdownCancel()
+	<-shutdownCtx.Done()
+	return nil
+}
+
+func (a *application) shutdown() error {
+	fmt.Println("shutdown")
 	return nil
 }
diff --git a/examples/simple/main.go b/examples/simple/main.go
index 1aae799..747a5d8 100644
--- a/examples/simple/main.go
+++ b/examples/simple/main.go
@@ -3,6 +3,7 @@ package main
 import (
 	"context"
 	"fmt"
+	"time"
 
 	"github.com/spf13/cobra"
 
@@ -20,15 +21,17 @@ func main() {
 		Title:                  "Main Test",
 		Banner:                 "",
 		Version:                "0.1.0-alpha.0+metadata.20101112",
-		EnableGracefulShutdown: false,
+		EnableGracefulShutdown: true,
 		OverrideRunE: func(cmd *cobra.Command, args []string) error {
 			fmt.Println("overrideRunE")
+			time.Sleep(5 * time.Second)
+			fmt.Println("overrideRunE done")
 			return nil
 		},
-		PersistentPreRunE:        nil,
-		PersistentPostRunE:       nil,
-		ShutdownSignals:          nil,
-		ShutdownTimeout:          0,
+		PersistentPreRunE:  nil,
+		PersistentPostRunE: nil,
+		// ShutdownSignals:          []os.Signal{syscall.SIGINT, syscall.SIGTERM, syscall.SIGHUP, syscall.SIGQUIT},
+		ShutdownTimeout:          1 * time.Second,
 		SubCommands:              nil,
 		SubCommandInitializeFunc: nil,
 		ValidArgs:                nil,