cre-1601: test extension to validate workflows to shards eassignments

Author: mchain0

pkg/workflows/ring/plugin.go

@@ -151,15 +151,31 @@ func (p *Plugin) collectShardInfo(aos []types.AttributedObservation) (shardHealt
 	return shardHealth, workflows, timestamps
 }
 
-func (p *Plugin) countHealthyShards(shardHealth map[uint32]int) uint32 {
-	var count uint32
+func (p *Plugin) getHealthyShards(shardHealth map[uint32]int) []uint32 {
+	var healthyShards []uint32
 	for shardID, votes := range shardHealth {
 		if votes > p.config.F {
-			count++
+			healthyShards = append(healthyShards, shardID)
 			p.store.SetShardHealth(shardID, true)
 		}
 	}
-	return max(p.minShardCount, min(count, p.maxShardCount))
+	slices.Sort(healthyShards)
+
+	// Apply min/max bounds on shard count
+	if uint32(len(healthyShards)) < p.minShardCount {
+		// Pad with sequential shard IDs if below minimum
+		for i := uint32(0); uint32(len(healthyShards)) < p.minShardCount; i++ {
+			if !slices.Contains(healthyShards, i) {
+				healthyShards = append(healthyShards, i)
+			}
+		}
+		slices.Sort(healthyShards)
+	} else if uint32(len(healthyShards)) > p.maxShardCount {
+		// Truncate to max (keep lowest shard IDs for determinism)
+		healthyShards = healthyShards[:p.maxShardCount]
+	}
+
+	return healthyShards
 }
 
 func (p *Plugin) Outcome(_ context.Context, outctx ocr3types.OutcomeContext, _ types.Query, aos []types.AttributedObservation) (ocr3types.Outcome, error) {
@@ -185,9 +201,9 @@ func (p *Plugin) Outcome(_ context.Context, outctx ocr3types.OutcomeContext, _ t
 
 	allWorkflows = uniqueSorted(allWorkflows)
 
-	healthyShardCount := p.countHealthyShards(currentShardHealth)
+	healthyShards := p.getHealthyShards(currentShardHealth)
 
-	nextState, err := p.calculateNextState(prior.State, healthyShardCount, now)
+	nextState, err := p.calculateNextState(prior.State, uint32(len(healthyShards)), now)
 	if err != nil {
 		return nil, err
 	}
@@ -196,7 +212,7 @@ func (p *Plugin) Outcome(_ context.Context, outctx ocr3types.OutcomeContext, _ t
 	// This must be a pure function of consensus-derived data to avoid protocol failures
 	routes := make(map[string]*pb.WorkflowRoute)
 	for _, wfID := range allWorkflows {
-		assignedShard := getShardForWorkflow(wfID, healthyShardCount)
+		assignedShard := getShardForWorkflow(wfID, healthyShards)
 		routes[wfID] = &pb.WorkflowRoute{
 			Shard: assignedShard,
 		}
@@ -207,7 +223,7 @@ func (p *Plugin) Outcome(_ context.Context, outctx ocr3types.OutcomeContext, _ t
 		Routes: routes,
 	}
 
-	p.lggr.Infow("Consensus Outcome", "healthyShards", healthyShardCount, "totalObservations", len(aos), "workflowCount", len(routes))
+	p.lggr.Infow("Consensus Outcome", "healthyShards", len(healthyShards), "totalObservations", len(aos), "workflowCount", len(routes))
 
 	return proto.MarshalOptions{Deterministic: true}.Marshal(outcome)
 }

pkg/workflows/ring/plugin_test.go

@@ -352,13 +352,13 @@ func makeObservations(t *testing.T, shardHealths []map[uint32]bool, workflows []
 	return aos
 }
 
-func TestPlugin_countHealthyShards(t *testing.T) {
+func TestPlugin_getHealthyShards(t *testing.T) {
 	tests := []struct {
 		name     string
 		min, max uint32
 		votes    map[uint32]int // shardID -> vote count
 		f        int
-		want     uint32
+		want     int
 	}{
 		{"below min", 3, 10, map[uint32]int{0: 2, 1: 2}, 1, 3},
 		{"above max", 1, 2, map[uint32]int{0: 2, 1: 2, 2: 2, 3: 2}, 1, 2},
@@ -374,8 +374,8 @@ func TestPlugin_countHealthyShards(t *testing.T) {
 				minShardCount: tc.min,
 				maxShardCount: tc.max,
 			}
-			got := plugin.countHealthyShards(tc.votes)
-			require.Equal(t, tc.want, got)
+			got := plugin.getHealthyShards(tc.votes)
+			require.Equal(t, tc.want, len(got))
 		})
 	}
 }

pkg/workflows/ring/store.go

@@ -73,9 +73,9 @@ func (s *Store) GetShardForWorkflow(ctx context.Context, workflowID string) (uin
 
 	// In steady state, compute locally using consistent hashing
 	if s.currentState == nil || s.isInSteadyState() {
-		shardCount := uint32(len(s.healthyShards))
+		healthyShards := slices.Clone(s.healthyShards)
 		s.mu.Unlock()
-		return getShardForWorkflow(workflowID, shardCount), nil
+		return getShardForWorkflow(workflowID, healthyShards), nil
 	}
 
 	// In transition state, enqueue request and wait for allocation

pkg/workflows/ring/store_test.go

@@ -9,7 +9,6 @@ import (
 	"github.com/stretchr/testify/require"
 )
 
-// TestStore_DeterministicHashing verifies that workflow assignments are deterministic
 func TestStore_DeterministicHashing(t *testing.T) {
 	store := NewStore()
 
@@ -35,7 +34,6 @@ func TestStore_DeterministicHashing(t *testing.T) {
 	require.True(t, shard1 >= 0 && shard1 <= 2, "Shard should be in healthy set")
 }
 
-// TestStore_ConsistentRingConsistency verifies that all nodes with same healthy shards agree
 func TestStore_ConsistentRingConsistency(t *testing.T) {
 	store1 := NewStore()
 	store2 := NewStore()
@@ -64,7 +62,6 @@ func TestStore_ConsistentRingConsistency(t *testing.T) {
 	}
 }
 
-// TestStore_Rebalancing verifies rebalancing when shard health changes
 func TestStore_Rebalancing(t *testing.T) {
 	store := NewStore()
 	ctx := context.Background()
@@ -93,9 +90,16 @@ func TestStore_Rebalancing(t *testing.T) {
 	healthyShards := store.GetHealthyShards()
 	require.Equal(t, 2, len(healthyShards), "Should have 2 healthy shards")
 	require.NotContains(t, healthyShards, uint32(1), "Shard 1 should not be healthy")
+
+	// Verify that workflows on healthy shards did not move
+	for wfID, originalShard := range assignments1 {
+		if originalShard == 0 || originalShard == 2 {
+			require.Equal(t, originalShard, assignments2[wfID],
+				"Workflow %s on healthy shard %d should not have moved", wfID, originalShard)
+		}
+	}
 }
 
-// TestStore_GetHealthyShards verifies that healthy shards list is correctly maintained
 func TestStore_GetHealthyShards(t *testing.T) {
 	store := NewStore()
 
@@ -111,7 +115,6 @@ func TestStore_GetHealthyShards(t *testing.T) {
 	require.Equal(t, []uint32{1, 2, 3}, healthyShards)
 }
 
-// TestStore_NilHashRingFallback verifies fallback when hash ring is uninitialized
 func TestStore_NilHashRingFallback(t *testing.T) {
 	store := NewStore()
 	ctx := context.Background()
@@ -122,7 +125,6 @@ func TestStore_NilHashRingFallback(t *testing.T) {
 	require.Equal(t, uint32(0), shard)
 }
 
-// TestStore_DistributionAcrossShards verifies that workflows are distributed across shards
 func TestStore_DistributionAcrossShards(t *testing.T) {
 	store := NewStore()
 	ctx := context.Background()
@@ -156,8 +158,6 @@ func sum(distribution map[uint32]int) int {
 	return total
 }
 
-// TestStore_PendingAllocsDuringTransition verifies that allocation requests block during transition
-// and are fulfilled when SetShardForWorkflow is called
 func TestStore_PendingAllocsDuringTransition(t *testing.T) {
 	store := NewStore()
 	store.SetAllShardHealth(map[uint32]bool{0: true, 1: true})

pkg/workflows/ring/utils.go

@@ -35,14 +35,14 @@ func consistentHashConfig() consistent.Config {
 	}
 }
 
-func getShardForWorkflow(workflowID string, shardCount uint32) uint32 {
-	if shardCount == 0 {
+func getShardForWorkflow(workflowID string, healthyShards []uint32) uint32 {
+	if len(healthyShards) == 0 {
 		return 0
 	}
 
-	members := make([]consistent.Member, shardCount)
-	for i := uint32(0); i < shardCount; i++ {
-		members[i] = ShardMember(strconv.FormatUint(uint64(i), 10))
+	members := make([]consistent.Member, len(healthyShards))
+	for i, shardID := range healthyShards {
+		members[i] = ShardMember(strconv.FormatUint(uint64(shardID), 10))
 	}
 
 	ring := consistent.New(members, consistentHashConfig())