fix[pipeline]: use existing context to propagate cancellation

Author: henrichter

internal/scheduling/cinder/filter_weigher_pipeline_controller.go

@@ -129,7 +129,7 @@ func (c *FilterWeigherPipelineController) process(ctx context.Context, decision
 		return err
 	}
 
-	result, err := pipeline.Run(request)
+	result, err := pipeline.Run(ctx, request)
 	if err != nil {
 		log.Error(err, "failed to run pipeline")
 		return err

internal/scheduling/lib/filter_monitor.go

@@ -38,6 +38,6 @@ func (fm *FilterMonitor[RequestType]) Init(ctx context.Context, client client.Cl
 }
 
 // Run the filter and observe its execution.
-func (fm *FilterMonitor[RequestType]) Run(traceLog *slog.Logger, request RequestType) (*FilterWeigherPipelineStepResult, error) {
-	return fm.monitor.RunWrapped(traceLog, request, fm.filter)
+func (fm *FilterMonitor[RequestType]) Run(ctx context.Context, traceLog *slog.Logger, request RequestType) (*FilterWeigherPipelineStepResult, error) {
+	return fm.monitor.RunWrapped(ctx, traceLog, request, fm.filter)
 }

internal/scheduling/lib/filter_monitor_test.go

@@ -23,7 +23,7 @@ func TestMonitorFilter(t *testing.T) {
 		InitFunc: func(ctx context.Context, cl client.Client, step v1alpha1.FilterSpec) error {
 			return nil
 		},
-		RunFunc: func(traceLog *slog.Logger, request mockFilterWeigherPipelineRequest) (*FilterWeigherPipelineStepResult, error) {
+		RunFunc: func(_ context.Context, traceLog *slog.Logger, request mockFilterWeigherPipelineRequest) (*FilterWeigherPipelineStepResult, error) {
 			return &FilterWeigherPipelineStepResult{
 				Activations: map[string]float64{"host1": 0.5, "host2": 1.0},
 			}, nil
@@ -79,7 +79,7 @@ func TestFilterMonitor_Init(t *testing.T) {
 func TestFilterMonitor_Run(t *testing.T) {
 	runCalled := false
 	mockFilter := &mockFilter[mockFilterWeigherPipelineRequest]{
-		RunFunc: func(traceLog *slog.Logger, request mockFilterWeigherPipelineRequest) (*FilterWeigherPipelineStepResult, error) {
+		RunFunc: func(_ context.Context, traceLog *slog.Logger, request mockFilterWeigherPipelineRequest) (*FilterWeigherPipelineStepResult, error) {
 			runCalled = true
 			return &FilterWeigherPipelineStepResult{
 				Activations: map[string]float64{"host1": 0.5, "host2": 1.0},
@@ -102,7 +102,7 @@ func TestFilterMonitor_Run(t *testing.T) {
 		Weights: map[string]float64{"host1": 0.1, "host2": 0.2, "host3": 0.3},
 	}
 
-	result, err := fm.Run(slog.Default(), request)
+	result, err := fm.Run(t.Context(), slog.Default(), request)
 	if err != nil {
 		t.Errorf("expected no error, got %v", err)
 	}

internal/scheduling/lib/filter_test.go

@@ -16,7 +16,7 @@ import (
 
 type mockFilter[RequestType FilterWeigherPipelineRequest] struct {
 	InitFunc func(ctx context.Context, client client.Client, step v1alpha1.FilterSpec) error
-	RunFunc  func(traceLog *slog.Logger, request RequestType) (*FilterWeigherPipelineStepResult, error)
+	RunFunc  func(ctx context.Context, traceLog *slog.Logger, request RequestType) (*FilterWeigherPipelineStepResult, error)
 }
 
 func (m *mockFilter[RequestType]) Init(ctx context.Context, client client.Client, step v1alpha1.FilterSpec) error {
@@ -25,11 +25,11 @@ func (m *mockFilter[RequestType]) Init(ctx context.Context, client client.Client
 	}
 	return m.InitFunc(ctx, client, step)
 }
-func (m *mockFilter[RequestType]) Run(traceLog *slog.Logger, request RequestType) (*FilterWeigherPipelineStepResult, error) {
+func (m *mockFilter[RequestType]) Run(ctx context.Context, traceLog *slog.Logger, request RequestType) (*FilterWeigherPipelineStepResult, error) {
 	if m.RunFunc == nil {
 		return &FilterWeigherPipelineStepResult{}, nil
 	}
-	return m.RunFunc(traceLog, request)
+	return m.RunFunc(ctx, traceLog, request)
 }
 
 // filterTestOptions implements FilterWeigherPipelineStepOpts for testing.

internal/scheduling/lib/filter_validation.go

@@ -30,8 +30,8 @@ func validateFilter[RequestType FilterWeigherPipelineRequest](filter Filter[Requ
 }
 
 // Run the filter and validate what happens.
-func (s *FilterValidator[RequestType]) Run(traceLog *slog.Logger, request RequestType) (*FilterWeigherPipelineStepResult, error) {
-	result, err := s.Filter.Run(traceLog, request)
+func (s *FilterValidator[RequestType]) Run(ctx context.Context, traceLog *slog.Logger, request RequestType) (*FilterWeigherPipelineStepResult, error) {
+	result, err := s.Filter.Run(ctx, traceLog, request)
 	if err != nil {
 		return nil, err
 	}

internal/scheduling/lib/filter_validation_test.go

@@ -151,7 +151,7 @@ func TestFilterValidator_Run(t *testing.T) {
 	for _, tt := range tests {
 		t.Run(tt.name, func(t *testing.T) {
 			filter := &mockFilter[mockFilterWeigherPipelineRequest]{
-				RunFunc: func(traceLog *slog.Logger, request mockFilterWeigherPipelineRequest) (*FilterWeigherPipelineStepResult, error) {
+				RunFunc: func(_ context.Context, traceLog *slog.Logger, request mockFilterWeigherPipelineRequest) (*FilterWeigherPipelineStepResult, error) {
 					return tt.runResult, tt.runError
 				},
 			}
@@ -161,7 +161,7 @@ func TestFilterValidator_Run(t *testing.T) {
 			}
 			traceLog := slog.Default()
 
-			result, err := validator.Run(traceLog, request)
+			result, err := validator.Run(t.Context(), traceLog, request)
 
 			if tt.expectError && err == nil {
 				t.Error("expected error but got nil")

internal/scheduling/lib/filter_weigher_pipeline.go

@@ -19,7 +19,7 @@ import (
 
 type FilterWeigherPipeline[RequestType FilterWeigherPipelineRequest] interface {
 	// Run the scheduling pipeline with the given request.
-	Run(request RequestType) (v1alpha1.DecisionResult, error)
+	Run(ctx context.Context, request RequestType) (v1alpha1.DecisionResult, error)
 }
 
 // Pipeline of scheduler steps.
@@ -132,6 +132,7 @@ func InitNewFilterWeigherPipeline[RequestType FilterWeigherPipelineRequest](
 // During this process, the request is mutated to only include the
 // remaining hosts.
 func (p *filterWeigherPipeline[RequestType]) runFilters(
+	ctx context.Context,
 	log *slog.Logger,
 	request RequestType,
 ) (filteredRequest RequestType) {
@@ -141,7 +142,7 @@ func (p *filterWeigherPipeline[RequestType]) runFilters(
 		filter := p.filters[filterName]
 		stepLog := log.With("filter", filterName)
 		stepLog.Info("scheduler: running filter")
-		result, err := filter.Run(stepLog, filteredRequest)
+		result, err := filter.Run(ctx, stepLog, filteredRequest)
 		if errors.Is(err, ErrStepSkipped) {
 			stepLog.Info("scheduler: filter skipped")
 			continue
@@ -160,6 +161,7 @@ func (p *filterWeigherPipeline[RequestType]) runFilters(
 
 // Execute weighers and collect their activations by step name.
 func (p *filterWeigherPipeline[RequestType]) runWeighers(
+	ctx context.Context,
 	log *slog.Logger,
 	filteredRequest RequestType,
 ) map[string]map[string]float64 {
@@ -173,7 +175,7 @@ func (p *filterWeigherPipeline[RequestType]) runWeighers(
 		wg.Go(func() {
 			stepLog := log.With("weigher", weigherName)
 			stepLog.Info("scheduler: running weigher")
-			result, err := weigher.Run(stepLog, filteredRequest)
+			result, err := weigher.Run(ctx, stepLog, filteredRequest)
 			if errors.Is(err, ErrStepSkipped) {
 				stepLog.Info("scheduler: weigher skipped")
 				return
@@ -243,7 +245,7 @@ func (s *filterWeigherPipeline[RequestType]) sortHostsByWeights(weights map[stri
 }
 
 // Evaluate the pipeline and return a list of hosts in order of preference.
-func (p *filterWeigherPipeline[RequestType]) Run(request RequestType) (v1alpha1.DecisionResult, error) {
+func (p *filterWeigherPipeline[RequestType]) Run(ctx context.Context, request RequestType) (v1alpha1.DecisionResult, error) {
 	slogArgs := request.GetTraceLogArgs()
 	slogArgsAny := make([]any, 0, len(slogArgs))
 	for _, arg := range slogArgs {
@@ -260,7 +262,7 @@ func (p *filterWeigherPipeline[RequestType]) Run(request RequestType) (v1alpha1.
 
 	// Run filters first to reduce the number of hosts.
 	// Any weights assigned to filtered out hosts are ignored.
-	filteredRequest := p.runFilters(traceLog, request)
+	filteredRequest := p.runFilters(ctx, traceLog, request)
 	traceLog.Info(
 		"scheduler: finished filters",
 		"remainingHosts", filteredRequest.GetHosts(),
@@ -271,7 +273,7 @@ func (p *filterWeigherPipeline[RequestType]) Run(request RequestType) (v1alpha1.
 	for _, host := range filteredRequest.GetHosts() {
 		remainingWeights[host] = inWeights[host]
 	}
-	stepWeights := p.runWeighers(traceLog, filteredRequest)
+	stepWeights := p.runWeighers(ctx, traceLog, filteredRequest)
 	outWeights := p.applyWeights(stepWeights, remainingWeights)
 	traceLog.Info("scheduler: output weights", "weights", outWeights)
 

internal/scheduling/lib/filter_weigher_pipeline_step.go

@@ -29,7 +29,7 @@ type FilterWeigherPipelineStep[RequestType FilterWeigherPipelineRequest] interfa
 	//
 	// A traceLog is provided that contains the global request id and should
 	// be used to log the step's execution.
-	Run(traceLog *slog.Logger, request RequestType) (*FilterWeigherPipelineStepResult, error)
+	Run(ctx context.Context, traceLog *slog.Logger, request RequestType) (*FilterWeigherPipelineStepResult, error)
 }
 
 // Common base for all steps that provides some functionality

internal/scheduling/lib/filter_weigher_pipeline_step_monitor.go

@@ -4,6 +4,7 @@
 package lib
 
 import (
+	"context"
 	"fmt"
 	"log/slog"
 	"maps"
@@ -63,6 +64,7 @@ func monitorStep[RequestType FilterWeigherPipelineRequest](stepName string, m Fi
 
 // Run the step and observe its execution.
 func (s *FilterWeigherPipelineStepMonitor[RequestType]) RunWrapped(
+	ctx context.Context,
 	traceLog *slog.Logger,
 	request RequestType,
 	step FilterWeigherPipelineStep[RequestType],
@@ -74,7 +76,7 @@ func (s *FilterWeigherPipelineStepMonitor[RequestType]) RunWrapped(
 	}
 
 	inWeights := request.GetWeights()
-	stepResult, err := step.Run(traceLog, request)
+	stepResult, err := step.Run(ctx, traceLog, request)
 	if err != nil {
 		return nil, err
 	}

internal/scheduling/lib/filter_weigher_pipeline_step_monitor_test.go

@@ -4,6 +4,7 @@
 package lib
 
 import (
+	"context"
 	"log/slog"
 	"os"
 	"testing"
@@ -28,7 +29,7 @@ func TestStepMonitorRun(t *testing.T) {
 		removedHostsObserver: removedHostsObserver,
 	}
 	step := &mockWeigher[mockFilterWeigherPipelineRequest]{
-		RunFunc: func(traceLog *slog.Logger, request mockFilterWeigherPipelineRequest) (*FilterWeigherPipelineStepResult, error) {
+		RunFunc: func(_ context.Context, traceLog *slog.Logger, request mockFilterWeigherPipelineRequest) (*FilterWeigherPipelineStepResult, error) {
 			return &FilterWeigherPipelineStepResult{
 				Activations: map[string]float64{"host1": 0.1, "host2": 1.0, "host3": 0.0},
 			}, nil
@@ -38,7 +39,7 @@ func TestStepMonitorRun(t *testing.T) {
 		Hosts:   []string{"host1", "host2", "host3"},
 		Weights: map[string]float64{"host1": 0.2, "host2": 0.1, "host3": 0.0},
 	}
-	if _, err := monitor.RunWrapped(slog.Default(), request, step); err != nil {
+	if _, err := monitor.RunWrapped(t.Context(), slog.Default(), request, step); err != nil {
 		t.Fatalf("Run() error = %v, want nil", err)
 	}
 	if len(removedHostsObserver.Observations) != 1 {