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/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,
-		// but this will require changing Shutdown method signature to accept context
-		// and return an error - caller can exit with code > 0 then.
-		// I.e. sth like server.Shutdown(context.WithTimeout(...), reason).
-		<-s.shutdownFinished
+		// Wait for server to shut down
+		select {
+		case <-s.shutdownFinished:
+			s.log.Debug("Finished waiting for server to shut down")
+		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)
 }