From 8a5c0cfdc00dc3ed2e59edf4dddb94bcbea3279a Mon Sep 17 00:00:00 2001
From: Will Browne <wbrowne@users.noreply.github.com>
Date: Fri, 19 Apr 2024 08:22:14 +0200
Subject: [PATCH] Plugins: Pass cancellable context during API server creation
 (#86545)

---
 pkg/cmd/grafana/apiserver/cmd.go             | 9 ++++++++-
 pkg/cmd/grafana/apiserver/server.go          | 5 +++--
 pkg/services/apiserver/standalone/factory.go | 4 ++--
 3 files changed, 13 insertions(+), 5 deletions(-)

diff --git a/pkg/cmd/grafana/apiserver/cmd.go b/pkg/cmd/grafana/apiserver/cmd.go
index cb3530b7c2c..9adab683f42 100644
--- a/pkg/cmd/grafana/apiserver/cmd.go
+++ b/pkg/cmd/grafana/apiserver/cmd.go
@@ -1,6 +1,7 @@
 package apiserver
 
 import (
+	"context"
 	"os"
 
 	"github.com/spf13/cobra"
@@ -54,8 +55,14 @@ func newCommandStartExampleAPIServer(o *APIServerOptions, stopCh <-chan struct{}
 			// TODO: Fix so that TracingOptions.ApplyTo happens before or during loadAPIGroupBuilders.
 			tracer := newLateInitializedTracingService()
 
+			ctx, cancel := context.WithCancel(c.Context())
+			go func() {
+				<-stopCh
+				cancel()
+			}()
+
 			// Load each group from the args
-			if err := o.loadAPIGroupBuilders(tracer, apis); err != nil {
+			if err := o.loadAPIGroupBuilders(ctx, tracer, apis); err != nil {
 				return err
 			}
 
diff --git a/pkg/cmd/grafana/apiserver/server.go b/pkg/cmd/grafana/apiserver/server.go
index ed064d1cef1..1ead6da20e5 100644
--- a/pkg/cmd/grafana/apiserver/server.go
+++ b/pkg/cmd/grafana/apiserver/server.go
@@ -1,6 +1,7 @@
 package apiserver
 
 import (
+	"context"
 	"fmt"
 	"io"
 	"net"
@@ -50,10 +51,10 @@ func newAPIServerOptions(out, errOut io.Writer) *APIServerOptions {
 	}
 }
 
-func (o *APIServerOptions) loadAPIGroupBuilders(tracer tracing.Tracer, apis []schema.GroupVersion) error {
+func (o *APIServerOptions) loadAPIGroupBuilders(ctx context.Context, tracer tracing.Tracer, apis []schema.GroupVersion) error {
 	o.builders = []builder.APIGroupBuilder{}
 	for _, gv := range apis {
-		api, err := o.factory.MakeAPIServer(tracer, gv)
+		api, err := o.factory.MakeAPIServer(ctx, tracer, gv)
 		if err != nil {
 			return err
 		}
diff --git a/pkg/services/apiserver/standalone/factory.go b/pkg/services/apiserver/standalone/factory.go
index 32009a224e0..0af7bf0132c 100644
--- a/pkg/services/apiserver/standalone/factory.go
+++ b/pkg/services/apiserver/standalone/factory.go
@@ -35,7 +35,7 @@ type APIServerFactory interface {
 	GetEnabled(runtime []RuntimeConfig) ([]schema.GroupVersion, error)
 
 	// Make an API server for a given group+version
-	MakeAPIServer(tracer tracing.Tracer, gv schema.GroupVersion) (builder.APIGroupBuilder, error)
+	MakeAPIServer(ctx context.Context, tracer tracing.Tracer, gv schema.GroupVersion) (builder.APIGroupBuilder, error)
 }
 
 // Zero dependency provider for testing
@@ -67,7 +67,7 @@ func (p *DummyAPIFactory) ApplyTo(config *genericapiserver.RecommendedConfig) er
 	return nil
 }
 
-func (p *DummyAPIFactory) MakeAPIServer(tracer tracing.Tracer, gv schema.GroupVersion) (builder.APIGroupBuilder, error) {
+func (p *DummyAPIFactory) MakeAPIServer(_ context.Context, tracer tracing.Tracer, gv schema.GroupVersion) (builder.APIGroupBuilder, error) {
 	if gv.Version != "v0alpha1" {
 		return nil, fmt.Errorf("only alpha supported now")
 	}