diff --git a/chasm/context.go b/chasm/context.go index 947ec53fd4..63869017ba 100644 --- a/chasm/context.go +++ b/chasm/context.go @@ -32,7 +32,7 @@ type Context interface { // Value returns the value associated with this context for key. The behavior is the same as context.Context.Value(). // Use WithContextValues RegistrableComponentOption to set key values pair for a component upon registration. // Registered key-value pairs will automatically be added to the Context whenever framework accesses the component. - // Alternatively, use ContextWithValue() to manually set values on Context. + // Alternatively, use ContextWithValue() to manually set values on Context which will take precedence over registered ones. Value(key any) any // Intent() OperationIntent @@ -145,7 +145,11 @@ func (c *immutableCtx) MetricsHandler() metrics.Handler { } func (c *immutableCtx) Value(key any) any { - return c.goContext().Value(key) + if v := c.goContext().Value(key); v != nil { + return v + } + + return c.root.registry.componentContextValue(key) } func (c *immutableCtx) withValue(key any, value any) Context { @@ -198,35 +202,3 @@ func ContextWithValue[C Context](c C, key any, value any) C { //nolint:revive // unchecked-type-assertion return any(c.withValue(key, value)).(C) } - -// AugmentContextForComponent returns a new Context with all context values -// associated with the given component in the registry added. -// This method should only be used by CHASM framework internal code, -// NOT CHASM library developers. -func AugmentContextForComponent[C Context]( - ctx C, - component any, - registry *Registry, -) C { - rc, ok := registry.componentFor(component) - if ok { - for key, value := range rc.contextValues { - ctx = ContextWithValue(ctx, key, value) - } - } - return ctx -} - -func augmentContextForArchetypeID[C Context]( - ctx C, - archetypeID ArchetypeID, - registry *Registry, -) C { - rc, ok := registry.ComponentByID(archetypeID) - if ok { - for key, value := range rc.contextValues { - ctx = ContextWithValue(ctx, key, value) - } - } - return ctx -} diff --git a/chasm/context_mock.go b/chasm/context_mock.go index b33c47b0b1..8ba9de072e 100644 --- a/chasm/context_mock.go +++ b/chasm/context_mock.go @@ -2,6 +2,7 @@ package chasm import ( "context" + "fmt" "slices" "sync" "time" @@ -24,9 +25,27 @@ type MockContext struct { HandleMetricsHandler func() metrics.Handler // GoCtx is the underlying context.Context used for context value lookups. - // Any values set on it will be available via the CHASM mock context's Value method. + // Any values set on it will be available via the CHASM mock context's Value method, + // and take precedence over any registered context values. // Defaults to context.Background() if nil. GoCtx context.Context + + registeredContextValues map[any]any +} + +func (c *MockContext) RegisterComponentContextValues( + keyValues map[any]any, +) { + if c.registeredContextValues == nil { + c.registeredContextValues = make(map[any]any) + } + for k, v := range keyValues { + if _, exists := c.registeredContextValues[k]; exists { + // nolint:forbidigo + panic(fmt.Sprintf("context value key already registered: %v", k)) + } + c.registeredContextValues[k] = v + } } func (c *MockContext) goContext() context.Context { diff --git a/chasm/engine.go b/chasm/engine.go index 40e05fcff6..ea78d9a42f 100644 --- a/chasm/engine.go +++ b/chasm/engine.go @@ -17,34 +17,34 @@ type Engine interface { StartExecution( context.Context, ComponentRef, - func(MutableContext, ArchetypeID, *Registry) (RootComponent, error), + func(MutableContext) (RootComponent, error), ...TransitionOption, ) (StartExecutionResult, error) UpdateWithStartExecution( context.Context, ComponentRef, - func(MutableContext, ArchetypeID, *Registry) (RootComponent, error), - func(MutableContext, Component, *Registry) error, + func(MutableContext) (RootComponent, error), + func(MutableContext, Component) error, ...TransitionOption, ) (EngineUpdateWithStartExecutionResult, error) UpdateComponent( context.Context, ComponentRef, - func(MutableContext, Component, *Registry) error, + func(MutableContext, Component) error, ...TransitionOption, ) ([]byte, error) ReadComponent( context.Context, ComponentRef, - func(Context, Component, *Registry) error, + func(Context, Component) error, ...TransitionOption, ) error PollComponent( context.Context, ComponentRef, - func(Context, Component, *Registry) (bool, error), + func(Context, Component) (bool, error), ...TransitionOption, ) ([]byte, error) @@ -214,12 +214,12 @@ func StartExecution[C RootComponent, I any]( result, err := engineFromContext(ctx).StartExecution( ctx, NewComponentRef[C](key), - func(ctx MutableContext, archetypeID ArchetypeID, registry *Registry) (_ RootComponent, retErr error) { - defer log.CapturePanic(ctx.Logger(), &retErr) + func(mutableContext MutableContext) (_ RootComponent, retErr error) { + defer log.CapturePanic(mutableContext.Logger(), &retErr) var c C var err error - c, err = startFn(augmentContextForArchetypeID(ctx, archetypeID, registry), input) + c, err = startFn(mutableContext, input) return c, err }, opts..., @@ -247,21 +247,21 @@ func UpdateWithStartExecution[C RootComponent, I any, O any]( result, err := engineFromContext(ctx).UpdateWithStartExecution( ctx, NewComponentRef[C](key), - func(ctx MutableContext, archetypeID ArchetypeID, registry *Registry) (_ RootComponent, retErr error) { - defer log.CapturePanic(ctx.Logger(), &retErr) + func(mutableContext MutableContext) (_ RootComponent, retErr error) { + defer log.CapturePanic(mutableContext.Logger(), &retErr) var c C var err error - c, err = startFn(augmentContextForArchetypeID(ctx, archetypeID, registry), input) + c, err = startFn(mutableContext, input) return c, err }, - func(ctx MutableContext, c Component, registry *Registry) (retErr error) { - defer log.CapturePanic(ctx.Logger(), &retErr) + func(mutableContext MutableContext, c Component) (retErr error) { + defer log.CapturePanic(mutableContext.Logger(), &retErr) var err error output, err = updateFn( c.(C), - AugmentContextForComponent(ctx, c, registry), + mutableContext, input, ) return err @@ -306,13 +306,13 @@ func UpdateComponent[C any, R []byte | ComponentRef, I any, O any]( newSerializedRef, err := engineFromContext(ctx).UpdateComponent( ctx, ref, - func(ctx MutableContext, c Component, registry *Registry) (retErr error) { - defer log.CapturePanic(ctx.Logger(), &retErr) + func(mutableContext MutableContext, c Component) (retErr error) { + defer log.CapturePanic(mutableContext.Logger(), &retErr) var err error output, err = updateFn( c.(C), - AugmentContextForComponent(ctx, c, registry), + mutableContext, input, ) return err @@ -345,13 +345,13 @@ func ReadComponent[C any, R []byte | ComponentRef, I any, O any]( err = engineFromContext(ctx).ReadComponent( ctx, ref, - func(ctx Context, c Component, registry *Registry) (retErr error) { - defer log.CapturePanic(ctx.Logger(), &retErr) + func(chasmContext Context, c Component) (retErr error) { + defer log.CapturePanic(chasmContext.Logger(), &retErr) var err error output, err = readFn( c.(C), - AugmentContextForComponent(ctx, c, registry), + chasmContext, input, ) return err @@ -386,12 +386,12 @@ func PollComponent[C any, R []byte | ComponentRef, I any, O any]( newSerializedRef, err := engineFromContext(ctx).PollComponent( ctx, ref, - func(ctx Context, c Component, registry *Registry) (_ bool, retErr error) { - defer log.CapturePanic(ctx.Logger(), &retErr) + func(chasmContext Context, c Component) (_ bool, retErr error) { + defer log.CapturePanic(chasmContext.Logger(), &retErr) out, satisfied, err := monotonicPredicate( c.(C), - AugmentContextForComponent(ctx, c, registry), + chasmContext, input, ) if satisfied { diff --git a/chasm/engine_mock.go b/chasm/engine_mock.go index 11383e34b5..a3d86b11ba 100644 --- a/chasm/engine_mock.go +++ b/chasm/engine_mock.go @@ -67,7 +67,7 @@ func (mr *MockEngineMockRecorder) NotifyExecution(arg0 any) *gomock.Call { } // PollComponent mocks base method. -func (m *MockEngine) PollComponent(arg0 context.Context, arg1 ComponentRef, arg2 func(Context, Component, *Registry) (bool, error), arg3 ...TransitionOption) ([]byte, error) { +func (m *MockEngine) PollComponent(arg0 context.Context, arg1 ComponentRef, arg2 func(Context, Component) (bool, error), arg3 ...TransitionOption) ([]byte, error) { m.ctrl.T.Helper() varargs := []any{arg0, arg1, arg2} for _, a := range arg3 { @@ -87,7 +87,7 @@ func (mr *MockEngineMockRecorder) PollComponent(arg0, arg1, arg2 any, arg3 ...an } // ReadComponent mocks base method. -func (m *MockEngine) ReadComponent(arg0 context.Context, arg1 ComponentRef, arg2 func(Context, Component, *Registry) error, arg3 ...TransitionOption) error { +func (m *MockEngine) ReadComponent(arg0 context.Context, arg1 ComponentRef, arg2 func(Context, Component) error, arg3 ...TransitionOption) error { m.ctrl.T.Helper() varargs := []any{arg0, arg1, arg2} for _, a := range arg3 { @@ -106,7 +106,7 @@ func (mr *MockEngineMockRecorder) ReadComponent(arg0, arg1, arg2 any, arg3 ...an } // StartExecution mocks base method. -func (m *MockEngine) StartExecution(arg0 context.Context, arg1 ComponentRef, arg2 func(MutableContext, ArchetypeID, *Registry) (RootComponent, error), arg3 ...TransitionOption) (StartExecutionResult, error) { +func (m *MockEngine) StartExecution(arg0 context.Context, arg1 ComponentRef, arg2 func(MutableContext) (RootComponent, error), arg3 ...TransitionOption) (StartExecutionResult, error) { m.ctrl.T.Helper() varargs := []any{arg0, arg1, arg2} for _, a := range arg3 { @@ -126,7 +126,7 @@ func (mr *MockEngineMockRecorder) StartExecution(arg0, arg1, arg2 any, arg3 ...a } // UpdateComponent mocks base method. -func (m *MockEngine) UpdateComponent(arg0 context.Context, arg1 ComponentRef, arg2 func(MutableContext, Component, *Registry) error, arg3 ...TransitionOption) ([]byte, error) { +func (m *MockEngine) UpdateComponent(arg0 context.Context, arg1 ComponentRef, arg2 func(MutableContext, Component) error, arg3 ...TransitionOption) ([]byte, error) { m.ctrl.T.Helper() varargs := []any{arg0, arg1, arg2} for _, a := range arg3 { @@ -146,7 +146,7 @@ func (mr *MockEngineMockRecorder) UpdateComponent(arg0, arg1, arg2 any, arg3 ... } // UpdateWithStartExecution mocks base method. -func (m *MockEngine) UpdateWithStartExecution(arg0 context.Context, arg1 ComponentRef, arg2 func(MutableContext, ArchetypeID, *Registry) (RootComponent, error), arg3 func(MutableContext, Component, *Registry) error, arg4 ...TransitionOption) (EngineUpdateWithStartExecutionResult, error) { +func (m *MockEngine) UpdateWithStartExecution(arg0 context.Context, arg1 ComponentRef, arg2 func(MutableContext) (RootComponent, error), arg3 func(MutableContext, Component) error, arg4 ...TransitionOption) (EngineUpdateWithStartExecutionResult, error) { m.ctrl.T.Helper() varargs := []any{arg0, arg1, arg2, arg3} for _, a := range arg4 { diff --git a/chasm/lib/callback/executors_test.go b/chasm/lib/callback/executors_test.go index 3a6da21aa7..9aa44fe9a1 100644 --- a/chasm/lib/callback/executors_test.go +++ b/chasm/lib/callback/executors_test.go @@ -229,7 +229,7 @@ func TestExecuteInvocationTaskNexus_Outcomes(t *testing.T) { gomock.Any(), gomock.Any(), gomock.Any(), - ).DoAndReturn(func(ctx context.Context, ref chasm.ComponentRef, readFn func(chasm.Context, chasm.Component, *chasm.Registry) error, opts ...chasm.TransitionOption) error { + ).DoAndReturn(func(ctx context.Context, ref chasm.ComponentRef, readFn func(chasm.Context, chasm.Component) error, opts ...chasm.TransitionOption) error { mockCtx := &chasm.MockContext{ HandleNow: func(component chasm.Component) time.Time { return timeSource.Now() @@ -238,14 +238,14 @@ func TestExecuteInvocationTaskNexus_Outcomes(t *testing.T) { return []byte{}, nil }, } - return readFn(mockCtx, callback, chasmRegistry) + return readFn(mockCtx, callback) }) mockEngine.EXPECT().UpdateComponent( gomock.Any(), gomock.Any(), gomock.Any(), - ).DoAndReturn(func(ctx context.Context, ref chasm.ComponentRef, updateFn func(chasm.MutableContext, chasm.Component, *chasm.Registry) error, opts ...chasm.TransitionOption) ([]any, error) { + ).DoAndReturn(func(ctx context.Context, ref chasm.ComponentRef, updateFn func(chasm.MutableContext, chasm.Component) error, opts ...chasm.TransitionOption) ([]any, error) { mockCtx := &chasm.MockMutableContext{ MockContext: chasm.MockContext{ HandleNow: func(component chasm.Component) time.Time { @@ -256,7 +256,7 @@ func TestExecuteInvocationTaskNexus_Outcomes(t *testing.T) { }, }, } - err := updateFn(mockCtx, callback, chasmRegistry) + err := updateFn(mockCtx, callback) return nil, err }) @@ -613,7 +613,7 @@ func TestExecuteInvocationTaskChasm_Outcomes(t *testing.T) { gomock.Any(), gomock.Any(), gomock.Any(), - ).DoAndReturn(func(ctx context.Context, ref chasm.ComponentRef, readFn func(chasm.Context, chasm.Component, *chasm.Registry) error, opts ...chasm.TransitionOption) error { + ).DoAndReturn(func(ctx context.Context, ref chasm.ComponentRef, readFn func(chasm.Context, chasm.Component) error, opts ...chasm.TransitionOption) error { // Create a mock context mockCtx := &chasm.MockContext{ HandleNow: func(component chasm.Component) time.Time { @@ -632,14 +632,14 @@ func TestExecuteInvocationTaskChasm_Outcomes(t *testing.T) { } // Call the readFn with our callback - return readFn(mockCtx, callback, chasmRegistry) + return readFn(mockCtx, callback) }) mockEngine.EXPECT().UpdateComponent( gomock.Any(), gomock.Any(), gomock.Any(), - ).DoAndReturn(func(ctx context.Context, ref chasm.ComponentRef, updateFn func(chasm.MutableContext, chasm.Component, *chasm.Registry) error, opts ...chasm.TransitionOption) ([]any, error) { + ).DoAndReturn(func(ctx context.Context, ref chasm.ComponentRef, updateFn func(chasm.MutableContext, chasm.Component) error, opts ...chasm.TransitionOption) ([]any, error) { // Create a mock mutable context mockCtx := &chasm.MockMutableContext{ MockContext: chasm.MockContext{ @@ -653,7 +653,7 @@ func TestExecuteInvocationTaskChasm_Outcomes(t *testing.T) { } // Call the updateFn with our callback - err := updateFn(mockCtx, callback, chasmRegistry) + err := updateFn(mockCtx, callback) return nil, err }) diff --git a/chasm/lib/scheduler/helper_test.go b/chasm/lib/scheduler/helper_test.go index 772ee15004..57818adeb2 100644 --- a/chasm/lib/scheduler/helper_test.go +++ b/chasm/lib/scheduler/helper_test.go @@ -299,8 +299,8 @@ func (e *testEnv) ExpectReadComponent(ctx chasm.Context, returnedComponent chasm e.t.Fatal("ExpectReadComponent requires withMockEngine() option") } e.MockEngine.EXPECT().ReadComponent(gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any()). - DoAndReturn(func(_ context.Context, _ chasm.ComponentRef, readFn func(chasm.Context, chasm.Component, *chasm.Registry) error, _ ...chasm.TransitionOption) error { - return readFn(ctx, returnedComponent, e.Registry) + DoAndReturn(func(_ context.Context, _ chasm.ComponentRef, readFn func(chasm.Context, chasm.Component) error, _ ...chasm.TransitionOption) error { + return readFn(ctx, returnedComponent) }).Times(1) } @@ -310,8 +310,8 @@ func (e *testEnv) ExpectUpdateComponent(ctx chasm.MutableContext, componentToUpd e.t.Fatal("ExpectUpdateComponent requires withMockEngine() option") } e.MockEngine.EXPECT().UpdateComponent(gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any()). - DoAndReturn(func(_ context.Context, _ chasm.ComponentRef, updateFn func(chasm.MutableContext, chasm.Component, *chasm.Registry) error, _ ...chasm.TransitionOption) ([]byte, error) { - err := updateFn(ctx, componentToUpdate, e.Registry) + DoAndReturn(func(_ context.Context, _ chasm.ComponentRef, updateFn func(chasm.MutableContext, chasm.Component) error, _ ...chasm.TransitionOption) ([]byte, error) { + err := updateFn(ctx, componentToUpdate) return nil, err }).Times(1) } diff --git a/chasm/registrable_component.go b/chasm/registrable_component.go index 24e59dfeaa..15b7f659e4 100644 --- a/chasm/registrable_component.go +++ b/chasm/registrable_component.go @@ -162,6 +162,10 @@ func WithSearchAttributes( // // This is useful for propagating values needed for those processing logic but are not avaiable via the // component's struct definition, such as configurations. +// +// Keys need to be globally unique across components. Conflicting keys across will cause component registration to fail. +// +// Manually added key-value pairs via ContextWithValue() will take precedence over registered context values. func WithContextValues( keyVals map[any]any, ) RegistrableComponentOption { diff --git a/chasm/registrable_task.go b/chasm/registrable_task.go index 846d3ab491..9a6f5e09ff 100644 --- a/chasm/registrable_task.go +++ b/chasm/registrable_task.go @@ -58,7 +58,7 @@ func NewRegistrableSideEffectTask[C any, T any]( registry *Registry, ) (bool, error) { return validator.Validate( - AugmentContextForComponent(ctx, component, registry), + ctx, component.(C), taskAttrs, taskData.(T), @@ -97,7 +97,7 @@ func NewRegistrablePureTask[C any, T any]( registry *Registry, ) (bool, error) { return validator.Validate( - AugmentContextForComponent(ctx, component, registry), + ctx, component.(C), taskAttrs, taskData.(T), @@ -111,7 +111,7 @@ func NewRegistrablePureTask[C any, T any]( registry *Registry, ) error { return executor.Execute( - AugmentContextForComponent(ctx, component, registry), + ctx, component.(C), taskAttrs, taskData.(T), diff --git a/chasm/registry.go b/chasm/registry.go index 4735e38766..910cdcf499 100644 --- a/chasm/registry.go +++ b/chasm/registry.go @@ -26,6 +26,10 @@ type ( rcByFqn map[string]*RegistrableComponent // fully qualified type name -> component rcByID map[uint32]*RegistrableComponent // component type ID -> component rcByGoType map[reflect.Type]*RegistrableComponent // component go type -> component + // rcContextValues is aggregated context values from all components, + // used for easy lookup when Context.Value(key) is called. + // Registration process will check for key conflicts and return error if same key is registered by multiple components. + rcContextValues map[any]valueWithFqn // rt stands for RegistrableTask. rtByFqn map[string]*RegistrableTask // fully qualified type name -> task @@ -39,6 +43,13 @@ type ( } ) +// valueWithFqn is a wrapper struct that associates a value with +// the fully qualified name (FQN) of the component that registered it. +type valueWithFqn struct { + v any + fqn string +} + func NewRegistry(logger log.Logger) *Registry { return &Registry{ libraries: make(map[string]Library), @@ -48,6 +59,7 @@ func NewRegistry(logger log.Logger) *Registry { rtByFqn: make(map[string]*RegistrableTask), rtByID: make(map[uint32]*RegistrableTask), rtByGoType: make(map[reflect.Type]*RegistrableTask), + rcContextValues: make(map[any]valueWithFqn), nexusServices: make(map[string]*nexus.Service), NexusEndpointProcessor: NewNexusEndpointProcessor(), logger: logger, @@ -244,6 +256,16 @@ func (r *Registry) registerComponent( return fmt.Errorf("component ID %d collision between %s and %s", id, fqn, existingComponent.fqType()) } + for key, value := range rc.contextValues { + if existingValue, ok := r.rcContextValues[key]; ok { + return fmt.Errorf("context value key %v registered by component %s conflicts with component %s", key, fqn, existingValue.fqn) + } + r.rcContextValues[key] = valueWithFqn{ + v: value, + fqn: fqn, + } + } + // rc.goType implements Component interface; therefore, it must be a struct. // This check to protect against the interface itself being registered. if !(rc.goType.Kind() == reflect.Struct || @@ -358,3 +380,10 @@ func (r *Registry) NexusServices() map[string]*nexus.Service { maps.Copy(services, r.nexusServices) return services } + +func (r *Registry) componentContextValue(key any) any { + if v, ok := r.rcContextValues[key]; ok { + return v.v + } + return nil +} diff --git a/chasm/tree.go b/chasm/tree.go index f981a67ed3..433beb0946 100644 --- a/chasm/tree.go +++ b/chasm/tree.go @@ -327,12 +327,7 @@ func newTreeInitSearchAttributesAndMemo( root *Node, registry *Registry, ) error { - immutableContext := augmentContextForArchetypeID( - NewContext(context.Background(), root), - root.ArchetypeID(), - registry, - ) - + immutableContext := NewContext(context.Background(), root) rootComponent, err := root.Component(immutableContext, ComponentRef{}) if err != nil { return err @@ -496,7 +491,7 @@ func (n *Node) validateAccessHelper(ctx Context) error { } componentValue, _ := n.value.(Component) //nolint:revive // unchecked-type-assertion - if componentValue.LifecycleState(AugmentContextForComponent(ctx, componentValue, n.registry)).IsClosed() { + if componentValue.LifecycleState(ctx).IsClosed() { return errAccessCheckFailed } @@ -1446,12 +1441,7 @@ func (n *Node) CloseTransaction() (NodesMutation, error) { TransitionCount: n.backend.NextTransitionCount(), } - immutableContext := augmentContextForArchetypeID( - NewContext(context.TODO(), n), - n.ArchetypeID(), - n.registry, - ) - + immutableContext := NewContext(context.TODO(), n) rootLifecycleChanged, err := n.closeTransactionHandleRootLifecycleChange(immutableContext) if err != nil { return NodesMutation{}, err @@ -2233,11 +2223,7 @@ func (n *Node) ApplyMutation( // // TODO: combine this with the logic in CloseTransactionForceUpdateVisibility // right that force update logic only applies to the active cluster. - immutableContext := augmentContextForArchetypeID( - NewContext(context.Background(), n), - n.ArchetypeID(), - n.registry, - ) + immutableContext := NewContext(context.TODO(), n) rootComponent, err := n.root().Component(immutableContext, ComponentRef{}) if err != nil { return err @@ -2562,16 +2548,11 @@ func (n *Node) Terminate( ) } - mutableContext := augmentContextForArchetypeID( - NewMutableContext(context.Background(), n.root()), - n.ArchetypeID(), - n.registry, - ) + mutableContext := NewMutableContext(context.TODO(), n.root()) component, err := n.Component(mutableContext, ComponentRef{}) if err != nil { return err } - rootComponent, ok := component.(RootComponent) if !ok { return softassert.UnexpectedInternalErr( diff --git a/service/history/chasm_engine.go b/service/history/chasm_engine.go index 09793ef73b..77cac2d219 100644 --- a/service/history/chasm_engine.go +++ b/service/history/chasm_engine.go @@ -111,7 +111,7 @@ func (e *ChasmEngine) NotifyExecution(key chasm.ExecutionKey) { func (e *ChasmEngine) StartExecution( ctx context.Context, executionRef chasm.ComponentRef, - startFn func(chasm.MutableContext, chasm.ArchetypeID, *chasm.Registry) (chasm.RootComponent, error), + startFn func(chasm.MutableContext) (chasm.RootComponent, error), opts ...chasm.TransitionOption, ) (chasm.StartExecutionResult, error) { options := e.constructTransitionOptions(opts...) @@ -122,7 +122,7 @@ func (e *ChasmEngine) StartExecution( func (e *ChasmEngine) startExecution( ctx context.Context, executionRef chasm.ComponentRef, - startFn func(chasm.MutableContext, chasm.ArchetypeID, *chasm.Registry) (chasm.RootComponent, error), + startFn func(chasm.MutableContext) (chasm.RootComponent, error), options chasm.TransitionOptions, ) (result chasm.StartExecutionResult, retErr error) { shardContext, err := e.getShardContext(executionRef) @@ -204,8 +204,8 @@ func (e *ChasmEngine) startExecution( func (e *ChasmEngine) UpdateWithStartExecution( ctx context.Context, executionRef chasm.ComponentRef, - startFn func(chasm.MutableContext, chasm.ArchetypeID, *chasm.Registry) (chasm.RootComponent, error), - updateFn func(chasm.MutableContext, chasm.Component, *chasm.Registry) error, + startFn func(chasm.MutableContext) (chasm.RootComponent, error), + updateFn func(chasm.MutableContext, chasm.Component) error, opts ...chasm.TransitionOption, ) (chasm.EngineUpdateWithStartExecutionResult, error) { options := e.constructTransitionOptions(opts...) @@ -216,8 +216,8 @@ func (e *ChasmEngine) UpdateWithStartExecution( func (e *ChasmEngine) updateWithStartExecution( ctx context.Context, executionRef chasm.ComponentRef, - startFn func(chasm.MutableContext, chasm.ArchetypeID, *chasm.Registry) (chasm.RootComponent, error), - updateFn func(chasm.MutableContext, chasm.Component, *chasm.Registry) error, + startFn func(chasm.MutableContext) (chasm.RootComponent, error), + updateFn func(chasm.MutableContext, chasm.Component) error, options chasm.TransitionOptions, ) (result chasm.EngineUpdateWithStartExecutionResult, retError error) { shardContext, err := e.getShardContext(executionRef) @@ -299,7 +299,7 @@ func (e *ChasmEngine) updateExecution( shardContext historyi.ShardContext, executionLease api.WorkflowLease, executionRef chasm.ComponentRef, - updateFn func(chasm.MutableContext, chasm.Component, *chasm.Registry) error, + updateFn func(chasm.MutableContext, chasm.Component) error, ) (chasm.ExecutionKey, []byte, error) { workflowKey := executionLease.GetContext().GetWorkflowKey() actualRef := executionRef @@ -321,8 +321,8 @@ func (e *ChasmEngine) startNewForClosedExecution( executionLease api.WorkflowLease, executionRef chasm.ComponentRef, archetypeID chasm.ArchetypeID, - startFn func(chasm.MutableContext, chasm.ArchetypeID, *chasm.Registry) (chasm.RootComponent, error), - updateFn func(chasm.MutableContext, chasm.Component, *chasm.Registry) error, + startFn func(chasm.MutableContext) (chasm.RootComponent, error), + updateFn func(chasm.MutableContext, chasm.Component) error, options chasm.TransitionOptions, ) (chasm.ExecutionKey, []byte, bool, error) { newExecutionParams, err := e.createNewExecutionWithUpdate( @@ -351,7 +351,7 @@ func (e *ChasmEngine) applyUpdateWithLease( shardContext historyi.ShardContext, executionLease api.WorkflowLease, ref chasm.ComponentRef, - updateFn func(chasm.MutableContext, chasm.Component, *chasm.Registry) error, + updateFn func(chasm.MutableContext, chasm.Component) error, ) ([]byte, error) { mutableState := executionLease.GetMutableState() chasmTree, ok := mutableState.ChasmTree().(*chasm.Node) @@ -369,7 +369,7 @@ func (e *ChasmEngine) applyUpdateWithLease( return nil, err } - if err := updateFn(mutableContext, component, e.registry); err != nil { + if err := updateFn(mutableContext, component); err != nil { return nil, err } @@ -395,8 +395,8 @@ func (e *ChasmEngine) startAndUpdateExecution( shardContext historyi.ShardContext, executionRef chasm.ComponentRef, archetypeID chasm.ArchetypeID, - startFn func(chasm.MutableContext, chasm.ArchetypeID, *chasm.Registry) (chasm.RootComponent, error), - updateFn func(chasm.MutableContext, chasm.Component, *chasm.Registry) error, + startFn func(chasm.MutableContext) (chasm.RootComponent, error), + updateFn func(chasm.MutableContext, chasm.Component) error, options chasm.TransitionOptions, ) (retKey chasm.ExecutionKey, retRef []byte, created bool, retErr error) { currentExecutionReleaseFn, err := e.lockCurrentExecution( @@ -450,7 +450,7 @@ func (e *ChasmEngine) startAndUpdateExecution( func (e *ChasmEngine) UpdateComponent( ctx context.Context, ref chasm.ComponentRef, - updateFn func(chasm.MutableContext, chasm.Component, *chasm.Registry) error, + updateFn func(chasm.MutableContext, chasm.Component) error, opts ...chasm.TransitionOption, ) ([]byte, error) { options := e.constructTransitionOptions(opts...) @@ -461,7 +461,7 @@ func (e *ChasmEngine) UpdateComponent( func (e *ChasmEngine) updateComponent( ctx context.Context, ref chasm.ComponentRef, - updateFn func(chasm.MutableContext, chasm.Component, *chasm.Registry) error, + updateFn func(chasm.MutableContext, chasm.Component) error, ) (updatedRef []byte, retError error) { shardContext, executionLease, err := e.getExecutionLease(ctx, ref) if err != nil { @@ -558,7 +558,7 @@ func (e *ChasmEngine) deleteExecution( func (e *ChasmEngine) ReadComponent( ctx context.Context, ref chasm.ComponentRef, - readFn func(chasm.Context, chasm.Component, *chasm.Registry) error, + readFn func(chasm.Context, chasm.Component) error, opts ...chasm.TransitionOption, ) error { options := e.constructTransitionOptions(opts...) @@ -568,7 +568,7 @@ func (e *ChasmEngine) ReadComponent( func (e *ChasmEngine) readComponent( ctx context.Context, ref chasm.ComponentRef, - readFn func(chasm.Context, chasm.Component, *chasm.Registry) error, + readFn func(chasm.Context, chasm.Component) error, ) error { _, executionLease, err := e.getExecutionLease(ctx, ref) if err != nil { @@ -595,7 +595,7 @@ func (e *ChasmEngine) readComponent( return err } - return readFn(chasmContext, component, e.registry) + return readFn(chasmContext, component) } // PollComponent waits until the supplied predicate is satisfied when evaluated against the @@ -614,7 +614,7 @@ func (e *ChasmEngine) readComponent( func (e *ChasmEngine) PollComponent( ctx context.Context, requestRef chasm.ComponentRef, - monotonicPredicate func(chasm.Context, chasm.Component, *chasm.Registry) (bool, error), + monotonicPredicate func(chasm.Context, chasm.Component) (bool, error), opts ...chasm.TransitionOption, ) ([]byte, error) { options := e.constructTransitionOptions(opts...) @@ -625,7 +625,7 @@ func (e *ChasmEngine) PollComponent( func (e *ChasmEngine) pollComponent( ctx context.Context, requestRef chasm.ComponentRef, - monotonicPredicate func(chasm.Context, chasm.Component, *chasm.Registry) (bool, error), + monotonicPredicate func(chasm.Context, chasm.Component) (bool, error), ) (retRef []byte, retError error) { var ch <-chan struct{} @@ -682,7 +682,7 @@ func (e *ChasmEngine) pollComponent( // iff there's no error and predicate evaluates to true. func (e *ChasmEngine) predicateSatisfied( ctx context.Context, - predicate func(chasm.Context, chasm.Component, *chasm.Registry) (bool, error), + predicate func(chasm.Context, chasm.Component) (bool, error), ref chasm.ComponentRef, executionLease api.WorkflowLease, ) ([]byte, error) { @@ -700,7 +700,7 @@ func (e *ChasmEngine) predicateSatisfied( if err != nil { return nil, err } - satisfied, err := predicate(chasmContext, component, e.registry) + satisfied, err := predicate(chasmContext, component) if err != nil { return nil, err } @@ -750,7 +750,7 @@ func (e *ChasmEngine) createNewExecution( shardContext historyi.ShardContext, executionRef chasm.ComponentRef, archetypeID chasm.ArchetypeID, - startFn func(chasm.MutableContext, chasm.ArchetypeID, *chasm.Registry) (chasm.RootComponent, error), + startFn func(chasm.MutableContext) (chasm.RootComponent, error), options chasm.TransitionOptions, ) (newExecutionParams, error) { return e.createNewExecutionWithUpdate( @@ -769,8 +769,8 @@ func (e *ChasmEngine) createNewExecutionWithUpdate( shardContext historyi.ShardContext, executionRef chasm.ComponentRef, archetypeID chasm.ArchetypeID, - startFn func(chasm.MutableContext, chasm.ArchetypeID, *chasm.Registry) (chasm.RootComponent, error), - updateFn func(chasm.MutableContext, chasm.Component, *chasm.Registry) error, + startFn func(chasm.MutableContext) (chasm.RootComponent, error), + updateFn func(chasm.MutableContext, chasm.Component) error, options chasm.TransitionOptions, ) (newExecutionParams, error) { executionRef.RunID = primitives.NewUUID().String() @@ -804,7 +804,7 @@ func (e *ChasmEngine) createNewExecutionWithUpdate( chasmContext := chasm.NewMutableContext(ctx, chasmTree) - rootComponent, err := startFn(chasmContext, archetypeID, e.registry) + rootComponent, err := startFn(chasmContext) if err != nil { return newExecutionParams{}, err } @@ -813,7 +813,7 @@ func (e *ChasmEngine) createNewExecutionWithUpdate( } if updateFn != nil { - if err = updateFn(chasmContext, rootComponent, e.registry); err != nil { + if err = updateFn(chasmContext, rootComponent); err != nil { return newExecutionParams{}, err } } diff --git a/service/history/chasm_engine_test.go b/service/history/chasm_engine_test.go index e9bd4ab81b..8f44c6d1e5 100644 --- a/service/history/chasm_engine_test.go +++ b/service/history/chasm_engine_test.go @@ -493,8 +493,8 @@ func (s *chasmEngineSuite) TestNewExecution_ConflictPolicy_TerminateExisting() { func (s *chasmEngineSuite) newTestExecutionFn( activityID string, -) func(chasm.MutableContext, chasm.ArchetypeID, *chasm.Registry) (chasm.RootComponent, error) { - return func(ctx chasm.MutableContext, _ chasm.ArchetypeID, _ *chasm.Registry) (chasm.RootComponent, error) { +) func(chasm.MutableContext) (chasm.RootComponent, error) { + return func(ctx chasm.MutableContext) (chasm.RootComponent, error) { return &testComponent{ ActivityInfo: &persistencespb.ActivityInfo{ ActivityId: activityID, @@ -715,7 +715,6 @@ func (s *chasmEngineSuite) TestUpdateComponent_Success() { func( ctx chasm.MutableContext, component chasm.Component, - _ *chasm.Registry, ) error { tc, ok := component.(*testComponent) s.True(ok) @@ -752,7 +751,6 @@ func (s *chasmEngineSuite) TestReadComponent_Success() { func( ctx chasm.Context, component chasm.Component, - _ *chasm.Registry, ) error { tc, ok := component.(*testComponent) s.True(ok) @@ -791,7 +789,7 @@ func (s *chasmEngineSuite) TestPollComponent_Success_NoWait() { newSerializedRef, err := s.engine.PollComponent( context.Background(), ref, - func(ctx chasm.Context, component chasm.Component, _ *chasm.Registry) (bool, error) { + func(ctx chasm.Context, component chasm.Component) (bool, error) { return true, nil }, ) @@ -889,7 +887,7 @@ func (s *chasmEngineSuite) testPollComponentWait(useEmptyRunID bool) { newSerializedRef, err := s.engine.PollComponent( ctx, pollRef, - func(ctx chasm.Context, component chasm.Component, _ *chasm.Registry) (bool, error) { + func(ctx chasm.Context, component chasm.Component) (bool, error) { tc, ok := component.(*testComponent) s.True(ok) satisfied := tc.ActivityInfo.ActivityId == activityID @@ -903,7 +901,7 @@ func (s *chasmEngineSuite) testPollComponentWait(useEmptyRunID bool) { _, err := s.engine.UpdateComponent( context.Background(), updateRef, - func(ctx chasm.MutableContext, component chasm.Component, _ *chasm.Registry) error { + func(ctx chasm.MutableContext, component chasm.Component) error { tc, ok := component.(*testComponent) s.True(ok) if satisfyPredicate { @@ -950,7 +948,6 @@ func (s *chasmEngineSuite) testPollComponentWait(useEmptyRunID bool) { func( ctx chasm.Context, component chasm.Component, - _ *chasm.Registry, ) error { tc, ok := component.(*testComponent) s.True(ok) @@ -1006,7 +1003,7 @@ func (s *chasmEngineSuite) TestPollComponent_StaleState() { _, err = s.engine.PollComponent( context.Background(), staleRef, - func(ctx chasm.Context, component chasm.Component, _ *chasm.Registry) (bool, error) { + func(ctx chasm.Context, component chasm.Component) (bool, error) { s.Fail("predicate should not be called with stale ref") return false, nil }, @@ -1050,7 +1047,6 @@ func (s *chasmEngineSuite) TestCloseTime_ReturnsNonZeroWhenCompleted() { func( ctx chasm.Context, component chasm.Component, - _ *chasm.Registry, ) error { // Verify CloseTime returns non-zero time when component is completed closeTime := ctx.ExecutionCloseTime() @@ -1093,7 +1089,7 @@ func (s *chasmEngineSuite) TestStateTransitionCount() { _, err := s.engine.UpdateComponent( context.Background(), ref, - func(ctx chasm.MutableContext, component chasm.Component, _ *chasm.Registry) error { + func(ctx chasm.MutableContext, component chasm.Component) error { tc, ok := component.(*testComponent) s.True(ok) tc.ActivityInfo.ActivityId = tv.ActivityID() @@ -1105,7 +1101,7 @@ func (s *chasmEngineSuite) TestStateTransitionCount() { err = s.engine.ReadComponent( context.Background(), ref, - func(ctx chasm.Context, component chasm.Component, _ *chasm.Registry) error { + func(ctx chasm.Context, component chasm.Component) error { s.Equal(initialCount+1, ctx.StateTransitionCount()) return nil }, @@ -1174,11 +1170,11 @@ func (s *chasmEngineSuite) TestUpdateWithStartExecution_ExistingRunning() { result, err := s.engine.UpdateWithStartExecution( context.Background(), chasm.NewComponentRef[*testComponent](executionKey), - func(ctx chasm.MutableContext, _ chasm.ArchetypeID, _ *chasm.Registry) (chasm.RootComponent, error) { + func(ctx chasm.MutableContext) (chasm.RootComponent, error) { s.Fail("newFn should not be called when execution exists and is running") return nil, nil }, - func(ctx chasm.MutableContext, component chasm.Component, _ *chasm.Registry) error { + func(ctx chasm.MutableContext, component chasm.Component) error { tc, ok := component.(*testComponent) s.True(ok) tc.ActivityInfo.ActivityId = "updated-" + tc.ActivityInfo.ActivityId @@ -1241,7 +1237,7 @@ func (s *chasmEngineSuite) TestUpdateWithStartExecution_NotFound() { result, err := s.engine.UpdateWithStartExecution( context.Background(), chasm.NewComponentRef[*testComponent](executionKey), - func(ctx chasm.MutableContext, _ chasm.ArchetypeID, _ *chasm.Registry) (chasm.RootComponent, error) { + func(ctx chasm.MutableContext) (chasm.RootComponent, error) { newFnCalled = true return &testComponent{ ActivityInfo: &persistencespb.ActivityInfo{ @@ -1249,7 +1245,7 @@ func (s *chasmEngineSuite) TestUpdateWithStartExecution_NotFound() { }, }, nil }, - func(ctx chasm.MutableContext, component chasm.Component, _ *chasm.Registry) error { + func(ctx chasm.MutableContext, component chasm.Component) error { updateFnCalled = true tc, ok := component.(*testComponent) s.True(ok) @@ -1331,7 +1327,7 @@ func (s *chasmEngineSuite) TestUpdateWithStartExecution_ExistingClosed() { result, err := s.engine.UpdateWithStartExecution( context.Background(), chasm.NewComponentRef[*testComponent](executionKey), - func(ctx chasm.MutableContext, _ chasm.ArchetypeID, _ *chasm.Registry) (chasm.RootComponent, error) { + func(ctx chasm.MutableContext) (chasm.RootComponent, error) { newFnCalled = true return &testComponent{ ActivityInfo: &persistencespb.ActivityInfo{ @@ -1339,7 +1335,7 @@ func (s *chasmEngineSuite) TestUpdateWithStartExecution_ExistingClosed() { }, }, nil }, - func(ctx chasm.MutableContext, component chasm.Component, _ *chasm.Registry) error { + func(ctx chasm.MutableContext, component chasm.Component) error { updateFnCalled = true // Apply the "update" to the newly created component tc, ok := component.(*testComponent) @@ -1399,11 +1395,11 @@ func (s *chasmEngineSuite) TestUpdateWithStartExecution_UpdateFnError() { _, err := s.engine.UpdateWithStartExecution( context.Background(), chasm.NewComponentRef[*testComponent](executionKey), - func(ctx chasm.MutableContext, _ chasm.ArchetypeID, _ *chasm.Registry) (chasm.RootComponent, error) { + func(ctx chasm.MutableContext) (chasm.RootComponent, error) { s.Fail("newFn should not be called") return nil, nil }, - func(ctx chasm.MutableContext, component chasm.Component, _ *chasm.Registry) error { + func(ctx chasm.MutableContext, component chasm.Component) error { return expectedErr }, ) @@ -1426,10 +1422,10 @@ func (s *chasmEngineSuite) TestUpdateWithStartExecution_NewFnError() { _, err := s.engine.UpdateWithStartExecution( context.Background(), chasm.NewComponentRef[*testComponent](executionKey), - func(ctx chasm.MutableContext, _ chasm.ArchetypeID, _ *chasm.Registry) (chasm.RootComponent, error) { + func(ctx chasm.MutableContext) (chasm.RootComponent, error) { return nil, expectedErr }, - func(ctx chasm.MutableContext, component chasm.Component, _ *chasm.Registry) error { + func(ctx chasm.MutableContext, component chasm.Component) error { s.Fail("updateFn should not be called when newFn fails") return nil }, @@ -1454,7 +1450,7 @@ func (s *chasmEngineSuite) TestUpdateWithStartExecution_UpdateFnErrorOnCreate() _, err := s.engine.UpdateWithStartExecution( context.Background(), chasm.NewComponentRef[*testComponent](executionKey), - func(ctx chasm.MutableContext, _ chasm.ArchetypeID, _ *chasm.Registry) (chasm.RootComponent, error) { + func(ctx chasm.MutableContext) (chasm.RootComponent, error) { newFnCalled = true return &testComponent{ ActivityInfo: &persistencespb.ActivityInfo{ @@ -1462,7 +1458,7 @@ func (s *chasmEngineSuite) TestUpdateWithStartExecution_UpdateFnErrorOnCreate() }, }, nil }, - func(ctx chasm.MutableContext, component chasm.Component, _ *chasm.Registry) error { + func(ctx chasm.MutableContext, component chasm.Component) error { return expectedErr }, ) @@ -1520,11 +1516,11 @@ func (s *chasmEngineSuite) TestUpdateWithStartExecution_UpdatePathVersionConflic _, err := s.engine.UpdateWithStartExecution( context.Background(), chasm.NewComponentRef[*testComponent](executionKey), - func(ctx chasm.MutableContext, _ chasm.ArchetypeID, _ *chasm.Registry) (chasm.RootComponent, error) { + func(ctx chasm.MutableContext) (chasm.RootComponent, error) { s.Fail("newFn should not be called when execution exists") return nil, nil }, - func(ctx chasm.MutableContext, component chasm.Component, _ *chasm.Registry) error { + func(ctx chasm.MutableContext, component chasm.Component) error { // updateFn is called before the version conflict is detected during persist. updateFnCalled = true tc, ok := component.(*testComponent) @@ -1553,7 +1549,7 @@ func (s *chasmEngineSuite) TestReadComponent_NotFound() { RunID: "11111111-2222-3333-4444-555555555555", }, ), - func(ctx chasm.Context, component chasm.Component, _ *chasm.Registry) error { + func(ctx chasm.Context, component chasm.Component) error { s.Fail("readFn should not be called") return nil }, diff --git a/service/history/visibility_queue_task_executor.go b/service/history/visibility_queue_task_executor.go index 8cbeda8020..abb665571f 100644 --- a/service/history/visibility_queue_task_executor.go +++ b/service/history/visibility_queue_task_executor.go @@ -425,12 +425,6 @@ func (t *visibilityQueueTaskExecutor) processChasmTask( return err } - visTaskContext = chasm.AugmentContextForComponent( - visTaskContext, - rootComponent, - t.shardContext.ChasmRegistry(), - ) - var chasmTaskQueue string if chasmSAProvider, ok := rootComponent.(chasm.VisibilitySearchAttributesProvider); ok { for _, chasmSA := range chasmSAProvider.SearchAttributes(visTaskContext) {