Move to RootComponent method

Author: awln-temporal

chasm/component.go

@@ -44,6 +44,9 @@ type TerminateComponentResponse struct{}
 // TODO: (not yet true) Visibility record will no longer be updated after RootComponent is closed.
 type RootComponent interface {
 	TerminableComponent
+
+	// ContextMetadata returns execution metadata to propagate to the request context.
+	ContextMetadata(Context) map[string]string
 }
 
 // Embed UnimplementedComponent to get forward compatibility

chasm/component_mock.go

@@ -106,6 +106,11 @@ func (a *Activity) LifecycleState(_ chasm.Context) chasm.LifecycleState {
 	}
 }
 
+func (a *Activity) ContextMetadata(_ chasm.Context) map[string]string {
+	// TODO: Export standalone activity context metadata.
+	return nil
+}
+
 // NewStandaloneActivity creates a new activity component and adds associated tasks to start execution.
 func NewStandaloneActivity(
 	ctx chasm.MutableContext,

chasm/context_metadata.go

@@ -288,6 +288,11 @@ func (s *Scheduler) LifecycleState(ctx chasm.Context) chasm.LifecycleState {
 	return chasm.LifecycleStateRunning
 }
 
+func (s *Scheduler) ContextMetadata(_ chasm.Context) map[string]string {
+	// TODO: Export scheduler context metadata.
+	return nil
+}
+
 // Terminate implements the chasm.RootComponent interface.
 func (s *Scheduler) Terminate(
 	_ chasm.MutableContext,

chasm/lib/activity/activity.go

@@ -115,6 +115,12 @@ func (s *PayloadStore) Cancel(
 	return CancelPayloadStoreResponse{}, nil
 }
 
+func (s *PayloadStore) ContextMetadata(_ chasm.Context) map[string]string {
+	return map[string]string{
+		string(componentCtxKey): componentCtxVal,
+	}
+}
+
 func (s *PayloadStore) AddPayload(
 	mutableContext chasm.MutableContext,
 	request AddPayloadRequest,

chasm/lib/scheduler/scheduler.go

@@ -46,6 +46,11 @@ func (w *Workflow) LifecycleState(
 	return chasm.LifecycleStateRunning
 }
 
+func (w *Workflow) ContextMetadata(_ chasm.Context) map[string]string {
+	// TODO: Export workflow metadata from the CHASM workflow root instead of CloseTransaction().
+	return nil
+}
+
 // ProcessCloseCallbacks triggers "WorkflowClosed" callbacks using the CHASM implementation.
 // It iterates through all callbacks and schedules WorkflowClosed ones that are in STANDBY state.
 func (w *Workflow) ProcessCloseCallbacks(ctx chasm.MutableContext) error {

chasm/lib/tests/payload.go

@@ -80,6 +80,7 @@ var (
 
 	_ VisibilitySearchAttributesProvider = (*TestComponent)(nil)
 	_ VisibilityMemoProvider             = (*TestComponent)(nil)
+	_ RootComponent                      = (*TestComponent)(nil)
 )
 
 func (tc *TestComponent) LifecycleState(_ Context) LifecycleState {
@@ -109,6 +110,11 @@ func (tc *TestComponent) Fail(_ MutableContext) {
 	tc.ComponentData.Status = enumspb.WORKFLOW_EXECUTION_STATUS_FAILED
 }
 
+func (tc *TestComponent) ContextMetadata(_ Context) map[string]string {
+	// TODO: Export context metadata from this test root.
+	return nil
+}
+
 // SearchAttributes implements VisibilitySearchAttributesProvider interface.
 func (tc *TestComponent) SearchAttributes(_ Context) []SearchAttributeKeyValue {
 	return []SearchAttributeKeyValue{

chasm/lib/workflow/workflow.go

@@ -128,27 +128,33 @@ func (e *ChasmEngine) setContextMetadata(
 		return chasmContext
 	}
 
-	provider, ok := rootComponent.(chasm.ContextMetadataProvider)
+	root, ok := rootComponent.(chasm.RootComponent)
 	if !ok {
+		softassert.Fail(
+			e.logger,
+			"root node must implement RootComponent interface",
+			tag.NewStringTag("component_type", fmt.Sprintf("%T", rootComponent)),
+		)
 		return chasmContext
 	}
 
-	for key, value := range provider.ContextMetadata(chasmContext) {
+	for key, value := range root.ContextMetadata(chasmContext) {
 		contextutil.ContextMetadataSet(ctx, key, value)
 	}
 
 	return chasmContext
 }
 
 func chasmTreeFromMutableState(
+	logger log.Logger,
 	mutableState historyi.MutableState,
 ) (*chasm.Node, error) {
 	chasmTree, ok := mutableState.ChasmTree().(*chasm.Node)
 	if !ok {
-		return nil, serviceerror.NewInternalf(
-			"CHASM tree implementation not properly wired up, encountered type: %T, expected type: %T",
-			mutableState.ChasmTree(),
-			&chasm.Node{},
+		return nil, softassert.UnexpectedInternalErr(
+			logger,
+			"CHASM tree implementation not properly wired up",
+			fmt.Errorf("encountered type: %T, expected type: %T", mutableState.ChasmTree(), &chasm.Node{}),
 		)
 	}
 	return chasmTree, nil
@@ -158,7 +164,7 @@ func (e *ChasmEngine) setContextMetadataFromMutableState(
 	ctx context.Context,
 	mutableState historyi.MutableState,
 ) {
-	chasmTree, err := chasmTreeFromMutableState(mutableState)
+	chasmTree, err := chasmTreeFromMutableState(e.logger, mutableState)
 	if err != nil {
 		e.logger.Error("Failed to resolve CHASM tree for context metadata", tag.Error(err))
 		return
@@ -413,7 +419,7 @@ func (e *ChasmEngine) applyUpdateWithLease(
 	updateFn func(chasm.MutableContext, chasm.Component) error,
 ) ([]byte, error) {
 	mutableState := executionLease.GetMutableState()
-	chasmTree, err := chasmTreeFromMutableState(mutableState)
+	chasmTree, err := chasmTreeFromMutableState(shardContext.GetLogger(), mutableState)
 	if err != nil {
 		return nil, err
 	}
@@ -641,7 +647,7 @@ func (e *ChasmEngine) readComponent(
 		executionLease.GetReleaseFn()(nil)
 	}()
 
-	chasmTree, err := chasmTreeFromMutableState(executionLease.GetMutableState())
+	chasmTree, err := chasmTreeFromMutableState(e.logger, executionLease.GetMutableState())
 	if err != nil {
 		return err
 	}
@@ -743,7 +749,7 @@ func (e *ChasmEngine) predicateSatisfied(
 	ref chasm.ComponentRef,
 	executionLease api.WorkflowLease,
 ) ([]byte, error) {
-	chasmTree, err := chasmTreeFromMutableState(executionLease.GetMutableState())
+	chasmTree, err := chasmTreeFromMutableState(e.logger, executionLease.GetMutableState())
 	if err != nil {
 		return nil, err
 	}

chasm/test_component_test.go

@@ -550,7 +550,7 @@ func (s *chasmEngineSuite) TestSetContextMetadata_StateAndRequestScopedValues()
 		},
 	)
 
-	chasmTree, err := chasmTreeFromMutableState(mutableState)
+	chasmTree, err := chasmTreeFromMutableState(s.mockShard.GetLogger(), mutableState)
 	s.NoError(err)
 
 	requestCtx := newTestMetadataContext("helper-request")
@@ -575,7 +575,7 @@ func (s *chasmEngineSuite) TestSetContextMetadata_NoProvider() {
 		},
 	)
 
-	chasmTree, err := chasmTreeFromMutableState(mutableState)
+	chasmTree, err := chasmTreeFromMutableState(s.mockShard.GetLogger(), mutableState)
 	s.NoError(err)
 
 	requestCtx := contextutil.WithMetadataContext(context.Background())
@@ -2002,7 +2002,7 @@ func (s *chasmEngineSuite) newTestMutableState(
 		s.mockShard.GetTimeSource().Now(),
 	)
 
-	chasmTree, err := chasmTreeFromMutableState(mutableState)
+	chasmTree, err := chasmTreeFromMutableState(s.mockShard.GetLogger(), mutableState)
 	s.NoError(err)
 	s.NoError(chasmTree.SetRootComponent(rootComponent))
 
@@ -2043,7 +2043,8 @@ var (
 
 	_ chasm.VisibilitySearchAttributesProvider = (*testComponent)(nil)
 	_ chasm.VisibilityMemoProvider             = (*testComponent)(nil)
-	_ chasm.ContextMetadataProvider            = (*testComponent)(nil)
+	_ chasm.RootComponent                      = (*testComponent)(nil)
+	_ chasm.RootComponent                      = (*testComponentNoMetadata)(nil)
 )
 
 type testRequestContextKey struct{}
@@ -2099,6 +2100,11 @@ func (l *testComponentNoMetadata) LifecycleState(_ chasm.Context) chasm.Lifecycl
 	return chasm.LifecycleStateRunning
 }
 
+func (l *testComponentNoMetadata) ContextMetadata(_ chasm.Context) map[string]string {
+	// TODO: Export context metadata from this root.
+	return nil
+}
+
 func (l *testComponentNoMetadata) Terminate(
 	_ chasm.MutableContext,
 	_ chasm.TerminateComponentRequest,