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.