Add CHASM Visibility component to CHASM Workflow

Author: rodrigozhou

chasm/lib/workflow/library.go

@@ -64,9 +64,13 @@ func workflowContextFromChasm(ctx chasm.Context) *workflowContext {
 
 func (l *library) Components() []*chasm.RegistrableComponent {
 	return []*chasm.RegistrableComponent{
-		chasm.NewRegistrableComponent[*Workflow](chasm.WorkflowComponentName, chasm.WithContextValues(map[any]any{
-			ctxKeyWorkflowContext: &workflowContext{registry: l.registry},
-		})),
+		chasm.NewRegistrableComponent[*Workflow](
+			chasm.WorkflowComponentName,
+			chasm.WithBusinessIDAlias("WorkflowId"),
+			chasm.WithContextValues(map[any]any{
+				ctxKeyWorkflowContext: &workflowContext{registry: l.registry},
+			}),
+		),
 		chasm.NewRegistrableComponent[*WorkflowUpdate]("update"),
 	}
 }

chasm/lib/workflow/workflow.go

@@ -12,6 +12,8 @@ import (
 	callbackspb "go.temporal.io/server/chasm/lib/callback/gen/callbackpb/v1"
 	"go.temporal.io/server/chasm/lib/nexusoperation"
 	chasmworkflowpb "go.temporal.io/server/chasm/lib/workflow/gen/workflowpb/v1"
+	"go.temporal.io/server/common/log/tag"
+	"go.temporal.io/server/common/softassert"
 	"go.temporal.io/server/service/history/historybuilder"
 	"google.golang.org/protobuf/types/known/emptypb"
 	"google.golang.org/protobuf/types/known/timestamppb"
@@ -27,6 +29,9 @@ type Workflow struct {
 	// MSPointer is a special in-memory field for accessing the underlying mutable state.
 	chasm.MSPointer
 
+	// Visibility to store custom search attributes and memo.
+	Visibility chasm.Field[*chasm.Visibility]
+
 	// Callbacks map is used to store the callbacks for the workflow.
 	Callbacks chasm.Map[string, *callback.Callback]
 
@@ -72,6 +77,73 @@ func (w *Workflow) Terminate(
 	return chasm.TerminateComponentResponse{}, serviceerror.NewInternal("workflow root Terminate should not be called")
 }
 
+// SearchAttributes returns the predefined search attributes set in the underlying mutable state.
+func (w *Workflow) SearchAttributes(ctx chasm.Context) []chasm.SearchAttributeKeyValue {
+	searchAttributes, err := w.GetPredefinedSearchAttributes()
+	softassert.That(
+		ctx.Logger(),
+		err == nil,
+		"failed to retrieve search attributes from mutable state execution info",
+		tag.Error(err),
+	)
+
+	var res []chasm.SearchAttributeKeyValue
+	for saName, value := range searchAttributes {
+		res = append(res, chasm.SearchAttributeKeyValue{
+			Alias: saName,
+			Field: saName,
+			Value: value,
+		})
+	}
+	return res
+}
+
+// CustomSearchAttributes returns the custom search attributes.
+func (w *Workflow) CustomSearchAttributes(ctx chasm.Context) map[string]*commonpb.Payload {
+	if vis, ok := w.Visibility.TryGet(ctx); ok {
+		return vis.CustomSearchAttributes(ctx)
+	}
+	return nil
+}
+
+// CustomMemo returns the custom memo.
+func (w *Workflow) CustomMemo(ctx chasm.Context) map[string]*commonpb.Payload {
+	if vis, ok := w.Visibility.TryGet(ctx); ok {
+		return vis.CustomMemo(ctx)
+	}
+	return nil
+}
+
+// UpsertCustomSearchAttributes merges the provided custom search attributes into the existing one.
+// For details of the merge, see [chasm.Visibility.MergeCustomSearchAttributes].
+func (w *Workflow) UpsertCustomSearchAttributes(
+	ctx chasm.MutableContext,
+	customSearchAttributes map[string]*commonpb.Payload,
+) error {
+	if vis, ok := w.Visibility.TryGet(ctx); ok {
+		vis.MergeCustomSearchAttributes(ctx, customSearchAttributes)
+	} else {
+		vis := chasm.NewVisibilityWithData(ctx, customSearchAttributes, nil)
+		w.Visibility = chasm.NewComponentField(ctx, vis)
+	}
+	return nil
+}
+
+// UpsertCustomMemo merges the provided custom memo into the existing one.
+// For details of the merge, see [chasm.Visibility.MergeCustomMemo].
+func (w *Workflow) UpsertCustomMemo(
+	ctx chasm.MutableContext,
+	customMemo map[string]*commonpb.Payload,
+) error {
+	if vis, ok := w.Visibility.TryGet(ctx); ok {
+		vis.MergeCustomMemo(ctx, customMemo)
+	} else {
+		vis := chasm.NewVisibilityWithData(ctx, nil, customMemo)
+		w.Visibility = chasm.NewComponentField(ctx, vis)
+	}
+	return nil
+}
+
 // ProcessCloseCallbacks triggers "WorkflowClosed" callbacks using the CHASM implementation.
 // It schedules all workflow-level and update-level callbacks that are in STANDBY state.
 func (w *Workflow) ProcessCloseCallbacks(ctx chasm.MutableContext) error {

chasm/ms_pointer.go

@@ -3,9 +3,12 @@ package chasm
 import (
 	"time"
 
+	commonpb "go.temporal.io/api/common/v1"
 	enumspb "go.temporal.io/api/enums/v1"
 	historypb "go.temporal.io/api/history/v1"
 	"go.temporal.io/server/common/nexus/nexusrpc"
+	"go.temporal.io/server/common/searchattribute/sadefs"
+	"google.golang.org/protobuf/proto"
 )
 
 // MSPointer is a special CHASM type which components can use to access their Node's underlying backend (i.e. mutable
@@ -61,3 +64,24 @@ func (m MSPointer) GetWorkflowTypeName() string {
 func (m MSPointer) GetNexusUpdateCompletion(ctx Context, updateID string, requestID string) (nexusrpc.CompleteOperationOptions, error) {
 	return m.backend.GetNexusUpdateCompletion(ctx.goContext(), updateID, requestID)
 }
+
+// GetPredefinedSearchAttributes retrieves the predefined search attributes from the underlying mutable state.
+func (m MSPointer) GetPredefinedSearchAttributes() (map[string]VisibilityValue, error) {
+	msSearchAttributes := m.backend.GetExecutionInfo().GetSearchAttributes()
+	predefinedWithType := make(map[string]*commonpb.Payload)
+	for saName, saPayload := range msSearchAttributes {
+		if saType, ok := predefinedSearchAttributes[saName]; ok {
+			if sadefs.GetMetadataType(saPayload) == enumspb.INDEXED_VALUE_TYPE_UNSPECIFIED {
+				saPayload = proto.CloneOf(saPayload)
+				sadefs.SetMetadataType(saPayload, saType)
+			}
+			predefinedWithType[saName] = saPayload
+		}
+	}
+	saMap, err := newSearchAttributesMapFromProto(
+		&commonpb.SearchAttributes{
+			IndexedFields: predefinedWithType,
+		},
+	)
+	return saMap.values, err
+}

chasm/node_backend_mock.go

@@ -22,6 +22,7 @@ var _ NodeBackend = (*MockNodeBackend)(nil)
 // fields (thread-safe).
 type MockNodeBackend struct {
 	// Optional function overrides. If nil, methods return zero-values.
+	HandleGetExecutionStateUpdated    func() bool
 	HandleGetExecutionState           func() *persistencespb.WorkflowExecutionState
 	HandleGetExecutionInfo            func() *persistencespb.WorkflowExecutionInfo
 	HandleGetCurrentVersion           func() int64
@@ -50,6 +51,13 @@ type MockNodeBackend struct {
 	}
 }
 
+func (m *MockNodeBackend) ExecutionStateUpdated() bool {
+	if m.HandleGetExecutionStateUpdated != nil {
+		return m.HandleGetExecutionStateUpdated()
+	}
+	return false
+}
+
 func (m *MockNodeBackend) GetExecutionState() *persistencespb.WorkflowExecutionState {
 	if m.HandleGetExecutionState != nil {
 		return m.HandleGetExecutionState()

chasm/search_attribute.go

@@ -9,6 +9,9 @@ import (
 	"go.temporal.io/server/common/searchattribute/sadefs"
 )
 
+// Make a copy of predefined search attributes for CHASM internal usage.
+var predefinedSearchAttributes = sadefs.Predefined()
+
 // CHASM Search Attribute User Guide:
 //
 // This contains CHASM search attribute field constants. These predefined fields correspond to the exact column name in Visibility storage.

chasm/tree.go

@@ -198,6 +198,7 @@ type (
 	// where MutableState is defined.
 	NodeBackend interface {
 		// TODO: Add methods needed from MutateState here.
+		ExecutionStateUpdated() bool
 		GetExecutionState() *persistencespb.WorkflowExecutionState
 		GetExecutionInfo() *persistencespb.WorkflowExecutionInfo
 		GetApproximatePersistedSize() int
@@ -1521,7 +1522,7 @@ func (n *Node) CloseTransaction() (NodesMutation, error) {
 		return NodesMutation{}, err
 	}
 
-	if n.isActiveStateDirty {
+	if n.isActiveStateDirty || rootLifecycleChanged {
 		if err := n.closeTransactionForceUpdateVisibility(immutableContext, rootLifecycleChanged); err != nil {
 			return NodesMutation{}, err
 		}
@@ -1591,7 +1592,7 @@ func (n *Node) closeTransactionHandleRootLifecycleChange(
 ) (bool, error) {
 	if n.backend.IsWorkflow() {
 		// Workflow manages its lifecycle directly in mutable state.
-		return false, nil
+		return n.backend.ExecutionStateUpdated(), nil
 	}
 
 	if n.valueState != valueStateNeedSerialize {

service/history/interfaces/mutable_state.go

@@ -178,6 +178,7 @@ type (
 		GetLastEventVersion() (int64, error)
 		GetExecutionInfo() *persistencespb.WorkflowExecutionInfo
 		GetExecutionState() *persistencespb.WorkflowExecutionState
+		ExecutionStateUpdated() bool
 		GetStartedWorkflowTask() *WorkflowTaskInfo
 		GetPendingWorkflowTask() *WorkflowTaskInfo
 		GetLastFirstEventIDTxnID() (int64, int64)

service/history/interfaces/mutable_state_mock.go

@@ -1113,6 +1113,10 @@ func (ms *MutableStateImpl) GetExecutionState() *persistencespb.WorkflowExecutio
 	return ms.executionState
 }
 
+func (ms *MutableStateImpl) ExecutionStateUpdated() bool {
+	return ms.stateInDB != ms.executionState.GetState()
+}
+
 func (ms *MutableStateImpl) FlushBufferedEvents() {
 	if ms.HasStartedWorkflowTask() {
 		return