runtime: add ExtraDiscoveryOpts to runtime.Params

v1/runtime/runtime.go

@@ -249,6 +249,9 @@ type Params struct {
 	// Evaluation performance is affected in that data doesn't need to be converted to AST during evaluation.
 	// Only applicable when using the default in-memory store, and not when used together with the DiskStorage option.
 	ReadAstValuesFromStore bool
+
+	// ExtraDiscoveryOpts allows for passing options to the discovery plugin, as instantiated by the runtime.
+	ExtraDiscoveryOpts []func(*discovery.Discovery)
 }
 
 func (p *Params) regoVersion() ast.RegoVersion {
@@ -476,7 +479,14 @@ func NewRuntime(ctx context.Context, params Params) (*Runtime, error) {
 		return nil, fmt.Errorf("config error: %w", err)
 	}
 
-	disco, err := discovery.New(manager, discovery.Factories(registeredPlugins), discovery.Metrics(metrics), discovery.BootConfig(bootConfig))
+	opts := make([]func(*discovery.Discovery), 0, len(params.ExtraDiscoveryOpts)+3)
+	opts = append(opts, params.ExtraDiscoveryOpts...)
+	opts = append(opts,
+		discovery.Factories(registeredPlugins),
+		discovery.Metrics(metrics),
+		discovery.BootConfig(bootConfig),
+	)
+	disco, err := discovery.New(manager, opts...)
 	if err != nil {
 		return nil, fmt.Errorf("config error: %w", err)
 	}

v1/runtime/runtime_test.go

@@ -9,6 +9,7 @@ import (
 	"context"
 	"encoding/json"
 	"fmt"
+	"io"
 	"maps"
 	"net/http"
 	"net/http/httptest"
@@ -23,6 +24,7 @@ import (
 
 	"github.com/open-policy-agent/opa/internal/file/archive"
 	"github.com/open-policy-agent/opa/v1/loader"
+	"github.com/open-policy-agent/opa/v1/plugins/discovery"
 
 	"github.com/open-policy-agent/opa/internal/report"
 	"github.com/open-policy-agent/opa/v1/logging"
@@ -1656,3 +1658,40 @@ func TestRuntimeWithExplicitBadMetricConfiguration(t *testing.T) {
 		}
 	})
 }
+
+func TestExtraDiscoveryOpts(t *testing.T) {
+	ctx, cancel := context.WithTimeout(context.Background(), 2*time.Millisecond)
+	defer cancel() // NOTE(sr): The timeout will have been reached by the time `done` is closed.
+
+	called := false
+	params := NewParams()
+	params.Output = io.Discard
+	params.Addrs = &[]string{"localhost:0"}
+	params.GracefulShutdownPeriod = 1
+	params.Logger = logging.NewNoOpLogger()
+	params.ExtraDiscoveryOpts = []func(*discovery.Discovery){
+		func(*discovery.Discovery) { called = true },
+	}
+
+	rt, err := NewRuntime(ctx, params)
+	if err != nil {
+		t.Fatalf("Unexpected error %v", err)
+	}
+
+	initChannel := rt.Manager.ServerInitializedChannel()
+	done := make(chan struct{})
+	go func() {
+		rt.StartServer(ctx)
+		close(done)
+	}()
+	<-done
+	select {
+	case <-initChannel:
+	default:
+		t.Fatal("expected ServerInitializedChannel to be closed")
+	}
+
+	if !called {
+		t.Error("expected extra discovery option to be called")
+	}
+}