feat: factorise server run code (kyverno#151)

Signed-off-by: Charles-Edouard Brétéché <charles.edouard@nirmata.com>

pkg/commands/inject/command.go

@@ -4,20 +4,17 @@ import (
 	"context"
 	"crypto/tls"
 	"encoding/json"
-	"errors"
-	"fmt"
 	"net/http"
 	"time"
 
+	"github.com/kyverno/kyverno-envoy-plugin/pkg/server"
 	"github.com/kyverno/kyverno-envoy-plugin/pkg/server/handlers"
 	"github.com/kyverno/kyverno-envoy-plugin/pkg/sidecar"
 	"github.com/kyverno/kyverno-envoy-plugin/pkg/signals"
 	"github.com/spf13/cobra"
-	"go.uber.org/multierr"
 	jsonpatch "gomodules.xyz/jsonpatch/v2"
 	admissionv1 "k8s.io/api/admission/v1"
 	corev1 "k8s.io/api/core/v1"
-	"k8s.io/apimachinery/pkg/util/wait"
 )
 
 func Command() *cobra.Command {
@@ -29,9 +26,9 @@ func Command() *cobra.Command {
 		Use:   "sidecar-injector",
 		Short: "Responsible for injecting sidecars into pod containers",
 		RunE: func(cmd *cobra.Command, args []string) error {
-			mux := setupMux(sidecarImage)
-			server := setupServer(address, mux)
-			return runServer(context.Background(), server, certFile, keyFile)
+			return signals.Do(context.Background(), func(ctx context.Context) error {
+				return server.Run(ctx, setupServer(address, setupMux(sidecarImage)), certFile, keyFile)
+			})
 		},
 	}
 	command.Flags().StringVar(&address, "address", ":9443", "Address to listen on")
@@ -84,33 +81,3 @@ func setupServer(addr string, handler http.Handler) *http.Server {
 		IdleTimeout:       5 * time.Minute,
 	}
 }
-
-func runServer(ctx context.Context, server *http.Server, certFile, keyFile string) error {
-	var group wait.Group
-	err := func() error {
-		signalsCtx, signalsCancel := signals.Context(ctx)
-		defer signalsCancel()
-		var shutdownErr error
-		group.StartWithContext(signalsCtx, func(ctx context.Context) {
-			<-ctx.Done()
-			fmt.Println("Shutting down server...")
-			shutdownCtx, shutdownCancel := context.WithTimeout(context.Background(), 10*time.Second)
-			defer shutdownCancel()
-			shutdownErr = server.Shutdown(shutdownCtx)
-		})
-		fmt.Printf("Starting server at %s...\n", server.Addr)
-		var serveErr error
-		if certFile != "" && keyFile != "" {
-			serveErr = server.ListenAndServeTLS(certFile, keyFile)
-		} else {
-			serveErr = server.ListenAndServe()
-		}
-		if errors.Is(serveErr, http.ErrServerClosed) {
-			serveErr = nil
-		}
-		return multierr.Combine(serveErr, shutdownErr)
-	}()
-	group.Wait()
-	fmt.Println("Server stopped")
-	return err
-}

pkg/server/run.go

@@ -0,0 +1,56 @@
+package server
+
+import (
+	"context"
+	"errors"
+	"fmt"
+	"net/http"
+	"time"
+
+	"go.uber.org/multierr"
+	"k8s.io/apimachinery/pkg/util/wait"
+)
+
+func Run(ctx context.Context, server *http.Server, certFile, keyFile string) error {
+	defer fmt.Println("Server stopped")
+	// track shutdown error
+	var shutdownErr error
+	// track serve error
+	serveErr := func(ctx context.Context) error {
+		// create a wait group
+		var group wait.Group
+		// wait all tasks in the group are over
+		defer group.Wait()
+		// create a cancellable context
+		ctx, cancel := context.WithCancel(ctx)
+		// cancel context at the end
+		defer cancel()
+		// shutdown server when context is cancelled
+		group.StartWithContext(ctx, func(ctx context.Context) {
+			// wait context cancelled
+			<-ctx.Done()
+			fmt.Println("Server shutting down...")
+			// create a context with timeout
+			ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
+			defer cancel()
+			shutdownErr = server.Shutdown(ctx)
+		})
+		serve := func() error {
+			fmt.Printf("Server starting at %s...\n", server.Addr)
+			if certFile != "" && keyFile != "" {
+				// server over https
+				return server.ListenAndServeTLS(certFile, keyFile)
+			} else {
+				// server over http
+				return server.ListenAndServe()
+			}
+		}
+		// server closed is not an error
+		if err := serve(); !errors.Is(err, http.ErrServerClosed) {
+			return err
+		}
+		return nil
+	}(ctx)
+	// return error if any
+	return multierr.Combine(serveErr, shutdownErr)
+}

pkg/signals/context.go

@@ -4,8 +4,29 @@ import (
 	"context"
 	"os/signal"
 	"syscall"
+
+	"k8s.io/apimachinery/pkg/util/wait"
 )
 
 func Context(ctx context.Context) (context.Context, context.CancelFunc) {
 	return signal.NotifyContext(ctx, syscall.SIGINT, syscall.SIGTERM)
 }
+
+func Do(ctx context.Context, f func(context.Context) error) error {
+	// create a wait group
+	var group wait.Group
+	// wait all tasks in the group are over
+	defer group.Wait()
+	// create a signal aware context
+	ctx, stop := Context(ctx)
+	// cancel context and restore signals behaviour
+	defer stop()
+	// wait until context is cancelled or signals are triggered
+	group.StartWithContext(ctx, func(ctx context.Context) {
+		// restore signals behaviour (context has been cancelled at this point)
+		defer stop()
+		// wait signals are triggered
+		<-ctx.Done()
+	})
+	return f(ctx)
+}

pkg/signals/context_test.go

@@ -0,0 +1,15 @@
+package signals
+
+import (
+	"context"
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+)
+
+func TestDo(t *testing.T) {
+	err := Do(context.Background(), func(ctx context.Context) error {
+		return nil
+	})
+	assert.NoError(t, err)
+}