History: Expose and call SetContextMetadata in readonly operations

Author: nikki-dag

common/contextutil/metadata.go

@@ -17,6 +17,13 @@ type (
 
 var metadataCtxKey = metadataContextKey{}
 
+const (
+	// MetadataKeyWorkflowType is the context metadata key for workflow type
+	MetadataKeyWorkflowType = "workflow-type"
+	// MetadataKeyWorkflowTaskQueue is the context metadata key for workflow task queue
+	MetadataKeyWorkflowTaskQueue = "workflow-task-queue"
+)
+
 // getMetadataContext extracts metadata context from golang context.
 func getMetadataContext(ctx context.Context) *metadataContext {
 	metadataCtx := ctx.Value(metadataCtxKey)
@@ -38,6 +45,12 @@ func WithMetadataContext(ctx context.Context) context.Context {
 	return context.WithValue(ctx, metadataCtxKey, metadataCtx)
 }
 
+// ContextHasMetadata returns true if the context has metadata support.
+// This can be used to debug whether a context has been properly initialized with metadata.
+func ContextHasMetadata(ctx context.Context) bool {
+	return getMetadataContext(ctx) != nil
+}
+
 // ContextMetadataSet sets a metadata key-value pair in the context.
 func ContextMetadataSet(ctx context.Context, key string, value any) bool {
 	metadataCtx := getMetadataContext(ctx)

common/contextutil/metadata_test.go

@@ -301,3 +301,63 @@ func TestMetadataContextIsolation(t *testing.T) {
 		assert.Equal(t, "value2", value2)
 	})
 }
+
+func TestContextHasMetadata(t *testing.T) {
+	t.Run("returns true when context has metadata", func(t *testing.T) {
+		ctx := WithMetadataContext(context.Background())
+		assert.True(t, ContextHasMetadata(ctx))
+	})
+
+	t.Run("returns false for context without metadata", func(t *testing.T) {
+		ctx := context.Background()
+		assert.False(t, ContextHasMetadata(ctx))
+	})
+
+	t.Run("returns true after setting metadata values", func(t *testing.T) {
+		ctx := WithMetadataContext(context.Background())
+		ContextMetadataSet(ctx, "key1", "value1")
+		ContextMetadataSet(ctx, "key2", "value2")
+
+		assert.True(t, ContextHasMetadata(ctx))
+	})
+
+	t.Run("returns true for empty metadata context", func(t *testing.T) {
+		ctx := WithMetadataContext(context.Background())
+		// No values set, but metadata context exists
+		assert.True(t, ContextHasMetadata(ctx))
+	})
+
+	t.Run("child context inherits metadata from parent", func(t *testing.T) {
+		parentCtx := WithMetadataContext(context.Background())
+		ContextMetadataSet(parentCtx, "key", "value")
+
+		type testContextKey string
+		childCtx := context.WithValue(parentCtx, testContextKey("other-key"), "other-value")
+
+		assert.True(t, ContextHasMetadata(parentCtx))
+		assert.True(t, ContextHasMetadata(childCtx))
+	})
+
+	t.Run("returns false for wrong type in context", func(t *testing.T) {
+		ctx := context.WithValue(context.Background(), metadataCtxKey, "wrong type")
+		assert.False(t, ContextHasMetadata(ctx))
+	})
+
+	t.Run("returns true for cancelled context with metadata", func(t *testing.T) {
+		ctx, cancel := context.WithCancel(context.Background())
+		ctx = WithMetadataContext(ctx)
+		cancel()
+
+		assert.True(t, ContextHasMetadata(ctx))
+	})
+
+	t.Run("multiple contexts with metadata are independent", func(t *testing.T) {
+		ctx1 := WithMetadataContext(context.Background())
+		ctx2 := WithMetadataContext(context.Background())
+		ctx3 := context.Background()
+
+		assert.True(t, ContextHasMetadata(ctx1))
+		assert.True(t, ContextHasMetadata(ctx2))
+		assert.False(t, ContextHasMetadata(ctx3))
+	})
+}

service/history/api/queryworkflow/api.go

@@ -80,6 +80,11 @@ func Invoke(
 		workflowLease.GetReleaseFn()(nil)
 	}()
 
+	// Context metadata is automatically set during mutable state transaction close for operations that mutate state.
+	// Since QueryWorkflow is readonly and never closes the transaction, we explicitly call SetContextMetadata
+	// here to ensure successful requests have workflow metadata populated in the context.
+	workflowLease.GetMutableState().SetContextMetadata(ctx)
+
 	req := request.GetRequest()
 	_, mutableStateStatus := workflowLease.GetMutableState().GetWorkflowStateStatus()
 	scope = scope.WithTags(metrics.StringTag("workflow_status", mutableStateStatus.String()))

service/history/api/respondworkflowtaskcompleted/api.go

@@ -180,6 +180,13 @@ func (handler *WorkflowTaskCompletedHandler) Invoke(
 			}
 		}
 
