feat(shard-distributor): add shard handover stats (#7495)

<!-- Describe what has changed in this PR -->
**What changed?**
* ShardHandoverStats have been added to AssignedState
* Add comments to Store's methods and State structures


<!-- Tell your future self why have you made these changes -->
**Why?**
* ShardHandoverStats are going to be used to track Shard Assignment
Distribution Latency and Shard Handover Latency which are going to be
used as health indicators of a cluster


<!-- How have you verified this change? Tested locally? Added a unit
test? Checked in staging env? -->
**How did you test it?**
* Unit tests


<!-- Assuming the worst case, what can be broken when deploying this
change to production? -->
**Potential risks**

<!-- Is it notable for release? e.g. schema updates, configuration or
data migration required? If so, please mention it, and also update
CHANGELOG.md -->
**Release notes**

<!-- Is there any documentation updates should be made for config,
https://cadenceworkflow.io/docs/operation-guide/setup/ ? If so, please
open an PR in https://github.com/cadence-workflow/cadence-docs -->
**Documentation Changes**

Author: arzonus

common/types/sharddistributor.go

@@ -22,9 +22,11 @@
 
 package types
 
-import "fmt"
+import (
+	"fmt"
+)
 
-//go:generate enumer -type=ExecutorStatus,ShardStatus,AssignmentStatus,MigrationMode -json -output sharddistributor_statuses_enumer_generated.go
+//go:generate enumer -type=ExecutorStatus,ShardStatus,AssignmentStatus,MigrationMode,HandoverType -json -output sharddistributor_statuses_enumer_generated.go
 
 type GetShardOwnerRequest struct {
 	ShardKey  string
@@ -198,6 +200,7 @@ func (v *ExecutorHeartbeatResponse) GetMigrationPhase() (o MigrationMode) {
 }
 
 type ShardAssignment struct {
+	// Status indicates the current assignment status of the shard.
 	Status AssignmentStatus `json:"status"`
 }
 
@@ -217,6 +220,30 @@ const (
 	AssignmentStatusREADY   AssignmentStatus = 1
 )
 
+// HandoverType is used to indicate the type of handover that occurred during shard reassignment.
+// Type is persisted to the DB with a string value mapping.
+// Beware - if we want to change the name - it should be backward compatible and should be done in two steps.
+type HandoverType int32
+
+const (
+	HandoverTypeINVALID HandoverType = 0
+
+	// HandoverTypeGRACEFUL
+	// Graceful handover indicates that the shard was transferred in a way that allowed
+	// the previous owner had a chance to finish processing before the shard was reassigned.
+	HandoverTypeGRACEFUL HandoverType = 1
+
+	// HandoverTypeEMERGENCY
+	// Emergency handover indicates that the shard was transferred abruptly without
+	// allowing the previous owner to finish processing.
+	HandoverTypeEMERGENCY HandoverType = 2
+)
+
+// Ptr returns a pointer to the HandoverType value.
+func (t HandoverType) Ptr() *HandoverType {
+	return &t
+}
+
 type MigrationMode int32
 
 const (

common/types/sharddistributor_statuses_enumer_generated.go

@@ -2,6 +2,7 @@ package process
 
 import (
 	"context"
+	"errors"
 	"fmt"
 	"maps"
 	"math/rand"
@@ -337,7 +338,6 @@ func (p *namespaceProcessor) rebalanceShards(ctx context.Context) (err error) {
 }
 
 func (p *namespaceProcessor) rebalanceShardsImpl(ctx context.Context, metricsLoopScope metrics.Scope) (err error) {
-
 	namespaceState, err := p.shardStore.GetState(ctx, p.namespaceCfg.Name)
 	if err != nil {
 		return fmt.Errorf("get state: %w", err)
@@ -378,8 +378,8 @@ func (p *namespaceProcessor) rebalanceShardsImpl(ctx context.Context, metricsLoo
 	}
 
 	p.addAssignmentsToNamespaceState(namespaceState, currentAssignments)
-
 	p.logger.Info("Applying new shard distribution.")
+
 	// Use the leader guard for the assign and delete operation.
 	err = p.shardStore.AssignShards(ctx, p.namespaceCfg.Name, store.AssignShardsRequest{
 		NewState:          namespaceState,
@@ -473,27 +473,105 @@ func (*namespaceProcessor) updateAssignments(shardsToReassign []string, activeEx
 }
 
 func (p *namespaceProcessor) addAssignmentsToNamespaceState(namespaceState *store.NamespaceState, currentAssignments map[string][]string) {
-	newState := make(map[string]store.AssignedState)
+	newState := make(map[string]store.AssignedState, len(currentAssignments))
+
 	for executorID, shards := range currentAssignments {
 		assignedShardsMap := make(map[string]*types.ShardAssignment)
+
 		for _, shardID := range shards {
 			assignedShardsMap[shardID] = &types.ShardAssignment{Status: types.AssignmentStatusREADY}
 		}
+
 		modRevision := int64(0) // Should be 0 if we have not seen it yet
 		if namespaceAssignments, ok := namespaceState.ShardAssignments[executorID]; ok {
 			modRevision = namespaceAssignments.ModRevision
 		}
 
 		newState[executorID] = store.AssignedState{
-			AssignedShards: assignedShardsMap,
-			LastUpdated:    p.timeSource.Now().UTC(),
-			ModRevision:    modRevision,
+			AssignedShards:     assignedShardsMap,
+			LastUpdated:        p.timeSource.Now().UTC(),
+			ModRevision:        modRevision,
+			ShardHandoverStats: p.addHandoverStatsToExecutorAssignedState(namespaceState, executorID, shards),
 		}
 	}
 
 	namespaceState.ShardAssignments = newState
 }
 
+func (p *namespaceProcessor) addHandoverStatsToExecutorAssignedState(
+	namespaceState *store.NamespaceState,
+	executorID string, shardIDs []string,
+) map[string]store.ShardHandoverStats {
+	var newStats = make(map[string]store.ShardHandoverStats)
+
+	// Prepare handover stats for each shard
+	for _, shardID := range shardIDs {
+		handoverStats := p.newHandoverStats(namespaceState, shardID, executorID)
+
+		// If there is no handover (first assignment), we skip adding handover stats
+		if handoverStats != nil {
+			newStats[shardID] = *handoverStats
+		}
+	}
+
+	return newStats
+}
+
+// newHandoverStats creates shard handover statistics if a handover occurred.
+func (p *namespaceProcessor) newHandoverStats(
+	namespaceState *store.NamespaceState,
+	shardID string,
+	newExecutorID string,
+) *store.ShardHandoverStats {
+	logger := p.logger.WithTags(
+		tag.ShardNamespace(p.namespaceCfg.Name),
+		tag.ShardKey(shardID),
+		tag.ShardExecutor(newExecutorID),
+	)
+
+	// Fetch previous shard owners from cache
+	prevExecutor, err := p.shardStore.GetShardOwner(context.Background(), p.namespaceCfg.Name, shardID)
+	if err != nil && !errors.Is(err, store.ErrShardNotFound) {
+		logger.Warn("failed to get shard owner for shard statistic", tag.Error(err))
+		return nil
+	}
+	// previous executor is not found in cache
+	// meaning this is the first assignment of the shard
+	// so we skip updating handover stats
+	if prevExecutor == nil {
+		return nil
+	}
+
+	// No change in assignment
+	// meaning no handover occurred
+	// skip updating handover stats
+	if prevExecutor.ExecutorID == newExecutorID {
+		return nil
+	}
+
+	// previous executor heartbeat is not found in namespace state
+	// meaning the executor has already been cleaned up
+	// skip updating handover stats
+	prevExecutorHeartbeat, ok := namespaceState.Executors[prevExecutor.ExecutorID]
+	if !ok {
+		logger.Info("previous executor heartbeat not found, skipping handover stats")
+		return nil
+	}
+
+	handoverType := types.HandoverTypeEMERGENCY
+
+	// Consider it a graceful handover if the previous executor was in DRAINING or DRAINED status
+	// otherwise, it's an emergency handover
+	if prevExecutorHeartbeat.Status == types.ExecutorStatusDRAINING || prevExecutorHeartbeat.Status == types.ExecutorStatusDRAINED {
+		handoverType = types.HandoverTypeGRACEFUL
+	}
+
+	return &store.ShardHandoverStats{
+		HandoverType:                      handoverType,
+		PreviousExecutorLastHeartbeatTime: prevExecutorHeartbeat.LastHeartbeat,
+	}
+}
+
 func (*namespaceProcessor) getActiveExecutors(namespaceState *store.NamespaceState, staleExecutors map[string]int64) []string {
 	var activeExecutors []string
 	for id, state := range namespaceState.Executors {