CHASM: Improve visibility component (#8900)

## What changed?
- Rename Set* methods to Merge*.
- Add Replace* methods.
- Automatically delete SA/Memo field when it’s empty.

## Why?
- Improve visibility component for custom SA & memo

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

chasm/lib/activity/activity.go

@@ -656,7 +656,7 @@ func (a *Activity) buildActivityExecutionInfo(ctx chasm.Context) (*apiactivitypb
 	}
 
 	sa := &commonpb.SearchAttributes{
-		IndexedFields: a.Visibility.Get(ctx).GetSearchAttributes(ctx),
+		IndexedFields: a.Visibility.Get(ctx).CustomSearchAttributes(ctx),
 	}
 
 	info := &apiactivitypb.ActivityExecutionInfo{

chasm/lib/scheduler/scheduler.go

@@ -156,8 +156,8 @@ func CreateScheduler(
 
 	// Update visibility with custom attributes.
 	visibility := sched.Visibility.Get(ctx)
-	visibility.SetSearchAttributes(ctx, req.FrontendRequest.GetSearchAttributes().GetIndexedFields())
-	visibility.SetMemo(ctx, req.FrontendRequest.GetMemo().GetFields())
+	visibility.MergeCustomSearchAttributes(ctx, req.FrontendRequest.GetSearchAttributes().GetIndexedFields())
+	visibility.MergeCustomMemo(ctx, req.FrontendRequest.GetMemo().GetFields())
 
 	return sched, &schedulerpb.CreateScheduleResponse{
 		FrontendResponse: &workflowservice.CreateScheduleResponse{
@@ -542,7 +542,7 @@ func (s *Scheduler) Describe(
 	}
 
 	visibility := s.Visibility.Get(ctx)
-	memo := visibility.GetMemo(ctx)
+	memo := visibility.CustomMemo(ctx)
 	delete(memo, visibilityMemoFieldInfo) // We don't need to return a redundant info block.
 
 	if s.Schedule.GetPolicies().GetOverlapPolicy() == enumspb.SCHEDULE_OVERLAP_POLICY_UNSPECIFIED {
@@ -562,7 +562,7 @@ func (s *Scheduler) Describe(
 			Info:             common.CloneProto(s.Info),
 			ConflictToken:    s.generateConflictToken(),
 			Memo:             &commonpb.Memo{Fields: memo},
-			SearchAttributes: &commonpb.SearchAttributes{IndexedFields: visibility.GetSearchAttributes(ctx)},
+			SearchAttributes: &commonpb.SearchAttributes{IndexedFields: visibility.CustomSearchAttributes(ctx)},
 		},
 	}, nil
 }
@@ -629,7 +629,7 @@ func (s *Scheduler) Update(
 
 		// Preserve the old custom memo in the new Visibility component.
 		oldVisibility := s.Visibility.Get(ctx)
-		oldMemo := oldVisibility.GetMemo(ctx)
+		oldMemo := oldVisibility.CustomMemo(ctx)
 
 		visibility := chasm.NewVisibilityWithData(ctx, req.FrontendRequest.GetSearchAttributes().GetIndexedFields(), oldMemo)
 		s.Visibility = chasm.NewComponentField(ctx, visibility)

chasm/visibility.go

@@ -131,6 +131,8 @@ type Visibility struct {
 
 	Data *persistencespb.ChasmVisibilityData
 
+	// Do NOT access those fields directly.
+	// Use the provided getters and setters instead.
 	SA   Field[*commonpb.SearchAttributes]
 	Memo Field[*commonpb.Memo]
 }
@@ -157,14 +159,19 @@ func NewVisibilityWithData(
 		Data: &persistencespb.ChasmVisibilityData{
 			TransitionCount: 0,
 		},
-		SA: NewDataField(
+	}
+
+	if len(customSearchAttributes) != 0 {
+		visibility.SA = NewDataField(
 			mutableContext,
 			&commonpb.SearchAttributes{IndexedFields: customSearchAttributes},
-		),
-		Memo: NewDataField(
+		)
+	}
+	if len(customMemo) != 0 {
+		visibility.Memo = NewDataField(
 			mutableContext,
 			&commonpb.Memo{Fields: customMemo},
-		),
+		)
 	}
 
 	visibility.generateTask(mutableContext)
@@ -175,7 +182,12 @@ func (v *Visibility) LifecycleState(_ Context) LifecycleState {
 	return LifecycleStateRunning
 }
 
-func (v *Visibility) GetSearchAttributes(
+// CustomSearchAttributes returns the stored custom search attribute fields.
+// Nil is returned if there are none.
+//
+// Returned map is NOT a deep copy of the underlying data, so do NOT modify it
+// directly, use Merge/ReplaceCustomSearchAttributes methods instead.
+func (v *Visibility) CustomSearchAttributes(
 	chasmContext Context,
 ) map[string]*commonpb.Payload {
 	sa, _ := v.SA.TryGet(chasmContext)
@@ -183,12 +195,22 @@ func (v *Visibility) GetSearchAttributes(
 	return sa.GetIndexedFields()
 }
 
-func (v *Visibility) SetSearchAttributes(
+// MergeCustomSearchAttributes merges the provided custom search attribute fields into the existing ones.
+//   - If a key in `customSearchAttributes` already exists,
+//     the value in `customSearchAttributes` replaces the existing value.
+//   - If a key in `customSearchAttributes` has nil or empty slice payload value,
+//     the key is deleted from the existing search attributes if it exists.
+//     If all search attributes are removed, the underlying search attributes node is deleted.
+//   - If `customSearchAttributes` is empty, this is a no-op.
+func (v *Visibility) MergeCustomSearchAttributes(
 	mutableContext MutableContext,
 	customSearchAttributes map[string]*commonpb.Payload,
 ) {
-	currentSA, ok := v.SA.TryGet(mutableContext)
+	if len(customSearchAttributes) == 0 {
+		return
+	}
 
+	currentSA, ok := v.SA.TryGet(mutableContext)
 	if !ok {
 		currentSA = &commonpb.SearchAttributes{}
 		v.SA = NewDataField(mutableContext, currentSA)
@@ -198,10 +220,43 @@ func (v *Visibility) SetSearchAttributes(
 		currentSA.GetIndexedFields(),
 		customSearchAttributes,
 	)
+	if len(currentSA.IndexedFields) == 0 {
+		v.SA = NewEmptyField[*commonpb.SearchAttributes]()
+	}
+
 	v.generateTask(mutableContext)
 }
 
-func (v *Visibility) GetMemo(
+// ReplaceCustomSearchAttributes replaces the existing custom search attribute fields with the provided ones.
+// If `customSearchAttributes` is empty, the underlying search attributes node is deleted.
+func (v *Visibility) ReplaceCustomSearchAttributes(
+	mutableContext MutableContext,
+	customSearchAttributes map[string]*commonpb.Payload,
+) {
+	if len(customSearchAttributes) == 0 {
+		_, ok := v.SA.TryGet(mutableContext)
+		if !ok {
+			// Already empty, no-op
+			return
+		}
+
+		v.SA = NewEmptyField[*commonpb.SearchAttributes]()
+	} else {
+		v.SA = NewDataField(
+			mutableContext,
+			&commonpb.SearchAttributes{IndexedFields: customSearchAttributes},
+		)
+	}
+
+	v.generateTask(mutableContext)
+}
+
+// CustomMemo returns the stored custom memo fields.
+// Nil is returned if there are none.
+//
+// Returned map is NOT a deep copy of the underlying data, so do NOT modify it
+// directly, use Merge/ReplaceCustomMemo methods instead.
+func (v *Visibility) CustomMemo(
 	chasmContext Context,
 ) map[string]*commonpb.Payload {
 	memo, _ := v.Memo.TryGet(chasmContext)
@@ -209,12 +264,22 @@ func (v *Visibility) GetMemo(
 	return memo.GetFields()
 }
 
-func (v *Visibility) SetMemo(
+// MergeCustomMemo merges the provided custom memo fields into the existing ones.
+//   - If a key in `customMemo` already exists,
+//     the value in `customMemo` replaces the existing value.
+//   - If a key in `customMemo` has nil or empty slice payload value,
+//     the key is deleted from the existing memo if it exists.
+//     If all memo fields are removed, the underlying memo node is deleted.
+//   - If `customMemo` is empty, this is a no-op.
+func (v *Visibility) MergeCustomMemo(
 	mutableContext MutableContext,
 	customMemo map[string]*commonpb.Payload,
 ) {
-	currentMemo, ok := v.Memo.TryGet(mutableContext)
+	if len(customMemo) == 0 {
+		return
+	}
 
+	currentMemo, ok := v.Memo.TryGet(mutableContext)
 	if !ok {
 		currentMemo = &commonpb.Memo{}
 		v.Memo = NewDataField(mutableContext, currentMemo)
@@ -224,6 +289,33 @@ func (v *Visibility) SetMemo(
 		currentMemo.GetFields(),
 		customMemo,
 	)
+	if len(currentMemo.Fields) == 0 {
+		v.Memo = NewEmptyField[*commonpb.Memo]()
+	}
+	v.generateTask(mutableContext)
+}
+
+// ReplaceCustomMemo replaces the existing custom memo fields with the provided ones.
+// If `customMemo` is empty, the underlying memo node is deleted.
+func (v *Visibility) ReplaceCustomMemo(
+	mutableContext MutableContext,
+	customMemo map[string]*commonpb.Payload,
+) {
+	if len(customMemo) == 0 {
+		_, ok := v.Memo.TryGet(mutableContext)
+		if !ok {
+			// Already empty, no-op
+			return
+		}
+
+		v.Memo = NewEmptyField[*commonpb.Memo]()
+	} else {
+		v.Memo = NewDataField(
+			mutableContext,
+			&commonpb.Memo{Fields: customMemo},
+		)
+	}
+
 	v.generateTask(mutableContext)
 }
 

chasm/visibility_test.go

@@ -68,8 +68,8 @@ func (s *visibilitySuite) TestLifeCycleState() {
 	s.Equal(LifecycleStateRunning, s.visibility.LifecycleState(s.mockMutableContext))
 }
 
-func (s *visibilitySuite) TestSearchAttributes() {
-	sa := s.visibility.GetSearchAttributes(s.mockMutableContext)
+func (s *visibilitySuite) TestMergeCustomSearchAttributes() {
+	sa := s.visibility.CustomSearchAttributes(s.mockMutableContext)
 	s.Empty(sa)
 
 	stringKey, stringVal := "stringKey", "stringValue"
@@ -77,7 +77,7 @@ func (s *visibilitySuite) TestSearchAttributes() {
 	floatKey, floatVal := "floatKey", 3.14
 
 	// Add SA via Visibility struct method.
-	s.visibility.SetSearchAttributes(
+	s.visibility.MergeCustomSearchAttributes(
 		s.mockMutableContext,
 		map[string]*commonpb.Payload{
 			stringKey: s.mustEncode(stringVal),
@@ -86,9 +86,9 @@ func (s *visibilitySuite) TestSearchAttributes() {
 		},
 	)
 	s.Len(s.mockMutableContext.Tasks, 1)
-	protorequire.ProtoEqual(s.T(), &persistencespb.ChasmVisibilityTaskData{TransitionCount: 2}, s.mockMutableContext.Tasks[0].Payload.(*persistencespb.ChasmVisibilityTaskData))
+	s.assertTaskPayload(2, s.mockMutableContext.Tasks[0].Payload)
 
-	sa = s.visibility.GetSearchAttributes(s.mockMutableContext)
+	sa = s.visibility.CustomSearchAttributes(s.mockMutableContext)
 	s.Len(sa, 3)
 
 	var actualStringVal string
@@ -107,21 +107,79 @@ func (s *visibilitySuite) TestSearchAttributes() {
 	s.Equal(floatVal, actualFloatVal)
 
 	// Test remove search attributes by setting payload to nil.
-	s.visibility.SetSearchAttributes(s.mockMutableContext, map[string]*commonpb.Payload{
+	s.visibility.MergeCustomSearchAttributes(s.mockMutableContext, map[string]*commonpb.Payload{
 		intKey:   s.mustEncode(intVal),
 		floatKey: nil,
 	})
 	s.NoError(err)
 	s.Len(s.mockMutableContext.Tasks, 2)
-	protorequire.ProtoEqual(s.T(), &persistencespb.ChasmVisibilityTaskData{TransitionCount: 3}, s.mockMutableContext.Tasks[1].Payload.(*persistencespb.ChasmVisibilityTaskData))
+	s.assertTaskPayload(3, s.mockMutableContext.Tasks[1].Payload)
 
-	sa = s.visibility.GetSearchAttributes(s.mockMutableContext)
+	sa = s.visibility.CustomSearchAttributes(s.mockMutableContext)
 	s.NoError(err)
 	s.Len(sa, 2, "intKey and stringKey should remain")
+
+	// Test removing all search attributes also removes the node.
+	s.visibility.MergeCustomSearchAttributes(s.mockMutableContext, map[string]*commonpb.Payload{
+		stringKey: nil,
+		intKey:    nil,
+	})
+	s.Len(s.mockMutableContext.Tasks, 3)
+	s.assertTaskPayload(4, s.mockMutableContext.Tasks[2].Payload)
+	_, ok := s.visibility.SA.TryGet(s.mockContext)
+	s.False(ok)
+	s.Nil(s.visibility.CustomSearchAttributes(s.mockContext))
 }
 
-func (s *visibilitySuite) TestMemo() {
-	memo := s.visibility.GetMemo(s.mockMutableContext)
+func (s *visibilitySuite) TestReplaceCustomSearchAttributes() {
+	stringKey, stringVal := "stringKey", "stringValue"
+	intKey, intVal := "intKey", 42
+	floatKey, floatVal := "floatKey", 3.14
+	byteKey, byteVal := "byteKey", []byte{0x01, 0x02, 0x03}
+
+	// Set up some initial SA.
+	s.visibility.ReplaceCustomSearchAttributes(
+		s.mockMutableContext,
+		map[string]*commonpb.Payload{
+			stringKey: s.mustEncode(stringVal),
+			intKey:    s.mustEncode(intVal),
+			floatKey:  s.mustEncode(floatVal),
+		},
+	)
+	s.Len(s.mockMutableContext.Tasks, 1)
+	s.assertTaskPayload(2, s.mockMutableContext.Tasks[0].Payload)
+
+	sa := s.visibility.CustomSearchAttributes(s.mockMutableContext)
+	s.Len(sa, 3)
+
+	// Set to a new set of SA, non-existing keys should be removed.
+	s.visibility.ReplaceCustomSearchAttributes(
+		s.mockMutableContext,
+		map[string]*commonpb.Payload{
+			floatKey: s.mustEncode(floatVal),
+			byteKey:  s.mustEncode(byteVal),
+		},
+	)
+	s.Len(s.mockMutableContext.Tasks, 2)
+	s.assertTaskPayload(3, s.mockMutableContext.Tasks[1].Payload)
+
+	sa = s.visibility.CustomSearchAttributes(s.mockMutableContext)
+	s.Len(sa, 2)
+
+	// Setting to an empty map should remove the node.
+	s.visibility.ReplaceCustomSearchAttributes(
+		s.mockMutableContext,
+		map[string]*commonpb.Payload{},
+	)
+	s.Len(s.mockMutableContext.Tasks, 3)
+	s.assertTaskPayload(4, s.mockMutableContext.Tasks[2].Payload)
+	_, ok := s.visibility.SA.TryGet(s.mockContext)
+	s.False(ok)
+	s.Nil(s.visibility.CustomSearchAttributes(s.mockContext))
+}
+
+func (s *visibilitySuite) TestMergeCustomMemo() {
+	memo := s.visibility.CustomMemo(s.mockMutableContext)
 	s.Empty(memo)
 
 	stringKey, stringVal := "stringKey", "stringValue"
@@ -129,15 +187,15 @@ func (s *visibilitySuite) TestMemo() {
 	floatKey, floatVal := "floatKey", 3.14
 
 	// Add memo via Visibility struct method.
-	s.visibility.SetMemo(s.mockMutableContext, map[string]*commonpb.Payload{
+	s.visibility.MergeCustomMemo(s.mockMutableContext, map[string]*commonpb.Payload{
 		stringKey: s.mustEncode(stringVal),
 		intKey:    s.mustEncode(intVal),
 		floatKey:  s.mustEncode(floatVal),
 	})
 	s.Len(s.mockMutableContext.Tasks, 1)
-	protorequire.ProtoEqual(s.T(), &persistencespb.ChasmVisibilityTaskData{TransitionCount: 2}, s.mockMutableContext.Tasks[0].Payload.(*persistencespb.ChasmVisibilityTaskData))
+	s.assertTaskPayload(2, s.mockMutableContext.Tasks[0].Payload)
 
-	memo = s.visibility.GetMemo(s.mockMutableContext)
+	memo = s.visibility.CustomMemo(s.mockMutableContext)
 	s.Len(memo, 3)
 
 	var actualStringVal string
@@ -156,15 +214,81 @@ func (s *visibilitySuite) TestMemo() {
 	s.Equal(floatVal, actualFloatVal)
 
 	// Test remove memo by setting payload to nil.
-	s.visibility.SetMemo(s.mockMutableContext, map[string]*commonpb.Payload{
+	s.visibility.MergeCustomMemo(s.mockMutableContext, map[string]*commonpb.Payload{
 		intKey:   s.mustEncode(intVal),
 		floatKey: nil,
 	})
 	s.Len(s.mockMutableContext.Tasks, 2)
-	protorequire.ProtoEqual(s.T(), &persistencespb.ChasmVisibilityTaskData{TransitionCount: 3}, s.mockMutableContext.Tasks[1].Payload.(*persistencespb.ChasmVisibilityTaskData))
+	s.assertTaskPayload(3, s.mockMutableContext.Tasks[1].Payload)
 
-	memo = s.visibility.GetMemo(s.mockMutableContext)
+	memo = s.visibility.CustomMemo(s.mockMutableContext)
 	s.Len(memo, 2, "intKey and stringKey should remain")
+
+	// Test removing all memo fields also removes the node.
+	s.visibility.MergeCustomMemo(s.mockMutableContext, map[string]*commonpb.Payload{
+		stringKey: nil,
+		intKey:    nil,
+	})
+	s.Len(s.mockMutableContext.Tasks, 3)
+	s.assertTaskPayload(4, s.mockMutableContext.Tasks[2].Payload)
+	_, ok := s.visibility.Memo.TryGet(s.mockContext)
+	s.False(ok)
+	s.Nil(s.visibility.CustomMemo(s.mockContext))
+}
+
+func (s *visibilitySuite) TestReplaceCustomMemo() {
+	stringKey, stringVal := "stringKey", "stringValue"
+	intKey, intVal := "intKey", 42
+	floatKey, floatVal := "floatKey", 3.14
+	byteKey, byteVal := "byteKey", []byte{0x01, 0x02, 0x03}
+
+	// Set up some initial memo fields.
+	s.visibility.ReplaceCustomMemo(
+		s.mockMutableContext,
+		map[string]*commonpb.Payload{
+			stringKey: s.mustEncode(stringVal),
+			intKey:    s.mustEncode(intVal),
+			floatKey:  s.mustEncode(floatVal),
+		},
+	)
+	s.Len(s.mockMutableContext.Tasks, 1)
+	s.assertTaskPayload(2, s.mockMutableContext.Tasks[0].Payload)
+
+	memo := s.visibility.CustomMemo(s.mockMutableContext)
+	s.Len(memo, 3)
+
+	// Set to a new set of memo fields, non-existing keys should be removed.
+	s.visibility.ReplaceCustomMemo(
+		s.mockMutableContext,
+		map[string]*commonpb.Payload{
+			floatKey: s.mustEncode(floatVal),
+			byteKey:  s.mustEncode(byteVal),
+		},
+	)
+	s.Len(s.mockMutableContext.Tasks, 2)
+	s.assertTaskPayload(3, s.mockMutableContext.Tasks[1].Payload)
+
+	memo = s.visibility.CustomMemo(s.mockMutableContext)
+	s.Len(memo, 2)
+
+	// Setting to an empty map should remove the node.
+	s.visibility.ReplaceCustomMemo(
+		s.mockMutableContext,
+		map[string]*commonpb.Payload{},
+	)
+	s.Len(s.mockMutableContext.Tasks, 3)
+	s.assertTaskPayload(4, s.mockMutableContext.Tasks[2].Payload)
+	_, ok := s.visibility.Memo.TryGet(s.mockContext)
+	s.False(ok)
+	s.Nil(s.visibility.CustomMemo(s.mockContext))
+}
+
+func (s *visibilitySuite) assertTaskPayload(expectedCount int64, taskPayload any) {
+	protorequire.ProtoEqual(
+		s.T(),
+		&persistencespb.ChasmVisibilityTaskData{TransitionCount: expectedCount},
+		taskPayload.(*persistencespb.ChasmVisibilityTaskData),
+	)
 }
 
 func (s *visibilitySuite) mustEncode(v any) *commonpb.Payload {

service/history/visibility_queue_task_executor.go

@@ -405,7 +405,7 @@ func (t *visibilityQueueTaskExecutor) processChasmTask(
 
 	searchattributes := make(map[string]*commonpb.Payload)
 
-	aliasedCustomSearchAttributes := visComponent.GetSearchAttributes(visTaskContext)
+	aliasedCustomSearchAttributes := visComponent.CustomSearchAttributes(visTaskContext)
 	for alias, value := range aliasedCustomSearchAttributes {
 		fieldName, err := customSaMapper.GetFieldName(alias, namespaceEntry.Name().String())
 		if err != nil {
@@ -428,7 +428,7 @@ func (t *visibilityQueueTaskExecutor) processChasmTask(
 	}
 
 	combinedMemo := make(map[string]*commonpb.Payload, 2)
-	userMemoMap := visComponent.GetMemo(visTaskContext)
+	userMemoMap := visComponent.CustomMemo(visTaskContext)
 	if len(userMemoMap) > 0 {
 		userMemoProto := &commonpb.Memo{Fields: userMemoMap}
 		userMemoPayload, err := payload.Encode(userMemoProto)