feat: Add a context.Context to the plugins.HAndle interface (kubernetes-sigs#1076)

* Added a context.Context to the plugins.Handle interface

Signed-off-by: Shmuel Kallner <[email protected]>

* Changes due to changes in internal APIs

Signed-off-by: Shmuel Kallner <[email protected]>

* Changes to tests due to changes in internal APIs

Signed-off-by: Shmuel Kallner <[email protected]>

---------

Signed-off-by: Shmuel Kallner <[email protected]>

Author: shmuelk

cmd/epp/runner/register.go

@@ -17,6 +17,8 @@ limitations under the License.
 package runner
 
 import (
+	"context"
+
 	"sigs.k8s.io/gateway-api-inference-extension/pkg/epp/plugins"
 	"sigs.k8s.io/gateway-api-inference-extension/pkg/epp/scheduling/framework/plugins/filter"
 	"sigs.k8s.io/gateway-api-inference-extension/pkg/epp/scheduling/framework/plugins/multi/prefix"
@@ -42,9 +44,15 @@ func RegisterAllPlugins() {
 
 // eppHandle is an implementation of the interface plugins.Handle
 type eppHandle struct {
+	ctx     context.Context
 	plugins plugins.HandlePlugins
 }
 
+// Context returns a context the plugins can use, if they need one
+func (h *eppHandle) Context() context.Context {
+	return h.ctx
+}
+
 // Plugins returns the sub-handle for working with instantiated plugins
 func (h *eppHandle) Plugins() plugins.HandlePlugins {
 	return h.plugins
@@ -79,8 +87,9 @@ func (h *eppHandlePlugins) GetAllPluginsWithNames() map[string]plugins.Plugin {
 	return h.thePlugins
 }
 
-func newEppHandle() *eppHandle {
+func newEppHandle(ctx context.Context) *eppHandle {
 	return &eppHandle{
+		ctx: ctx,
 		plugins: &eppHandlePlugins{
 			thePlugins: map[string]plugins.Plugin{},
 		},

cmd/epp/runner/runner.go

@@ -248,7 +248,7 @@ func (r *Runner) Run(ctx context.Context) error {
 
 	// ===================================================================
 
-	err = r.parseConfiguration()
+	err = r.parseConfiguration(ctx)
 	if err != nil {
 		setupLog.Error(err, "Failed to parse the configuration")
 		return err
@@ -344,14 +344,14 @@ func (r *Runner) initializeScheduler(datastore datastore.Datastore) (*scheduling
 	return scheduler, nil
 }
 
-func (r *Runner) parseConfiguration() error {
+func (r *Runner) parseConfiguration(ctx context.Context) error {
 	if len(*configText) != 0 || len(*configFile) != 0 {
 		theConfig, err := loader.LoadConfig([]byte(*configText), *configFile)
 		if err != nil {
 			return fmt.Errorf("failed to load the configuration - %w", err)
 		}
 
-		epp := newEppHandle()
+		epp := newEppHandle(ctx)
 
 		err = loader.LoadPluginReferences(theConfig.Plugins, epp)
 		if err != nil {

pkg/epp/common/config/loader/configloader_test.go

@@ -205,11 +205,12 @@ func TestLoadConfiguration(t *testing.T) {
 }
 
 func TestLoadPluginReferences(t *testing.T) {
+	ctx := context.Background()
 	theConfig, err := LoadConfig([]byte(successConfigText), "")
 	if err != nil {
 		t.Fatalf("LoadConfig returned unexpected error: %v", err)
 	}
-	handle := utils.NewTestHandle()
+	handle := utils.NewTestHandle(ctx)
 	err = LoadPluginReferences(theConfig.Plugins, handle)
 	if err != nil {
 		t.Fatalf("LoadPluginReferences returned unexpected error: %v", err)
@@ -227,15 +228,15 @@ func TestLoadPluginReferences(t *testing.T) {
 	if err != nil {
 		t.Fatalf("LoadConfig returned unexpected error: %v", err)
 	}
-	err = LoadPluginReferences(theConfig.Plugins, utils.NewTestHandle())
+	err = LoadPluginReferences(theConfig.Plugins, utils.NewTestHandle(ctx))
 	if err == nil {
 		t.Fatalf("LoadPluginReferences did not return the expected error")
 	}
 }
 
 func TestInstantiatePlugin(t *testing.T) {
 	plugSpec := configapi.PluginSpec{Type: "plover"}
-	_, err := instantiatePlugin(plugSpec, utils.NewTestHandle())
+	_, err := instantiatePlugin(plugSpec, utils.NewTestHandle(context.Background()))
 	if err == nil {
 		t.Fatalf("InstantiatePlugin did not return the expected error")
 	}
@@ -286,6 +287,8 @@ func TestLoadSchedulerConfig(t *testing.T) {
 
 	registerNeededPlgugins()
 
+	ctx := context.Background()
+
 	for _, test := range tests {
 		theConfig, err := LoadConfig([]byte(test.configText), "")
 		if err != nil {
@@ -294,7 +297,7 @@ func TestLoadSchedulerConfig(t *testing.T) {
 			}
 			t.Fatalf("LoadConfig returned unexpected error: %v", err)
 		}
-		handle := utils.NewTestHandle()
+		handle := utils.NewTestHandle(ctx)
 		err = LoadPluginReferences(theConfig.Plugins, handle)
 		if err != nil {
 			if test.wantErr {

pkg/epp/plugins/plugins.go

@@ -16,6 +16,8 @@ limitations under the License.
 
 package plugins
 
+import "context"
+
 // Plugin defines the interface for a plugin.
 // This interface should be embedded in all plugins across the code.
 type Plugin interface {
@@ -27,6 +29,9 @@ type Plugin interface {
 
 // Handle provides plugins a set of standard data and tools to work with
 type Handle interface {
+	// Context returns a context the plugins can use, if they need one
+	Context() context.Context
+
 	// Plugins returns the sub-handle for working with instantiated plugins
 	Plugins() HandlePlugins
 }

pkg/epp/scheduling/framework/plugins/filter/filter_test.go

@@ -406,7 +406,7 @@ func TestDecisionTreeFilterFactory(t *testing.T) {
 
 	kvCacheScorer := scorer.NewKVCacheScorer()
 
-	testHandle := utils.NewTestHandle()
+	testHandle := utils.NewTestHandle(context.Background())
 
 	testHandle.Plugins().AddPlugin("leastKvCache", leastKvCacheFilter)
 	testHandle.Plugins().AddPlugin("leastQueue", leastQueueFilter)

test/utils/handle.go

@@ -16,13 +16,23 @@ limitations under the License.
 
 package utils
 
-import "sigs.k8s.io/gateway-api-inference-extension/pkg/epp/plugins"
+import (
+	"context"
+
+	"sigs.k8s.io/gateway-api-inference-extension/pkg/epp/plugins"
+)
 
 // testHandle is an implmentation of plugins.Handle for test purposes
 type testHandle struct {
+	ctx     context.Context
 	plugins plugins.HandlePlugins
 }
 
+// Context returns a context the plugins can use, if they need one
+func (h *testHandle) Context() context.Context {
+	return h.ctx
+}
+
 func (h *testHandle) Plugins() plugins.HandlePlugins {
 	return h.plugins
 }
@@ -51,8 +61,9 @@ func (h *testHandlePlugins) GetAllPluginsWithNames() map[string]plugins.Plugin {
 	return h.thePlugins
 }
 
-func NewTestHandle() plugins.Handle {
+func NewTestHandle(ctx context.Context) plugins.Handle {
 	return &testHandle{
+		ctx: ctx,
 		plugins: &testHandlePlugins{
 			thePlugins: map[string]plugins.Plugin{},
 		},