refactor to sync request-thread confinement

Concludes the EPP multithreading model by removing background
goroutines, channels, and select statements during plugin execution.
All data producers now run synchronously and sequentially on the main
request-handling goroutine.

Introduces three-phase request orchestration in the Director relative to
the admission boundary. Request attributes are wrapped in an immutable,
read-only snapshot during wait-queue blocking to allow lockless
concurrent reads by queue policies, and restored to a writable local map
post-admission.

Implements individual plugin panic recovery via deferred recovery stubs
to isolate plugin crashes and protect router process stability.

Signed-off-by: Luke Van Drie <lukevandrie@google.com>

Author: LukeAVanDrie

cmd/epp/runner/runner.go

@@ -690,15 +690,16 @@ func (r *Runner) parseConfigurationPhaseTwo(ctx context.Context, rawConfig *conf
 		}
 	}
 
-	// Sort data plugins in DAG order (topological sort). Also check DAG for cycles.
+	// Sort and partition data plugins into Pre-Admission and Post-Admission in DAG order.
 	// This must run after auto-created producers are added so they are included in the ordering.
-	dag, err := datalayer.ValidateAndOrderDataDependencies(handle.GetAllPlugins())
+	preAdmission, postAdmission, err := datalayer.CompilePipeline(handle.GetAllPlugins())
 	if err != nil {
 		return nil, fmt.Errorf("failed to load the configuration - %w", err)
 	}
 
-	// The plugins will be executed in topologically sorted order to ensure that data is produced before it is consumed.
-	r.requestControlConfig.OrderDataProducerPlugins(dag)
+	// Map and register the topologically sorted pre-admission and post-admission data producers.
+	r.requestControlConfig.OrderPreAdmissionDataProducers(preAdmission)
+	r.requestControlConfig.OrderPostAdmissionDataProducers(postAdmission)
 
 	r.parser = handlers.NewParser(cfg.ParserConfig)
 	logger.Info("loaded configuration from file/text successfully")

pkg/epp/requestcontrol/director.go

@@ -213,25 +213,42 @@ func (d *Director) HandleRequest(ctx context.Context, reqCtx *handlers.RequestCo
 		FairnessID:       fairnessID,
 		Objectives:       requestObjectives,
 		RequestSizeBytes: reqCtx.RequestSize,
-		Attributes:       fwkdl.NewAttributes(),
+		Attributes:       fwkdl.NewLocalAttributes(), // Lockless map confined to the main request-handling goroutine
 	}
 
 	logger = logger.WithValues("objectiveKey", reqCtx.ObjectiveKey, "incomingModelName", reqCtx.IncomingModelName, "targetModelName", reqCtx.TargetModelName, "priority", infObjective.Spec.Priority)
 	ctx = log.IntoContext(ctx, logger)
 	logger.V(logutil.DEBUG).Info("LLM request assembled")
 
+	// --- Pre-Admission Data Production (Pre-Queue) ---
+	if err := d.runPreAdmissionDataProducers(ctx, reqCtx.SchedulingRequest); err != nil {
+		return reqCtx, err
+	}
+
 	if err := d.runPreAdmissionPlugins(ctx, reqCtx.SchedulingRequest); err != nil {
 		return reqCtx, err
 	}
+
 	if reqCtx.SchedulingRequest.FairnessID == "" {
 		reqCtx.SchedulingRequest.FairnessID = metadata.DefaultFairnessID
 	}
 
+	// --- Admission Control Queuing (The Queue Boundary) ---
+	// Capture a reference to the original writable attributes map.
+	localAttributesSnapshot := reqCtx.SchedulingRequest.Attributes
+
+	// Wrap the attributes into a thread-safe read-only snapshot during queue blocking.
+	reqCtx.SchedulingRequest.Attributes = fwkdl.NewReadOnlyAttributes(localAttributesSnapshot)
+
 	// Admit may block until flow control admits the request.
 	if err := d.admissionController.Admit(ctx, reqCtx, priority); err != nil {
 		return reqCtx, err
 	}
 
