Send full requirements to modules that accept them

Author: nolag

pkg/workflows/host/module.go

@@ -23,6 +23,13 @@ type Module interface {
 	Execute(ctx context.Context, request *sdkpb.ExecuteRequest, handler ExecutionHelper) (*sdkpb.ExecutionResult, error)
 }
 
+type RequirementEnforcingModule interface {
+	Module
+
+	// SetRequirements must respect the requirements for the execution until it completes
+	SetRequirements(executionId string, requirements *sdkpb.Requirements)
+}
+
 // ExecutionHelper Implemented by those running the host, for example the Workflow Engine
 type ExecutionHelper interface {
 	// CallCapability blocking call to the Workflow Engine

pkg/workflows/host/requirement_selecting_module.go

@@ -27,64 +27,64 @@ func (l *lazyModule) ensureStarted() {
 	})
 }
 
+// NewRequirementSelectingModule creates a module that routes trigger executions
+// based on subscription requirements. main is prepended as modules[0]; additional
+// modules follow. Subscribe always runs on modules[0].
 func NewRequirementSelectingModule(main ModuleAndHandler, additional []ModuleAndHandler) Module {
-	wrapped := make([]*lazyModule, len(additional))
-	for i := range additional {
-		wrapped[i] = &lazyModule{ModuleAndHandler: additional[i]}
-	}
-	return &requirementSelectingModule{
-		main:       main,
-		additional: wrapped,
+	modules := make([]*lazyModule, 1+len(additional))
+	modules[0] = &lazyModule{ModuleAndHandler: main}
+	for i, a := range additional {
+		modules[1+i] = &lazyModule{ModuleAndHandler: a}
 	}
+	return &requirementSelectingModule{modules: modules}
+}
+
+type triggerInfo struct {
+	moduleIdx    int
+	requirements *sdk.Requirements
 }
 
 type requirementSelectingModule struct {
-	main       ModuleAndHandler
-	additional []*lazyModule
-	// triggerID → index into additional
+	modules []*lazyModule
+	// triggerID → triggerInfo
 	cache sync.Map
 }
 
 func (r *requirementSelectingModule) Start() {
-	r.main.Start()
+	r.modules[0].ensureStarted()
 }
 
 func (r *requirementSelectingModule) Close() {
-	r.main.Close()
-	for _, m := range r.additional {
+	for _, m := range r.modules {
 		if m.started {
 			m.Close()
 		}
 	}
 }
 
 func (r *requirementSelectingModule) IsLegacyDAG() bool {
-	return r.main.IsLegacyDAG()
+	return r.modules[0].IsLegacyDAG()
 }
 
 func (r *requirementSelectingModule) Execute(ctx context.Context, request *sdk.ExecuteRequest, handler ExecutionHelper) (*sdk.ExecutionResult, error) {
-	if triggerID, ok := extractTriggerID(request); ok {
-		if idx, cached := r.cache.Load(triggerID); cached {
-			return r.additional[idx.(int)].Execute(ctx, request, handler)
-		}
-		return r.main.Execute(ctx, request, handler)
+	if request.GetTrigger() == nil {
+		return r.subscribe(ctx, request, handler)
 	}
+	return r.trigger(ctx, request, handler)
+}
 
-	// Subscribe: run main, then build triggerID→module cache from subscription requirements
-	result, err := r.main.Execute(ctx, request, handler)
+func (r *requirementSelectingModule) subscribe(ctx context.Context, request *sdk.ExecuteRequest, handler ExecutionHelper) (*sdk.ExecutionResult, error) {
+	result, err := r.modules[0].Execute(ctx, request, handler)
 	if err != nil {
 		return nil, err
 	}
 
 	for i, sub := range result.GetTriggerSubscriptions().GetSubscriptions() {
-		if sub.Requirements == nil || CheckRequirements(ctx, r.main.RequirementsHandler, sub.Requirements) {
-			continue
-		}
 		matched := false
-		for j, m := range r.additional {
+		for j, m := range r.modules {
 			if CheckRequirements(ctx, m.RequirementsHandler, sub.Requirements) {
 				m.ensureStarted()
-				r.cache.Store(uint64(i), j)
+				r.cache.Store(uint64(i), triggerInfo{moduleIdx: j, requirements: sub.Requirements})
 				matched = true
 				break
 			}
@@ -97,11 +97,18 @@ func (r *requirementSelectingModule) Execute(ctx context.Context, request *sdk.E
 	return result, nil
 }
 
-func extractTriggerID(req *sdk.ExecuteRequest) (uint64, bool) {
-	if t := req.GetTrigger(); t != nil {
-		return t.Id, true
+func (r *requirementSelectingModule) trigger(ctx context.Context, request *sdk.ExecuteRequest, handler ExecutionHelper) (*sdk.ExecutionResult, error) {
+	trigger := request.GetTrigger()
+	if val, cached := r.cache.Load(trigger.Id); cached {
+		info := val.(triggerInfo)
+		m := r.modules[info.moduleIdx]
+		if rem, ok := m.Module.(RequirementEnforcingModule); ok && info.requirements != nil {
+			rem.SetRequirements(handler.GetWorkflowExecutionID(), info.requirements)
+		}
+
+		return m.Execute(ctx, request, handler)
 	}
-	return 0, false
+	return r.modules[0].Execute(ctx, request, handler)
 }
 
 var _ Module = &requirementSelectingModule{}

pkg/workflows/host/requirement_selecting_module_test.go

@@ -27,6 +27,15 @@ func (s *stubModule) Execute(ctx context.Context, req *sdk.ExecuteRequest, h Exe
 	return s.executeFn(ctx, req, h)
 }
 
+type requirementEnforcingStub struct {
+	*stubModule
+	setRequirementsFn func(string, *sdk.Requirements)
+}
+
+func (s *requirementEnforcingStub) SetRequirements(executionID string, requirements *sdk.Requirements) {
+	s.setRequirementsFn(executionID, requirements)
+}
+
 func noop()      {}
 func noopClose() {}
 
@@ -362,6 +371,62 @@ func TestRequirementSelectingModule_Execute(t *testing.T) {
 		assert.Equal(t, want, got)
 		assert.Equal(t, int32(1), atomic.LoadInt32(&mainTriggerCalls), "trigger should run on main")
 	})
+
+	t.Run("cached trigger sets requirements before execute", func(t *testing.T) {
+		teeReqs := &sdk.Requirements{Tee: &sdk.Tee{}}
+		want := &sdk.ExecutionResult{}
+		executionID := "wf-exec-1"
+
+		main := ModuleAndHandler{
+			Module: &stubModule{
+				startFn: noop,
+				executeFn: func(_ context.Context, _ *sdk.ExecuteRequest, _ ExecutionHelper) (*sdk.ExecutionResult, error) {
+					return subscribeResult(subWithReqs(teeReqs)), nil
+				},
+			},
+			RequirementsHandler: RequirementsHandler{Tee: func(context.Context, *sdk.Tee) bool { return false }},
+		}
+
+		var calls []string
+		var gotReqs *sdk.Requirements
+		var gotExecutionID string
+		enforcingAdd := &requirementEnforcingStub{
+			stubModule: &stubModule{
+				startFn: noop,
+				closeFn: noopClose,
+				executeFn: func(context.Context, *sdk.ExecuteRequest, ExecutionHelper) (*sdk.ExecutionResult, error) {
+					calls = append(calls, "execute")
+					return want, nil
+				},
+			},
+			setRequirementsFn: func(id string, requirements *sdk.Requirements) {
+				calls = append(calls, "set")
+				gotExecutionID = id
+				gotReqs = requirements
+			},
+		}
+		add := ModuleAndHandler{
+			Module:              enforcingAdd,
+			RequirementsHandler: RequirementsHandler{Tee: func(context.Context, *sdk.Tee) bool { return true }},
+		}
+
+		m := NewRequirementSelectingModule(main, []ModuleAndHandler{add})
+		m.Start()
+
+		helper := &MockExecutionHelper{}
+		helper.On("GetWorkflowExecutionID").Return(executionID).Once()
+
+		_, err := m.Execute(t.Context(), subscribeRequest(), nil)
+		require.NoError(t, err)
+
+		got, err := m.Execute(t.Context(), triggerRequest(0), helper)
+		require.NoError(t, err)
+		assert.Equal(t, want, got)
+		assert.Equal(t, []string{"set", "execute"}, calls)
+		assert.Equal(t, executionID, gotExecutionID)
+		assert.Same(t, teeReqs, gotReqs)
+		helper.AssertExpectations(t)
+	})
 }
 
 func TestRequirementSelectingModule_TriggerCache(t *testing.T) {

pkg/workflows/host/requirements_gen/requirements_helper.go.tmpl

@@ -14,6 +14,10 @@ type RequirementsHandler struct {
 // non-nil field in req, returning false if any are false, or if the handler is nil.
 // Unknown fields on the proto also result in a false return value.
 func CheckRequirements(ctx context.Context, handler RequirementsHandler, req *sdk.Requirements) bool {
+	if req == nil {
+		return true
+	}
+
 	if len(req.ProtoReflect().GetUnknown()) != 0 {
 		return false
 	}