Server: Implement timeout waiting for server to shut down (#33333)

* tests: Undo cleanup in goroutine

Signed-off-by: Arve Knudsen <arve.knudsen@gmail.com>

* Server: Implement timeout waiting for it to shut down

Signed-off-by: Arve Knudsen <arve.knudsen@gmail.com>

pkg/cmd/grafana-server/main.go

@@ -1,6 +1,7 @@
 package main
 
 import (
+	"context"
 	"errors"
 	"flag"
 	"fmt"
@@ -156,7 +157,9 @@ func executeServer(configFile, homePath, pidFile, packaging string, traceDiagnos
 		return err
 	}
 
-	go listenToSystemSignals(s)
+	ctx := context.Background()
+
+	go listenToSystemSignals(ctx, s)
 
 	if err := s.Run(); err != nil {
 		code := s.ExitCode(err)
@@ -179,7 +182,7 @@ func validPackaging(packaging string) string {
 	return "unknown"
 }
 
-func listenToSystemSignals(s *server.Server) {
+func listenToSystemSignals(ctx context.Context, s *server.Server) {
 	signalChan := make(chan os.Signal, 1)
 	sighupChan := make(chan os.Signal, 1)
 
@@ -193,7 +196,11 @@ func listenToSystemSignals(s *server.Server) {
 				fmt.Fprintf(os.Stderr, "Failed to reload loggers: %s\n", err)
 			}
 		case sig := <-signalChan:
-			s.Shutdown(fmt.Sprintf("System signal: %s", sig))
+			ctx, cancel := context.WithTimeout(ctx, 30*time.Second)
+			defer cancel()
+			if err := s.Shutdown(ctx, fmt.Sprintf("System signal: %s", sig)); err != nil {
+				fmt.Fprintf(os.Stderr, "Timed out waiting for server to shut down\n")
+			}
 			return
 		}
 	}

pkg/server/server.go

@@ -222,17 +222,23 @@ func (s *Server) Run() error {
 // Shutdown initiates Grafana graceful shutdown. This shuts down all
 // running background services. Since Run blocks Shutdown supposed to
 // be run from a separate goroutine.
-func (s *Server) Shutdown(reason string) {
+func (s *Server) Shutdown(ctx context.Context, reason string) error {
+	var err error
 	s.shutdownOnce.Do(func() {
 		s.log.Info("Shutdown started", "reason", reason)
 		// Call cancel func to stop services.
 		s.shutdownFn()
-		// Can introduce termination timeout here if needed over incoming Context,
+		// Wait for server to shut down
-		// but this will require changing Shutdown method signature to accept context
+		select {
-		// and return an error - caller can exit with code > 0 then.
+		case <-s.shutdownFinished:
-		// I.e. sth like server.Shutdown(context.WithTimeout(...), reason).
+			s.log.Debug("Finished waiting for server to shut down")
-		<-s.shutdownFinished
+		case <-ctx.Done():
+			s.log.Warn("Timed out while waiting for server to shut down")
+			err = fmt.Errorf("timeout waiting for shutdown")
+		}
 	})
+
+	return err
 }
 
 // ExitCode returns an exit code for a given error.

pkg/server/server_test.go

@@ -4,6 +4,7 @@ import (
 	"context"
 	"errors"
 	"testing"
+	"time"
 
 	"github.com/grafana/grafana/pkg/registry"
 
@@ -83,6 +84,8 @@ func TestServer_Run_Error(t *testing.T) {
 }
 
 func TestServer_Shutdown(t *testing.T) {
+	ctx := context.Background()
+
 	s := testServer()
 	services := []*registry.Descriptor{
 		{
@@ -100,14 +103,24 @@ func TestServer_Shutdown(t *testing.T) {
 		services: services,
 	}
 
+	ch := make(chan error)
+
 	go func() {
+		defer close(ch)
+
 		// Wait until all services launched.
 		for _, svc := range services {
 			<-svc.Instance.(*testService).started
 		}
-		s.Shutdown("test interrupt")
+		ctx, cancel := context.WithTimeout(ctx, 3*time.Second)
+		defer cancel()
+		err := s.Shutdown(ctx, "test interrupt")
+		ch <- err
 	}()
 	err := s.Run()
 	require.NoError(t, err)
 	require.Zero(t, s.ExitCode(err))
+
+	err = <-ch
+	require.NoError(t, err)
 }

pkg/tests/testinfra/testinfra.go

@@ -1,6 +1,7 @@
 package testinfra
 
 import (
+	"context"
 	"fmt"
 	"io/ioutil"
 	"net"
@@ -9,6 +10,7 @@ import (
 	"path/filepath"
 	"strings"
 	"testing"
+	"time"
 
 	"github.com/grafana/grafana/pkg/infra/fs"
 	"github.com/grafana/grafana/pkg/models"
@@ -25,6 +27,8 @@ import (
 func StartGrafana(t *testing.T, grafDir, cfgPath string, sqlStore *sqlstore.SQLStore) string {
 	t.Helper()
 
+	ctx := context.Background()
+
 	origSQLStore := registry.GetService(sqlstore.ServiceName)
 	t.Cleanup(func() {
 		registry.Register(origSQLStore)
@@ -58,7 +62,11 @@ func StartGrafana(t *testing.T, grafDir, cfgPath string, sqlStore *sqlstore.SQLS
 		}
 	}()
 	t.Cleanup(func() {
-		go server.Shutdown("test cleanup")
+		ctx, cancel := context.WithTimeout(ctx, 3*time.Second)
+		defer cancel()
+		if err := server.Shutdown(ctx, "test cleanup"); err != nil {
+			t.Error("Timed out waiting on server to shut down")
+		}
 	})
 
 	// Wait for Grafana to be ready