Use buffered channel based approach to synchronize go routines. Also add unit tests

Author: rahulgurnani

pkg/epp/requestcontrol/plugin_executor.go

@@ -27,48 +27,62 @@ import (
 
 // executePluginsAsDAG executes PrepareData plugins as a DAG based on their dependencies asynchronously.
 // So, a plugin is executed only after all its dependencies have been executed.
-// If there is a cycle or other error in the DAG, it returns an error.
+// If there is a cycle or any plugin fails with error, it returns an error.
 func executePluginsAsDAG(plugins []PrepareDataPlugin, ctx context.Context, request *schedulingtypes.LLMRequest, pods []schedulingtypes.Pod) error {
 	// Build the DAG
 	// The error validation happens on startup when loading the config. So, here there should not be any error.
 	dag, err := prepareDataGraph(plugins)
 	if err != nil {
 		return err
 	}
-	// Execute the DAG
-
-	// Channels to signal plugin execution completion.
-	pluginExecuted := make(map[string]chan error)
+	// Create a readonly map of plugin name to plugin instance.
 	nameToNode := map[string]PrepareDataPlugin{}
 	for _, plugin := range plugins {
-		pluginExecuted[plugin.TypedName().String()] = make(chan error)
 		nameToNode[plugin.TypedName().String()] = plugin
 	}
+	// Execute the DAG
 
+	// Channels to signal plugin execution completion.
+	pluginExecuted := make(map[string]chan error)
+	// The capacity of the channel is equal to the number of dependents + 1 (for itself).
+	capacityMap := make(map[string]int)
+	for pluginName := range dag {
+		capacityMap[pluginName]++
+		for _, dep := range dag[pluginName] {
+			capacityMap[dep]++
+		}
+	}
+	for pluginName, capacity := range capacityMap {
+		pluginExecuted[pluginName] = make(chan error, capacity)
+	}
 	for pluginName, dependents := range dag {
 		// Execute plugins based on dependencies.
 		//  Wait for the dependencies to complete before executing a plugin.
-		go func() {
-			for _, dep := range dependents {
-				err, open := <-pluginExecuted[dep]
-				if !open {
-					continue
-				}
-				if err != nil {
-					// If a dependency failed, propagate the error and do not execute this plugin.
-					pluginExecuted[pluginName] <- fmt.Errorf("dependency plugin %s failed: %w", dep, err)
+		go func(pN string, depds []string) {
+			// Wait for dependencies to complete.
+			for _, dep := range depds {
+				// Wait for the dependency plugin to signal completion.
+				if err, closed := <-pluginExecuted[dep]; closed {
+					if err != nil {
+						pluginExecuted[pN] <- fmt.Errorf("dependency plugin %s failed: %w", dep, err)
+						return
+					}
 				}
 			}
-			// Signal that the plugin has been executed.
-			defer close(pluginExecuted[pluginName])
-
-			pluginExecuted[pluginName] <- nameToNode[pluginName].PrepareRequestData(ctx, request, pods)
-		}()
+			res := nameToNode[pN].PrepareRequestData(ctx, request, pods)
+			for range cap(pluginExecuted[pN]) {
+				// Notify all dependents about the completion.
+				pluginExecuted[pN] <- res
+			}
+		}(pluginName, dependents)
 	}
+
+	// Check for errors in plugin execution.
+	// This will also ensure that all plugins have completed execution before returning.
 	for pluginName := range dag {
 		err := <-pluginExecuted[pluginName]
 		if err != nil {
-			return errors.New("prepare data plugin " + pluginName + " failed: " + err.Error())
+			return fmt.Errorf("prepare data plugin %s failed: %v", pluginName, err)
 		}
 	}
 	return nil
@@ -78,7 +92,6 @@ func executePluginsAsDAG(plugins []PrepareDataPlugin, ctx context.Context, reque
 func prepareDataPluginsWithTimeout(timeout time.Duration, plugins []PrepareDataPlugin,
 	ctx context.Context, request *schedulingtypes.LLMRequest, pods []schedulingtypes.Pod) error {
 	errCh := make(chan error, 1)
-	// Execute plugins sequentially in a separate goroutine
 	go func() {
 		errCh <- executePluginsAsDAG(plugins, ctx, request, pods)
 	}()

pkg/epp/requestcontrol/plugin_executor_test.go

@@ -19,6 +19,7 @@ package requestcontrol
 import (
 	"context"
 	"errors"
+	"sync"
 	"testing"
 	"time"
 
@@ -157,3 +158,138 @@ func TestPrepareDataPluginsWithTimeout(t *testing.T) {
 		})
 	}
 }
+
+type dagTestPlugin struct {
+	mockPrepareRequestDataPlugin
+	produces map[string]any
+	consumes map[string]any
+	execTime time.Time
+	mu       sync.Mutex
+}
+
+func (p *dagTestPlugin) PrepareRequestData(ctx context.Context, request *schedulingtypes.LLMRequest, pods []schedulingtypes.Pod) error {
+	p.mu.Lock()
+	defer p.mu.Unlock()
+	p.execTime = time.Now()
+	return p.mockPrepareRequestDataPlugin.PrepareRequestData(ctx, request, pods)
+}
+
+func (p *dagTestPlugin) Produces() map[string]any {
+	return p.produces
+}
+
+func (p *dagTestPlugin) Consumes() map[string]any {
+	return p.consumes
+}
+
+func TestExecutePluginsAsDAG(t *testing.T) {
+	pluginA := &dagTestPlugin{
+		mockPrepareRequestDataPlugin: mockPrepareRequestDataPlugin{name: "A", delay: 20 * time.Millisecond},
+		produces:                     map[string]any{"keyA": nil},
+	}
+	pluginB := &dagTestPlugin{
+		mockPrepareRequestDataPlugin: mockPrepareRequestDataPlugin{name: "B"},
+		consumes:                     map[string]any{"keyA": nil},
+		produces:                     map[string]any{"keyB": nil},
+	}
+	pluginC := &dagTestPlugin{
+		mockPrepareRequestDataPlugin: mockPrepareRequestDataPlugin{name: "C"},
+		consumes:                     map[string]any{"keyB": nil},
+	}
+	pluginD := &dagTestPlugin{
+		mockPrepareRequestDataPlugin: mockPrepareRequestDataPlugin{name: "D"},
+		consumes:                     map[string]any{"keyA": nil},
+	}
+	pluginE := &dagTestPlugin{
+		mockPrepareRequestDataPlugin: mockPrepareRequestDataPlugin{name: "E"},
+	}
+	pluginFail := &dagTestPlugin{
+		mockPrepareRequestDataPlugin: mockPrepareRequestDataPlugin{name: "Fail", returnErr: errors.New("plugin failed")},
+		produces:                     map[string]any{"keyFail": nil},
+	}
+	pluginDependsOnFail := &dagTestPlugin{
+		mockPrepareRequestDataPlugin: mockPrepareRequestDataPlugin{name: "DependsOnFail"},
+		consumes:                     map[string]any{"keyFail": nil},
+	}
+
+	testCases := []struct {
+		name      string
+		plugins   []PrepareDataPlugin
+		expectErr bool
+		checkFunc func(t *testing.T, plugins []PrepareDataPlugin)
+	}{
+		{
+			name:    "no plugins",
+			plugins: []PrepareDataPlugin{},
+		},
+		{
+			name:    "simple linear dependency (A -> B -> C)",
+			plugins: []PrepareDataPlugin{pluginA, pluginB, pluginC},
+			checkFunc: func(t *testing.T, plugins []PrepareDataPlugin) {
+				pA := plugins[0].(*dagTestPlugin)
+				pB := plugins[1].(*dagTestPlugin)
+				pC := plugins[2].(*dagTestPlugin)
+
+				assert.True(t, pA.executed, "Plugin A should have been executed")
+				assert.True(t, pB.executed, "Plugin B should have been executed")
+				assert.True(t, pC.executed, "Plugin C should have been executed")
+
+				assert.True(t, pB.execTime.After(pA.execTime), "Plugin B should execute after A")
+				assert.True(t, pC.execTime.After(pB.execTime), "Plugin C should execute after B")
+			},
+		},
+		{
+			name:    "DAG with multiple dependencies (A -> B, A -> D) and one independent (E)",
+			plugins: []PrepareDataPlugin{pluginA, pluginB, pluginD, pluginE},
+			checkFunc: func(t *testing.T, plugins []PrepareDataPlugin) {
+				pA := plugins[0].(*dagTestPlugin)
+				pB := plugins[1].(*dagTestPlugin)
+				pD := plugins[2].(*dagTestPlugin)
+				pE := plugins[3].(*dagTestPlugin)
+
+				assert.True(t, pA.executed, "Plugin A should have been executed")
+				assert.True(t, pB.executed, "Plugin B should have been executed")
+				assert.True(t, pD.executed, "Plugin D should have been executed")
+				assert.True(t, pE.executed, "Plugin E should have been executed")
+
+				assert.True(t, pB.execTime.After(pA.execTime), "Plugin B should execute after A")
+				assert.True(t, pD.execTime.After(pA.execTime), "Plugin D should execute after A")
+			},
+		},
+		{
+			name:      "dependency fails",
+			plugins:   []PrepareDataPlugin{pluginFail, pluginDependsOnFail},
+			expectErr: true,
+			checkFunc: func(t *testing.T, plugins []PrepareDataPlugin) {
+				pF := plugins[0].(*dagTestPlugin)
+				pDOF := plugins[1].(*dagTestPlugin)
+
+				assert.True(t, pF.executed, "Failing plugin should have been executed")
+				assert.False(t, pDOF.executed, "Plugin depending on fail should not be executed")
+			},
+		},
+	}
+
+	for _, tc := range testCases {
+		t.Run(tc.name, func(t *testing.T) {
+			// Reset execution state for plugins
+			for _, p := range tc.plugins {
+				plugin := p.(*dagTestPlugin)
+				plugin.executed = false
+				plugin.execTime = time.Time{}
+			}
+
+			err := executePluginsAsDAG(tc.plugins, context.Background(), &schedulingtypes.LLMRequest{}, nil)
+
+			if tc.expectErr {
+				assert.Error(t, err)
+			} else {
+				assert.NoError(t, err)
+			}
+
+			if tc.checkFunc != nil {
+				tc.checkFunc(t, tc.plugins)
+			}
+		})
+	}
+}