Add CHASM Visibility registration and task processing

chasm/lib/tests/library.go

@@ -16,7 +16,13 @@ func (l *library) Name() string {
 
 func (l *library) Components() []*chasm.RegistrableComponent {
 	return []*chasm.RegistrableComponent{
-		chasm.NewRegistrableComponent[*PayloadStore]("payloadStore"),
+		chasm.NewRegistrableComponent[*PayloadStore]("payloadStore",
+			chasm.WithSearchAttributes(
+				PayloadTotalCountSearchAttribute,
+				PayloadTotalSizeSearchAttribute,
+				chasm.SearchAttributeTemporalScheduledByID,
+			),
+		),
 	}
 }
 

chasm/lib/tests/payload.go

@@ -6,7 +6,6 @@ import (
 	"go.temporal.io/server/chasm"
 	"go.temporal.io/server/chasm/lib/tests/gen/testspb/v1"
 	"go.temporal.io/server/common"
-	"go.temporal.io/server/common/searchattribute"
 	"google.golang.org/protobuf/types/known/timestamppb"
 )
 
@@ -15,15 +14,16 @@ const (
 	TotalSizeMemoFieldName  = "TotalSize"
 )
 
-// TODO: Register proper SA for TotalCount and TotalSize
-// For now, CHASM framework does NOT support Per-Component SearchAttributes
-// so just update a random existing pre-defined SA to make sure the logic works.
 const (
-	TestKeywordSAFieldName  = searchattribute.TemporalScheduledById
-	TestKeywordSAFieldValue = "test-keyword-value"
+	TestScheduleID           = "TestScheduleID"
+	PayloadTotalCountSAAlias = "PayloadTotalCount"
+	PayloadTotalSizeSAAlias  = "PayloadTotalSize"
 )
 
 var (
+	PayloadTotalCountSearchAttribute = chasm.NewSearchAttributeInt(PayloadTotalCountSAAlias, chasm.SearchAttributeFieldInt01)
+	PayloadTotalSizeSearchAttribute  = chasm.NewSearchAttributeInt(PayloadTotalSizeSAAlias, chasm.SearchAttributeFieldInt02)
+
 	_ chasm.VisibilitySearchAttributesProvider = (*PayloadStore)(nil)
 	_ chasm.VisibilityMemoProvider             = (*PayloadStore)(nil)
 )
@@ -147,12 +147,12 @@ func (s *PayloadStore) LifecycleState(
 
 // SearchAttributes implements chasm.VisibilitySearchAttributesProvider interface
 func (s *PayloadStore) SearchAttributes(
-	_ chasm.Context,
-) map[string]chasm.VisibilityValue {
-	// TODO: UpsertSearchAttribute as well when CHASM framework supports Per-Component SearchAttributes
-	// For now, we just update a random existing pre-defined SA to make sure the logic works.
-	return map[string]chasm.VisibilityValue{
-		TestKeywordSAFieldName: chasm.VisibilityValueString(TestKeywordSAFieldValue),
+	ctx chasm.Context,
+) []chasm.SearchAttributeKeyValue {
+	return []chasm.SearchAttributeKeyValue{
+		PayloadTotalCountSearchAttribute.Value(s.State.TotalCount),
+		PayloadTotalSizeSearchAttribute.Value(s.State.TotalSize),
+		chasm.SearchAttributeTemporalScheduledByID.Value(TestScheduleID),
 	}
 }
 

chasm/registrable_component.go

@@ -1,7 +1,10 @@
 package chasm
 
 import (
+	"fmt"
 	"reflect"
+
+	enumspb "go.temporal.io/api/enums/v1"
 )
 
 type (
@@ -13,6 +16,8 @@ type (
 		ephemeral     bool
 		singleCluster bool
 		shardingFn    func(EntityKey) string
+
+		searchAttributesMapper *VisibilitySearchAttributesMapper
 	}
 
 	RegistrableComponentOption func(*RegistrableComponent)
@@ -56,6 +61,40 @@ func WithShardingFn(
 	}
 }
 
+func WithSearchAttributes(
+	searchAttributes ...SearchAttribute,
+) RegistrableComponentOption {
+	return func(rc *RegistrableComponent) {
+		if len(searchAttributes) == 0 {
+			return
+		}
+		rc.searchAttributesMapper = &VisibilitySearchAttributesMapper{
+			aliasToField: make(map[string]string, len(searchAttributes)),
+			fieldToAlias: make(map[string]string, len(searchAttributes)),
+			saTypeMap:    make(map[string]enumspb.IndexedValueType, len(searchAttributes)),
+		}
+
+		for _, sa := range searchAttributes {
+			alias := sa.definition().alias
+			field := sa.definition().field
+			valueType := sa.definition().valueType
+
+			if _, ok := rc.searchAttributesMapper.aliasToField[alias]; ok {
+				//nolint:forbidigo
+				panic(fmt.Sprintf("registrable component validation error: search attribute alias %q is already defined", alias))
+			}
+			if _, ok := rc.searchAttributesMapper.fieldToAlias[field]; ok {
+				//nolint:forbidigo
+				panic(fmt.Sprintf("registrable component validation error: search attribute field %q is already defined", field))
+			}
+
+			rc.searchAttributesMapper.aliasToField[alias] = field
+			rc.searchAttributesMapper.fieldToAlias[field] = alias
+			rc.searchAttributesMapper.saTypeMap[field] = valueType
+		}
+	}
+}
+
 // fqType returns the fully qualified name of the component, which is a combination of
 // the library name and the component type. This is used to uniquely identify
 // the component in the registry.