cre-1601: consistent hashing and plugin test

Author: mchain0

pkg/workflows/shardorchestrator/plugin.go

@@ -44,8 +44,8 @@ func NewPlugin(store *Store, config ocr3types.ReportingPluginConfig, lggr logger
 	if cfg == nil {
 		cfg = &ConsensusConfig{
 			MinShardCount: 1,
-			MaxShardCount: 10,
-			BatchSize:     100,
+			MaxShardCount: 100,
+			BatchSize:     1000,
 		}
 	}
 
@@ -101,8 +101,8 @@ func (p *Plugin) ObservationQuorum(_ context.Context, _ ocr3types.OutcomeContext
 
 func (p *Plugin) Outcome(_ context.Context, outctx ocr3types.OutcomeContext, _ types.Query, aos []types.AttributedObservation) (ocr3types.Outcome, error) {
 	prevOutcome := &pb.Outcome{}
-	if err := proto.Unmarshal(outctx.PreviousOutcome, prevOutcome); err != nil {
-		p.lggr.Warnf("failed to unmarshal previous outcome: %v", err)
+	if err := proto.Unmarshal(outctx.PreviousOutcome, prevOutcome); err != nil || prevOutcome.State == nil {
+		p.lggr.Warnf("failed to unmarshal previous outcome or state is nil")
 		prevOutcome = &pb.Outcome{
 			State: &pb.RoutingState{
 				Id: 0,
@@ -116,8 +116,9 @@ func (p *Plugin) Outcome(_ context.Context, outctx ocr3types.OutcomeContext, _ t
 
 	currentShardHealth := make(map[uint32]int)
 	totalObservations := len(aos)
+	allWorkflows := []string{}
 
-	// Collect shard health observations
+	// Collect shard health observations and workflows
 	for _, ao := range aos {
 		observation := &pb.Observation{}
 		if err := proto.Unmarshal(ao.Observation, observation); err != nil {
@@ -130,13 +131,29 @@ func (p *Plugin) Outcome(_ context.Context, outctx ocr3types.OutcomeContext, _ t
 				currentShardHealth[shardID]++
 			}
 		}
+
+		// Collect workflow IDs
+		allWorkflows = append(allWorkflows, observation.Hashes...)
 	}
 
+	// Deduplicate workflows
+	workflowMap := make(map[string]bool)
+	for _, wf := range allWorkflows {
+		workflowMap[wf] = true
+	}
+	allWorkflows = make([]string, 0, len(workflowMap))
+	for wf := range workflowMap {
+		allWorkflows = append(allWorkflows, wf)
+	}
+	slices.Sort(allWorkflows) // Ensure deterministic order
+
 	// Determine desired shard count based on observations
 	healthyShardCount := uint32(0)
-	for _, count := range currentShardHealth {
+	for shardID, count := range currentShardHealth {
 		if count > int(p.config.F) {
 			healthyShardCount++
+			// Update store with healthy shard
+			p.store.SetShardHealth(shardID, true)
 		}
 	}
 
@@ -148,6 +165,15 @@ func (p *Plugin) Outcome(_ context.Context, outctx ocr3types.OutcomeContext, _ t
 		healthyShardCount = p.maxShardCount
 	}
 
+	// Use deterministic hashing to assign workflows to shards
+	routes := make(map[string]*pb.WorkflowRoute)
+	for _, wfID := range allWorkflows {
+		assignedShard := p.store.GetShardForWorkflow(wfID)
+		routes[wfID] = &pb.WorkflowRoute{
+			Shard: assignedShard,
+		}
+	}
+
 	// Update routing state
 	outcome := &pb.Outcome{
 		State: &pb.RoutingState{
@@ -156,17 +182,10 @@ func (p *Plugin) Outcome(_ context.Context, outctx ocr3types.OutcomeContext, _ t
 				RoutableShards: healthyShardCount,
 			},
 		},
-		Routes: make(map[string]*pb.WorkflowRoute),
+		Routes: routes,
 	}
 
-	p.lggr.Infow("Consensus Outcome", "healthyShards", healthyShardCount, "totalObservations", totalObservations)
-
-	// Copy previous routes
-	if prevOutcome.Routes != nil {
-		for wfID, route := range prevOutcome.Routes {
-			outcome.Routes[wfID] = route
-		}
-	}
+	p.lggr.Infow("Consensus Outcome", "healthyShards", healthyShardCount, "totalObservations", totalObservations, "workflowCount", len(routes))
 
 	return proto.MarshalOptions{Deterministic: true}.Marshal(outcome)
 }

pkg/workflows/shardorchestrator/plugin_test.go

@@ -0,0 +1,117 @@
+package shardorchestrator
+
+import (
+	"testing"
+
+	"github.com/smartcontractkit/libocr/commontypes"
+	"github.com/stretchr/testify/require"
+	"google.golang.org/protobuf/proto"
+
+	"github.com/smartcontractkit/chainlink-common/pkg/logger"
+	"github.com/smartcontractkit/chainlink-common/pkg/workflows/shardorchestrator/pb"
+	"github.com/smartcontractkit/libocr/offchainreporting2/types"
+	"github.com/smartcontractkit/libocr/offchainreporting2plus/ocr3types"
+)
+
+func TestPlugin_OutcomeWithMultiNodeObservations(t *testing.T) {
+	lggr := logger.Test(t)
+	store := NewStore()
+	store.SetAllShardHealth(map[uint32]bool{0: true, 1: true, 2: true})
+
+	config := ocr3types.ReportingPluginConfig{
+		N: 4, F: 1,
+		OffchainConfig:                          []byte{},
+		MaxDurationObservation:                  0,
+		MaxDurationShouldAcceptAttestedReport:   0,
+		MaxDurationShouldTransmitAcceptedReport: 0,
+	}
+
+	plugin, err := NewPlugin(store, config, lggr, nil)
+	require.NoError(t, err)
+
+	ctx := t.Context()
+	outcomeCtx := ocr3types.OutcomeContext{PreviousOutcome: []byte("")}
+
+	// Observations from 4 NOPs reporting health and workflows
+	observations := []struct {
+		name        string
+		shardHealth map[uint32]bool
+		workflows   []string
+	}{
+		{
+			name:        "NOP 0",
+			shardHealth: map[uint32]bool{0: true, 1: true, 2: true},
+			workflows:   []string{"wf-A", "wf-B", "wf-C"},
+		},
+		{
+			name:        "NOP 1",
+			shardHealth: map[uint32]bool{0: true, 1: true, 2: true},
+			workflows:   []string{"wf-B", "wf-C", "wf-D"},
+		},
+		{
+			name:        "NOP 2",
+			shardHealth: map[uint32]bool{0: true, 1: true, 2: false}, // shard 2 unhealthy
+			workflows:   []string{"wf-A", "wf-C"},
+		},
+		{
+			name:        "NOP 3",
+			shardHealth: map[uint32]bool{0: true, 1: true, 2: true},
+			workflows:   []string{"wf-A", "wf-B", "wf-D"},
+		},
+	}
+
+	// Build attributed observations
+	aos := make([]types.AttributedObservation, 0)
+	for _, obs := range observations {
+		pbObs := &pb.Observation{
+			Status: obs.shardHealth,
+			Hashes: obs.workflows,
+		}
+		rawObs, err := proto.Marshal(pbObs)
+		require.NoError(t, err)
+
+		aos = append(aos, types.AttributedObservation{
+			Observation: rawObs,
+			Observer:    commontypes.OracleID(len(aos)),
+		})
+	}
+
+	// Execute Outcome phase
+	outcome, err := plugin.Outcome(ctx, outcomeCtx, nil, aos)
+	require.NoError(t, err)
+	require.NotNil(t, outcome)
+
+	// Verify outcome
+	outcomeProto := &pb.Outcome{}
+	err = proto.Unmarshal(outcome, outcomeProto)
+	require.NoError(t, err)
+
+	// Check consensus results
+	t.Logf("Outcome - ID: %d, HealthyShards: %d", outcomeProto.State.Id, outcomeProto.State.GetRoutableShards())
+	t.Logf("Workflows assigned: %d", len(outcomeProto.Routes))
+
+	// Verify all workflows are assigned
+	expectedWorkflows := map[string]bool{"wf-A": true, "wf-B": true, "wf-C": true, "wf-D": true}
+	require.Equal(t, len(expectedWorkflows), len(outcomeProto.Routes))
+	for wf := range expectedWorkflows {
+		route, exists := outcomeProto.Routes[wf]
+		require.True(t, exists, "workflow %s should be assigned", wf)
+		require.True(t, route.Shard <= 2, "shard should be healthy (0-2)")
+		t.Logf("  %s → shard %d", wf, route.Shard)
+	}
+
+	// Verify determinism: run again, should get same assignments
+	outcome2, err := plugin.Outcome(ctx, outcomeCtx, nil, aos)
+	require.NoError(t, err)
+
+	outcomeProto2 := &pb.Outcome{}
+	err = proto.Unmarshal(outcome2, outcomeProto2)
+	require.NoError(t, err)
+
+	// Same workflows → same shards
+	for wf, route1 := range outcomeProto.Routes {
+		route2, exists := outcomeProto2.Routes[wf]
+		require.True(t, exists)
+		require.Equal(t, route1.Shard, route2.Shard, "workflow %s should assign to same shard", wf)
+	}
+}

pkg/workflows/shardorchestrator/store.go

@@ -1,33 +1,110 @@
 package shardorchestrator
 
 import (
-	"maps"
+	"slices"
+	"strconv"
 	"sync"
-	"time"
+
+	"github.com/buraksezer/consistent"
+	"github.com/cespare/xxhash/v2"
 )
 
-var DefaultRequestTimeout = 20 * time.Minute
+// xxhashHasher implements the consistent.Hasher interface using xxhash
+type xxhashHasher struct{}
+
+func (h xxhashHasher) Sum64(data []byte) uint64 {
+	return xxhash.Sum64(data)
+}
+
+// ShardMember implements consistent.Member for shard IDs
+type ShardMember string
+
+func (m ShardMember) String() string {
+	return string(m)
+}
 
 // Store manages shard routing state and workflow mappings
 type Store struct {
-	routingState map[string]uint32 // workflow_id -> shard_id
-	shardHealth  map[uint32]bool   // shard_id -> is_healthy
-	mu           sync.Mutex
+	routingState   map[string]uint32      // workflow_id -> shard_id
+	shardHealth    map[uint32]bool        // shard_id -> is_healthy
+	consistentHash *consistent.Consistent // Consistent hash ring for routing
+	healthyShards  []uint32               // Sorted list of healthy shards
+	mu             sync.Mutex
 }
 
 func NewStore() *Store {
 	return &Store{
-		routingState: make(map[string]uint32),
-		shardHealth:  make(map[uint32]bool),
-		mu:           sync.Mutex{},
+		routingState:  make(map[string]uint32),
+		shardHealth:   make(map[uint32]bool),
+		healthyShards: make([]uint32, 0),
+		mu:            sync.Mutex{},
+	}
+}
+
+// consistentHashConfig returns the configuration for consistent hashing
+// Matches prototype: PartitionCount=997 (prime), ReplicationFactor=50, Load=1.1
+func consistentHashConfig() consistent.Config {
+	return consistent.Config{
+		PartitionCount:    997, // Prime number for better distribution
+		ReplicationFactor: 50,  // Number of replicas per node
+		Load:              1.1, // Load factor for bounded loads
+		Hasher:            xxhashHasher{},
+	}
+}
+
+// updateConsistentHash rebuilds the consistent hash ring based on healthy shards
+func (s *Store) updateConsistentHash() {
+	s.mu.Lock()
+	defer s.mu.Unlock()
+
+	// Get list of healthy shards and create members
+	healthyMembers := make([]consistent.Member, 0)
+	s.healthyShards = make([]uint32, 0)
+
+	for shardID, healthy := range s.shardHealth {
+		if healthy {
+			healthyMembers = append(healthyMembers, ShardMember(strconv.FormatUint(uint64(shardID), 10)))
+			s.healthyShards = append(s.healthyShards, shardID)
+		}
+	}
+
+	// Sort for determinism
+	slices.Sort(s.healthyShards)
+
+	// If no healthy shards, add shard 0 as fallback
+	if len(healthyMembers) == 0 {
+		healthyMembers = append(healthyMembers, ShardMember("0"))
+		s.healthyShards = []uint32{0}
 	}
+
+	// Create consistent hash ring
+	s.consistentHash = consistent.New(healthyMembers, consistentHashConfig())
 }
 
-func (s *Store) GetShardForWorkflow(workflowID string) (uint32, bool) {
+// GetShardForWorkflow deterministically assigns a workflow to a shard using consistent hashing.
+// The assignment uses the same algorithm as the prototype: xxhash + consistent hashing ring.
+func (s *Store) GetShardForWorkflow(workflowID string) uint32 {
 	s.mu.Lock()
 	defer s.mu.Unlock()
-	shardID, ok := s.routingState[workflowID]
-	return shardID, ok
+
+	if s.consistentHash == nil {
+		// Fallback if hash ring not initialized
+		return 0
+	}
+
+	// Use consistent hashing to find the member for this workflow
+	member := s.consistentHash.LocateKey([]byte(workflowID))
+	if member == nil {
+		return 0
+	}
+
+	// Parse shard ID from member name
+	shardID, err := strconv.ParseUint(member.String(), 10, 32)
+	if err != nil {
+		return 0
+	}
+
+	return uint32(shardID)
 }
 
 func (s *Store) SetShardForWorkflow(workflowID string, shardID uint32) {
@@ -40,28 +117,40 @@ func (s *Store) GetShardHealth() map[uint32]bool {
 	s.mu.Lock()
 	defer s.mu.Unlock()
 	copied := make(map[uint32]bool)
-	maps.Copy(copied, s.shardHealth)
+	for k, v := range s.shardHealth {
+		copied[k] = v
+	}
 	return copied
 }
 
 func (s *Store) SetShardHealth(shardID uint32, healthy bool) {
 	s.mu.Lock()
-	defer s.mu.Unlock()
 	s.shardHealth[shardID] = healthy
+	s.mu.Unlock()
+
+	// Rebuild consistent hash ring when shard health changes
+	s.updateConsistentHash()
 }
 
 func (s *Store) SetAllShardHealth(health map[uint32]bool) {
 	s.mu.Lock()
-	defer s.mu.Unlock()
 	s.shardHealth = make(map[uint32]bool)
-	maps.Copy(s.shardHealth, health)
+	for k, v := range health {
+		s.shardHealth[k] = v
+	}
+	s.mu.Unlock()
+
+	// Rebuild consistent hash ring
+	s.updateConsistentHash()
 }
 
 func (s *Store) GetAllRoutingState() map[string]uint32 {
 	s.mu.Lock()
 	defer s.mu.Unlock()
 	copied := make(map[string]uint32)
-	maps.Copy(copied, s.routingState)
+	for k, v := range s.routingState {
+		copied[k] = v
+	}
 	return copied
 }
 
@@ -82,3 +171,10 @@ func (s *Store) GetHealthyShardCount() int {
 	}
 	return count
 }
+
+// GetHealthyShards returns a sorted list of healthy shards for inspection
+func (s *Store) GetHealthyShards() []uint32 {
+	s.mu.Lock()
+	defer s.mu.Unlock()
+	return slices.Clone(s.healthyShards)
+}