+	// --- Post-Admission Routing & Scheduling (Post-Queue) ---
+	// Restore the writable local attributes map populated with Pre-Admission traits.
+	reqCtx.SchedulingRequest.Attributes = localAttributesSnapshot
+
 	endpointCandidates := d.endpointCandidates.Locate(ctx, reqCtx.Request.Metadata)
 	if len(endpointCandidates) == 0 {
 		return reqCtx, errcommon.Error{
@@ -241,11 +258,12 @@ func (d *Director) HandleRequest(ctx context.Context, reqCtx *handlers.RequestCo
 	}
 
 	snapshotOfCandidatePods := d.toSchedulerEndpoints(endpointCandidates)
-	// Prepare per request data by running DataProducer plugins.
-	err = d.runDataProducerPlugins(ctx, reqCtx.SchedulingRequest, snapshotOfCandidatePods)
+	// Prepare per request data by running Post-Admission DataProducer plugins sequentially on the
+	// main request-handling goroutine
+	err = d.runPostAdmissionDataProducers(ctx, reqCtx.SchedulingRequest, snapshotOfCandidatePods)
 	if err != nil {
-		// Don't fail the request if DataProducer plugins fail.
-		logger.Error(err, "failed to prepare per request data")
+		// Don't fail the request if Post-Admission DataProducer plugins fail.
+		logger.Error(err, "failed to prepare post-admission per request data")
 	}
 
 	// Run admit request plugins
@@ -523,12 +541,19 @@ func (d *Director) runPreAdmissionPlugins(ctx context.Context, request *fwksched
 	return nil
 }
 
-func (d *Director) runDataProducerPlugins(ctx context.Context,
-	request *fwksched.InferenceRequest, endpoints []fwksched.Endpoint) error {
-	if len(d.requestControlPlugins.dataProducerPlugins) == 0 {
+func (d *Director) runPreAdmissionDataProducers(ctx context.Context, request *fwksched.InferenceRequest) error {
+	if len(d.requestControlPlugins.preAdmissionDataProducers) == 0 {
+		return nil
+	}
+	emptyEndpoints := []fwksched.Endpoint{}
+	return dataProducerPluginsWithTimeout(ctx, dataProducerTimeout, d.requestControlPlugins.preAdmissionDataProducers, request, emptyEndpoints)
+}
+
+func (d *Director) runPostAdmissionDataProducers(ctx context.Context, request *fwksched.InferenceRequest, endpoints []fwksched.Endpoint) error {
+	if len(d.requestControlPlugins.postAdmissionDataProducers) == 0 {
 		return nil
 	}
-	return dataProducerPluginsWithTimeout(ctx, dataProducerTimeout, d.requestControlPlugins.dataProducerPlugins, request, endpoints)
+	return dataProducerPluginsWithTimeout(ctx, dataProducerTimeout, d.requestControlPlugins.postAdmissionDataProducers, request, endpoints)
 }
 
 func (d *Director) runAdmissionPlugins(ctx context.Context,

pkg/epp/requestcontrol/director_test.go

@@ -65,10 +65,14 @@ var (
 // --- Mocks ---
 
 type mockAdmissionController struct {
-	admitErr error
+	admitErr      error
+	admitCallback func(request *fwksched.InferenceRequest)
 }
 
-func (m *mockAdmissionController) Admit(context.Context, *handlers.RequestContext, int) error {
+func (m *mockAdmissionController) Admit(ctx context.Context, reqCtx *handlers.RequestContext, priority int) error {
+	if m.admitCallback != nil {
+		m.admitCallback(reqCtx.SchedulingRequest)
+	}
 	return m.admitErr
 }
 
@@ -141,6 +145,27 @@ func newMockDataProducerPlugin(name string) *mockDataProducerPlugin {
 	}
 }
 
+type mockDataProducer struct {
+	name            string
+	produces        map[fwkplugin.DataKey]any
+	produceCallback func(ctx context.Context, request *fwksched.InferenceRequest, endpoints []fwksched.Endpoint) error
+}
+
+func (m *mockDataProducer) TypedName() fwkplugin.TypedName {
+	return fwkplugin.TypedName{Type: "mock-data-producer", Name: m.name}
+}
+
+func (m *mockDataProducer) Produces() map[fwkplugin.DataKey]any {
+	return m.produces
+}
+
+func (m *mockDataProducer) Produce(ctx context.Context, request *fwksched.InferenceRequest, endpoints []fwksched.Endpoint) error {
+	if m.produceCallback != nil {
+		return m.produceCallback(ctx, request, endpoints)
+	}
+	return nil
+}
+
 type mockAdmissionPlugin struct {
 	typedName   fwkplugin.TypedName
 	denialError error
@@ -328,7 +353,10 @@ func TestDirector_HandleRequest(t *testing.T) {
 		preRequestPlugin        *mockPreRequestPlugin
 		wantMutatedBody         map[string]any
 		fairnessIDHeader        string // If non-empty, set as metadata.FlowFairnessIDKey on the incoming request.
+		preAdmissionProducers   []fwkrc.DataProducer
+		postAdmissionProducers  []fwkrc.DataProducer
 		wantFairnessID          string // If non-empty, asserted against returnedReqCtx.SchedulingRequest.FairnessID.
+		checkFn                 func(t *testing.T, reqCtx *handlers.RequestContext)
 	}{
 		{
 			name: "successful completions request",
@@ -734,6 +762,86 @@ func TestDirector_HandleRequest(t *testing.T) {
 			wantErrCode:            errcommon.Internal,
 			inferenceObjectiveName: objectiveName,
 		},
+		{
+			name: "pre-admission data producer failure aborts request",
+			reqBodyMap: map[string]any{
+				"model":  model,
+				"prompt": "prompt with failed pre-producer",
+			},
+			mockAdmissionController: &mockAdmissionController{admitErr: nil},
+			inferenceObjectiveName:  objectiveName,
+			preAdmissionProducers: []fwkrc.DataProducer{
+				&mockDataProducer{
+					name: "failed-pre-producer",
+					produceCallback: func(ctx context.Context, request *fwksched.InferenceRequest, endpoints []fwksched.Endpoint) error {
+						return errcommon.Error{
+							Code: errcommon.Internal,
+							Msg:  "simulated pre-admission data producer error",
+						}
+					},
+				},
+			},
+			wantErrCode: errcommon.Internal,
+		},
+		{
+			name: "attributes wrapped in read-only container during queue blocking",
+			reqBodyMap: map[string]any{
+				"model":  model,
+				"prompt": "prompt for queue blocking safety",
+			},
+			mockAdmissionController: &mockAdmissionController{
+				admitErr: nil,
+				admitCallback: func(req *fwksched.InferenceRequest) {
+					assert.PanicsWithValue(t, "cannot write to ReadOnlyAttributes", func() {
+						req.Attributes.Put("illegal_key", "illegal_value")
+					})
+					val, ok := req.Attributes.Get("pre_key")
+					assert.True(t, ok)
+					assert.Equal(t, "pre_val", val)
+				},
+			},
+			schedulerMockSetup: func(m *mockScheduler) {
+				m.scheduleResults = defaultSuccessfulScheduleResults
+			},
+			inferenceObjectiveName: objectiveName,
+			preAdmissionProducers: []fwkrc.DataProducer{
+				&mockDataProducer{
+					name: "pre-producer",
+					produces: map[fwkplugin.DataKey]any{fwkplugin.NewDataKey("pre_key", "pre-producer"): nil},
+					produceCallback: func(ctx context.Context, request *fwksched.InferenceRequest, endpoints []fwksched.Endpoint) error {
+						request.Attributes.Put("pre_key", "pre_val")
+						return nil
+					},
+				},
+			},
+			postAdmissionProducers: []fwkrc.DataProducer{
+				&mockDataProducer{
+					name: "post-producer",
+					produces: map[fwkplugin.DataKey]any{fwkplugin.NewDataKey("post_key", "post-producer"): nil},
+					produceCallback: func(ctx context.Context, request *fwksched.InferenceRequest, endpoints []fwksched.Endpoint) error {
+						request.Attributes.Put("post_key", "post_val")
+						return nil
+					},
+				},
+			},
+			checkFn: func(t *testing.T, returnedReqCtx *handlers.RequestContext) {
+				require.NotNil(t, returnedReqCtx.SchedulingRequest)
+				assert.NotPanics(t, func() {
+					returnedReqCtx.SchedulingRequest.Attributes.Put("writable_now", "success")
+				})
+				v1, ok := returnedReqCtx.SchedulingRequest.Attributes.Get("pre_key")
+				assert.True(t, ok)
+				assert.Equal(t, "pre_val", v1)
+
+				v2, ok := returnedReqCtx.SchedulingRequest.Attributes.Get("post_key")
+				assert.True(t, ok)
+				assert.Equal(t, "post_val", v2)
+
+				v3, ok := returnedReqCtx.SchedulingRequest.Attributes.Get("writable_now")
+				assert.True(t, ok)
+				assert.Equal(t, "success", v3)
+			},
+		},
 	}
 
 	period := time.Second
@@ -781,6 +889,12 @@ func TestDirector_HandleRequest(t *testing.T) {
 					test.schedulerMockSetup(mockSched)
 				}
 				config := NewConfig()
+				if len(test.preAdmissionProducers) > 0 {
+					config = config.WithPreAdmissionDataProducers(test.preAdmissionProducers...)
+				}
+				if len(test.postAdmissionProducers) > 0 {
+					config = config.WithPostAdmissionDataProducers(test.postAdmissionProducers...)
+				}
 				if test.dataProducerPlugin != nil {
 					config = config.WithDataProducerPlugins(test.dataProducerPlugin)
 				}
@@ -870,6 +984,9 @@ func TestDirector_HandleRequest(t *testing.T) {
 						t.Errorf("reqCtx.Request.RawBody mismatch (-want +got):\n%s", diff)
 					}
 				}
+				if test.checkFn != nil {
+					test.checkFn(t, returnedReqCtx)
+				}
 				assert.Equal(t, len(reqCtx.Request.RawBody), reqCtx.RequestSize)
 			})
 		}
@@ -1448,7 +1565,7 @@ type testResponseStreaming struct {
 	respsOnStreaming      []*fwkrc.Response
 	targetPodsOnStreaming []string
 
-	// Legacy fields for existing tests if any, but better to update them
+	// Legacy fields maintained for backward compatibility.
 	lastRespOnStreaming      *fwkrc.Response
 	lastTargetPodOnStreaming string
 }

pkg/epp/requestcontrol/plugin_executor.go

@@ -22,42 +22,49 @@ import (
 	"fmt"
 	"time"
 
+	"sigs.k8s.io/controller-runtime/pkg/log"
+
 	fwkrc "github.com/llm-d/llm-d-router/pkg/epp/framework/interface/requestcontrol"
 	fwksched "github.com/llm-d/llm-d-router/pkg/epp/framework/interface/scheduling"
 )
 
-// executePluginsAsDAG executes DataProducer 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 any plugin fails with error, it returns an error.
+// executePluginsAsDAG executes DataProducer plugins sequentially based on their dependencies.
+// If any plugin fails with error or panics, it returns an error.
 func executePluginsAsDAG(ctx context.Context, plugins []fwkrc.DataProducer, request *fwksched.InferenceRequest, endpoints []fwksched.Endpoint) error {
 	for _, plugin := range plugins {
-		if err := plugin.Produce(ctx, request, endpoints); err != nil {
+		if err := ctx.Err(); err != nil {
+			return err
+		}
+
+		err := func() (prodErr error) {
+			defer func() {
+				if r := recover(); r != nil {
+					prodErr = fmt.Errorf("panic in DataProducer %q: %v", plugin.TypedName().String(), r)
+					log.FromContext(ctx).Error(prodErr, "panic caught during plugin execution", "plugin", plugin.TypedName().String())
+				}
+			}()
+			return plugin.Produce(ctx, request, endpoints)
+		}()
+		if err != nil {
 			return fmt.Errorf("DataProducer %q failed: %w", plugin.TypedName().String(), err)
 		}
 	}
 	return nil
 }
 
-// dataProducerPluginsWithTimeout executes DataProducer plugins with a timeout.
-// The child context is cancelled when the timeout fires so plugins can observe cancellation
-// (e.g. abort outbound HTTP calls) and avoid committing state after the director has moved on.
+// dataProducerPluginsWithTimeout executes DataProducer plugins sequentially with a timeout.
+// The child context is cancelled when the timeout fires so plugins can observe cancellation.
 func dataProducerPluginsWithTimeout(ctx context.Context, timeout time.Duration, plugins []fwkrc.DataProducer,
 	request *fwksched.InferenceRequest, endpoints []fwksched.Endpoint) error {
 	ctx, cancel := context.WithTimeout(ctx, timeout)
 	defer cancel()
 
-	errCh := make(chan error, 1)
-	go func() {
-		errCh <- executePluginsAsDAG(ctx, plugins, request, endpoints)
-	}()
-
-	select {
-	case err := <-errCh:
-		return err
-	case <-ctx.Done():
+	err := executePluginsAsDAG(ctx, plugins, request, endpoints)
+	if err != nil {
 		if errors.Is(ctx.Err(), context.DeadlineExceeded) {
 			return fmt.Errorf("DataProducer execution timed out: %w", ctx.Err())
 		}
-		return ctx.Err()
+		return err
 	}
+	return nil
 }

pkg/epp/requestcontrol/plugin_executor_test.go

@@ -94,8 +94,7 @@ func (p *ctxObservingPlugin) Produces() map[fwkplugin.DataKey]any { return nil }
 // child context passed to plugins is cancelled with DeadlineExceeded when the
 // timeout fires. Without this cancellation, a slow plugin would continue
 // executing past the director's deadline and potentially commit state after
-// downstream hooks have already observed an "empty" state — the root cause of
-// the orphan-decrement drift we're fixing in the predicted-latency producer.
+// downstream hooks have already observed an "empty" state, preventing state drift.
 func TestDataProducerPluginsWithTimeout_CancelsPluginContext(t *testing.T) {
 	plugin := &ctxObservingPlugin{name: "slow", block: time.Second}
 	plugin.wg.Add(1)
@@ -349,3 +348,38 @@ func TestExecutePluginsAsDAG(t *testing.T) {
 		})
 	}
 }
+
+type panickingDataProducerPlugin struct {
+	name     string
+	executed bool
+}
+
+func (m *panickingDataProducerPlugin) TypedName() fwkplugin.TypedName {
+	return fwkplugin.TypedName{Type: "mock", Name: m.name}
+}
+
+func (m *panickingDataProducerPlugin) Produce(context.Context, *fwksched.InferenceRequest, []fwksched.Endpoint) error {
+	m.executed = true
+	panic("simulate database connection failure")
+}
+
+func (m *panickingDataProducerPlugin) Produces() map[fwkplugin.DataKey]any {
+	return nil
+}
+
+func TestExecutePluginsAsDAG_PanicRecovery(t *testing.T) {
+	panicker := &panickingDataProducerPlugin{name: "panicker"}
+	nextPlugin := &executorMockDataProducerPlugin{name: "next"}
+
+	err := executePluginsAsDAG(
+		context.Background(),
+		[]fwkrc.DataProducer{panicker, nextPlugin},
+		&fwksched.InferenceRequest{},
+		nil,
+	)
+
+	assert.Error(t, err)
+	assert.Contains(t, err.Error(), "panic in DataProducer \"panicker/mock\": simulate database connection failure")
+	assert.True(t, panicker.executed)
+	assert.False(t, nextPlugin.executed, "Subsequent plugin should not be executed after a panic")
+}