+		if len(request.Commands) == 0 && len(request.Messages) > 0 {
+			// Context metadata is automatically set during mutable state transaction close. For RespondWorkflowTaskCompleted
+			// with only `update.Rejection` messages, the transaction is never closed. We explicitly call SetContextMetadata
+			// here to ensure readonly messages have workflow metadata populated in the context.
+			ms.SetContextMetadata(ctx)
+		}
+
 		workflowLease.GetReleaseFn()(errForRelease)
 	}()
 

service/history/interfaces/mutable_state.go

@@ -289,6 +289,7 @@ type (
 		UpdateBuildIdAssignment(buildId string) error
 		ApplyBuildIdRedirect(startingTaskScheduledEventId int64, buildId string, redirectCounter int64) error
 		RefreshExpirationTimeoutTask(ctx context.Context) error
+		SetContextMetadata(context.Context)
 
 		GetHistorySize() int64
 		AddHistorySize(size int64)

service/history/interfaces/mutable_state_mock.go

@@ -6993,19 +6993,19 @@ type closeTransactionResult struct {
 	chasmNodesMutation chasm.NodesMutation
 }
 
-func (ms *MutableStateImpl) setMetaDataMap(
+func (ms *MutableStateImpl) SetContextMetadata(
 	ctx context.Context,
 ) {
 	switch ms.chasmTree.ArchetypeID() {
 	case chasm.WorkflowArchetypeID, chasm.UnspecifiedArchetypeID:
 		// Set workflow type
 		if wfType := ms.GetWorkflowType(); wfType != nil && wfType.GetName() != "" {
-			contextutil.ContextMetadataSet(ctx, "workflow-type", wfType.GetName())
+			contextutil.ContextMetadataSet(ctx, contextutil.MetadataKeyWorkflowType, wfType.GetName())
 		}
 
 		// Set workflow task queue
 		if ms.executionInfo != nil && ms.executionInfo.TaskQueue != "" {
-			contextutil.ContextMetadataSet(ctx, "workflow-task-queue", ms.executionInfo.TaskQueue)
+			contextutil.ContextMetadataSet(ctx, contextutil.MetadataKeyWorkflowTaskQueue, ms.executionInfo.TaskQueue)
 		}
 
 		// TODO: To set activity_type/activity_task_queue metadata, the history gRPC handler should
@@ -7019,7 +7019,7 @@ func (ms *MutableStateImpl) closeTransaction(
 	ctx context.Context,
 	transactionPolicy historyi.TransactionPolicy,
 ) (closeTransactionResult, error) {
-	ms.setMetaDataMap(ctx)
+	ms.SetContextMetadata(ctx)
 
 	if err := ms.closeTransactionWithPolicyCheck(
 		transactionPolicy,

service/history/workflow/mutable_state_impl.go

@@ -34,6 +34,7 @@ import (
 	"go.temporal.io/server/chasm"
 	"go.temporal.io/server/common"
 	"go.temporal.io/server/common/cluster"
+	"go.temporal.io/server/common/contextutil"
 	"go.temporal.io/server/common/definition"
 	"go.temporal.io/server/common/dynamicconfig"
 	"go.temporal.io/server/common/failure"
@@ -6110,3 +6111,40 @@ func (s *mutableStateSuite) TestCHASMNodeSize() {
 	expectedTotalSize += len(newNodeKey) + newNode.Size()
 	s.Equal(expectedTotalSize, mutableState.GetApproximatePersistedSize())
 }
+
+func (s *mutableStateSuite) TestSetContextMetadata() {
+	s.mockEventsCache.EXPECT().PutEvent(gomock.Any(), gomock.Any()).Times(1)
+
+	ctx := contextutil.WithMetadataContext(context.Background())
+	workflowType := "test-workflow-type"
+	taskQueue := "test-task-queue"
+
+	execution := &commonpb.WorkflowExecution{
+		WorkflowId: tests.WorkflowID,
+		RunId:      tests.RunID,
+	}
+
+	_, err := s.mutableState.AddWorkflowExecutionStartedEvent(
+		execution,
+		&historyservice.StartWorkflowExecutionRequest{
+			NamespaceId: tests.NamespaceID.String(),
+			StartRequest: &workflowservice.StartWorkflowExecutionRequest{
+				WorkflowId:   tests.WorkflowID,
+				WorkflowType: &commonpb.WorkflowType{Name: workflowType},
+				TaskQueue:    &taskqueuepb.TaskQueue{Name: taskQueue},
+			},
+		},
+	)
+	s.NoError(err)
+
+	s.mutableState.SetContextMetadata(ctx)
+
+	// Verify metadata was set correctly
+	wfType, ok := contextutil.ContextMetadataGet(ctx, contextutil.MetadataKeyWorkflowType)
+	s.True(ok)
+	s.Equal(workflowType, wfType)
+
+	tq, ok := contextutil.ContextMetadataGet(ctx, contextutil.MetadataKeyWorkflowTaskQueue)
+	s.True(ok)
+	s.Equal(taskQueue, tq)
+}