Add CHASM Visibility registration and task processing (#8491)

## What changed?
Adds CHASM search attribute provider and mapper implementation. To use
CHASM search attributes, embed your component with
ComponentSearchAttributesProvider and register the search attributes to
the CHASM Registry.

To define a component search attribute, define a SearchAttribute as a
package/global scoped variable (CHASM search attributes should reference
exported variables like `SearchAttributeFieldInt01`,
`SearchAttributeFieldDateTime01`:
```
var (
	TestComponentStartTimeSearchAttribute = NewSearchAttributeTime(testComponentStartTimeSAKey, SearchAttributeFieldDateTime01)

	_ VisibilitySearchAttributesProvider = (*TestComponent)(nil)
	_ VisibilityMemoProvider             = (*TestComponent)(nil)
	_ VisibilitySearchAttributesMapper   = (*TestComponent)(nil)
)
```

To update SearchAttributes during an execution of a CHASM archetype,
implement the provider `SearchAttributes` in the root component, and
call the `NewValue` method on a SearchAttribute within the Provider
method to generate a new SearchAttributeValue. On SearchAttribute
updates during CHASM transactions, the framework will detect changes and
submit a Visibility task for persistence.
```
type TestComponent struct {
	UnimplementedComponent

	...
	Visibility Field[*Visibility]
}


func (tc *TestComponent) SearchAttributes(ctx Context) []SearchAttributeValue {
	return []SearchAttributeValue{
             TestComponentStartTimeSearchAttribute.NewValue(time)
        }
}
```

To register the search attribute with the CHASM Registry,
`WithSearchAttributes` must be passed as a RegistrableComponentOption to
the root component of the library. This is required to support
Visibility queries.

```
return []*chasm.RegistrableComponent{
		chasm.NewRegistrableComponent[*Scheduler]("scheduler", WIthSearchAttributes([]SearchAttributeDefinition{TestComponentStartTimeSearchAttribute}),
		chasm.NewRegistrableComponent[*Generator]("generator"),
		chasm.NewRegistrableComponent[*Invoker]("invoker"),
		chasm.NewRegistrableComponent[*Backfiller]("backfiller"),
	}
```



## Why?
Required to support CHASM Search Attributes. 

## 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)

## Potential risks
TBD

---------

Co-authored-by: Rodrigo Zhou <[email protected]>

Author: awln-temporal

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.