Add context metadata provider interface for root components to implement (#9773)

## What changed?
Add context metadata provider interface for root components to
implement.

## Why?
Allow Component instances a method to provide key value pairs to be
propagated in ContextMetadata based off Component state.

## How did you test it?
- [X] built
- [X] run locally and tested manually
- [X] covered by existing tests
- [X] added new unit test(s)
- [ ] added new functional test(s)

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/lib/activity/activity.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/scheduler/scheduler.go

@@ -121,6 +121,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/tests/payload.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/workflow/workflow.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/test_component_test.go

@@ -11,6 +11,7 @@ import (
 	enumsspb "go.temporal.io/server/api/enums/v1"
 	"go.temporal.io/server/chasm"
 	"go.temporal.io/server/common"
+	"go.temporal.io/server/common/contextutil"
 	"go.temporal.io/server/common/convert"
 	"go.temporal.io/server/common/definition"
 	"go.temporal.io/server/common/headers"
@@ -108,6 +109,69 @@ func (e *ChasmEngine) NotifyExecution(key chasm.ExecutionKey) {
 	e.notifier.Notify(key)
 }
 
+func (e *ChasmEngine) setContextMetadata(
+	ctx context.Context,
+	chasmTree *chasm.Node,
+) chasm.Context {
+	chasmContext := chasm.NewContext(ctx, chasmTree)
+
+	rootComponent, err := chasmTree.Component(chasmContext, chasm.ComponentRef{})
+	if err != nil {
+		executionKey := chasmContext.ExecutionKey()
+		e.logger.Error(
+			"Failed to resolve CHASM root component for context metadata",
+			tag.WorkflowNamespaceID(executionKey.NamespaceID),
+			tag.WorkflowID(executionKey.BusinessID),
+			tag.WorkflowRunID(executionKey.RunID),
+			tag.Error(err),
+		)
+		return chasmContext
+	}
+
+	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 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, 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
+}
+
+func (e *ChasmEngine) setContextMetadataFromMutableState(
+	ctx context.Context,
+	mutableState historyi.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
+	}
+	e.setContextMetadata(ctx, chasmTree)
+}
+
 func (e *ChasmEngine) StartExecution(
 	ctx context.Context,
 	executionRef chasm.ComponentRef,
@@ -177,6 +241,7 @@ func (e *ChasmEngine) startExecution(
 		return chasm.StartExecutionResult{}, err
 	}
 	if !hasCurrentRun {
+		e.setContextMetadataFromMutableState(ctx, newExecutionParams.mutableState)
 		serializedRef, err := newExecutionParams.executionRef.Serialize(e.registry)
 		if err != nil {
 			// Created is true here because persistAsBrandNew succeeded, but we failed to serialize the ref.
@@ -354,13 +419,9 @@ func (e *ChasmEngine) applyUpdateWithLease(
 	updateFn func(chasm.MutableContext, chasm.Component) error,
 ) ([]byte, error) {
 	mutableState := executionLease.GetMutableState()
-	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{},
-		)
+	chasmTree, err := chasmTreeFromMutableState(shardContext.GetLogger(), mutableState)
+	if err != nil {
+		return nil, err
 	}
 
 	mutableContext := chasm.NewMutableContext(ctx, chasmTree)
@@ -375,6 +436,8 @@ func (e *ChasmEngine) applyUpdateWithLease(
 
 	// TODO: Support WithSpeculative() TransitionOption.
 
+	e.setContextMetadata(ctx, chasmTree)
+
 	if err := executionLease.GetContext().UpdateWorkflowExecutionAsActive(
 		ctx,
 		shardContext,
@@ -438,6 +501,8 @@ func (e *ChasmEngine) startAndUpdateExecution(
 		return chasm.ExecutionKey{}, nil, false, currentRunInfo.CurrentWorkflowConditionFailedError
 	}
 
+	e.setContextMetadataFromMutableState(ctx, newExecutionParams.mutableState)
+
 	serializedRef, err := newExecutionParams.executionRef.Serialize(e.registry)
 
 	return newExecutionParams.executionRef.ExecutionKey, serializedRef, true, err
@@ -502,6 +567,8 @@ func (e *ChasmEngine) deleteExecution(
 	mutableState := executionLease.GetMutableState()
 	we := mutableState.GetWorkflowKey()
 
+	e.setContextMetadataFromMutableState(ctx, mutableState)
+
 	log.With(shardContext.GetLogger(),
 		tag.WorkflowNamespaceID(ref.NamespaceID),
 		tag.WorkflowID(we.WorkflowID),
@@ -580,16 +647,12 @@ func (e *ChasmEngine) readComponent(
 		executionLease.GetReleaseFn()(nil)
 	}()
 
-	chasmTree, ok := executionLease.GetMutableState().ChasmTree().(*chasm.Node)
-	if !ok {
-		return serviceerror.NewInternalf(
-			"CHASM tree implementation not properly wired up, encountered type: %T, expected type: %T",
-			executionLease.GetMutableState().ChasmTree(),
-			&chasm.Node{},
-		)
+	chasmTree, err := chasmTreeFromMutableState(e.logger, executionLease.GetMutableState())
+	if err != nil {
+		return err
 	}
 
-	chasmContext := chasm.NewContext(ctx, chasmTree)
+	chasmContext := e.setContextMetadata(ctx, chasmTree)
 	component, err := chasmTree.Component(chasmContext, ref)
 	if err != nil {
 		return err
@@ -686,16 +749,13 @@ func (e *ChasmEngine) predicateSatisfied(
 	ref chasm.ComponentRef,
 	executionLease api.WorkflowLease,
 ) ([]byte, error) {
-	chasmTree, ok := executionLease.GetMutableState().ChasmTree().(*chasm.Node)
-	if !ok {
-		return nil, serviceerror.NewInternalf(
-			"CHASM tree implementation not properly wired up, encountered type: %T, expected type: %T",
-			executionLease.GetMutableState().ChasmTree(),
-			&chasm.Node{},
-		)
+	chasmTree, err := chasmTreeFromMutableState(e.logger, executionLease.GetMutableState())
+	if err != nil {
+		return nil, err
 	}
 
-	chasmContext := chasm.NewContext(ctx, chasmTree)
+	chasmContext := e.setContextMetadata(ctx, chasmTree)
+
 	component, err := chasmTree.Component(chasmContext, ref)
 	if err != nil {
 		return nil, err
@@ -1052,6 +1112,8 @@ func (e *ChasmEngine) handleReusePolicy(
 		return chasm.StartExecutionResult{}, err
 	}
 
+	e.setContextMetadataFromMutableState(ctx, newExecutionParams.mutableState)
+
 	serializedRef, err := newExecutionParams.executionRef.Serialize(e.registry)
 	if err != nil {
 		return chasm.StartExecutionResult{ExecutionKey: newExecutionParams.executionRef.ExecutionKey, Created: true}, err