cre-1601: store in two states, steady and transition; enque for allocation trigger post round;

Author: mchain0

pkg/workflows/ring/plugin.go

@@ -101,11 +101,17 @@ func (p *Plugin) Query(_ context.Context, _ ocr3types.OutcomeContext) (types.Que
 func (p *Plugin) Observation(_ context.Context, _ ocr3types.OutcomeContext, _ types.Query) (types.Observation, error) {
 	shardHealth := p.store.GetShardHealth()
 
+	// Collect workflow IDs from cache and pending allocation requests
 	allWorkflowIDs := make([]string, 0)
 	for wfID := range p.store.GetAllRoutingState() {
 		allWorkflowIDs = append(allWorkflowIDs, wfID)
 	}
-	slices.Sort(allWorkflowIDs)
+
+	// Include any pending allocation requests (workflows waiting for assignment during transition)
+	pendingAllocs := p.store.GetPendingAllocations()
+	allWorkflowIDs = append(allWorkflowIDs, pendingAllocs...)
+
+	allWorkflowIDs = uniqueSorted(allWorkflowIDs)
 
 	observation := &pb.Observation{
 		ShardHealthStatus: shardHealth,

pkg/workflows/ring/store.go

@@ -1,23 +1,39 @@
 package ring
 
 import (
+	"context"
 	"slices"
 	"sync"
+
+	"github.com/smartcontractkit/chainlink-common/pkg/workflows/ring/pb"
 )
 
+// AllocationRequest represents a pending workflow allocation request during transition
+type AllocationRequest struct {
+	WorkflowID string
+	Result     chan uint32
+}
+
 // Store manages shard routing state and workflow mappings
 type Store struct {
-	routingState  map[string]uint32 // workflow_id -> shard_id
+	routingState  map[string]uint32 // workflow_id -> shard_id (cache of allocated workflows)
 	shardHealth   map[uint32]bool   // shard_id -> is_healthy
 	healthyShards []uint32          // Sorted list of healthy shards
-	mu            sync.Mutex
+	currentState  *pb.RoutingState  // Current routing state (steady or transition)
+
+	pendingAllocs map[string][]chan uint32 // workflow_id -> waiting channels
+	allocRequests chan AllocationRequest   // Channel for new allocation requests
+
+	mu sync.Mutex
 }
 
 func NewStore() *Store {
 	return &Store{
 		routingState:  make(map[string]uint32),
 		shardHealth:   make(map[uint32]bool),
 		healthyShards: make([]uint32, 0),
+		pendingAllocs: make(map[string][]chan uint32),
+		allocRequests: make(chan AllocationRequest, 1000),
 		mu:            sync.Mutex{},
 	}
 }
@@ -44,19 +60,104 @@ func (s *Store) updateHealthyShards() {
 	}
 }
 
-// GetShardForWorkflow deterministically assigns a workflow to a shard using consistent hashing.
-func (s *Store) GetShardForWorkflow(workflowID string) uint32 {
+// GetShardForWorkflow returns the shard for a workflow.
+// In steady state, it uses consistent hashing from cache.
+// In transition state, it enqueues an allocation request and waits for OCR to process it.
+func (s *Store) GetShardForWorkflow(ctx context.Context, workflowID string) (uint32, error) {
 	s.mu.Lock()
-	shardCount := uint32(len(s.healthyShards))
+
+	// Check if already allocated in cache
+	if shard, ok := s.routingState[workflowID]; ok {
+		s.mu.Unlock()
+		return shard, nil
+	}
+
+	// In steady state, compute locally using consistent hashing
+	if s.currentState == nil || s.isInSteadyState() {
+		shardCount := uint32(len(s.healthyShards))
+		s.mu.Unlock()
+		return getShardForWorkflow(workflowID, shardCount), nil
+	}
+
+	// In transition state, enqueue request and wait for allocation
+	resultCh := make(chan uint32, 1)
+	s.pendingAllocs[workflowID] = append(s.pendingAllocs[workflowID], resultCh)
 	s.mu.Unlock()
 
-	return getShardForWorkflow(workflowID, shardCount)
+	// Submit allocation request
+	select {
+	case s.allocRequests <- AllocationRequest{WorkflowID: workflowID, Result: resultCh}:
+	case <-ctx.Done():
+		return 0, ctx.Err()
+	}
+
+	// Wait for result
+	select {
+	case shard := <-resultCh:
+		return shard, nil
+	case <-ctx.Done():
+		return 0, ctx.Err()
+	}
+}
+
+func (s *Store) isInSteadyState() bool {
+	if s.currentState == nil {
+		return true
+	}
+	_, ok := s.currentState.State.(*pb.RoutingState_RoutableShards)
+	return ok
 }
 
 func (s *Store) SetShardForWorkflow(workflowID string, shardID uint32) {
 	s.mu.Lock()
 	defer s.mu.Unlock()
+
 	s.routingState[workflowID] = shardID
+
+	// Signal any waiting allocation requests
+	if waiters, ok := s.pendingAllocs[workflowID]; ok {
+		for _, ch := range waiters {
+			select {
+			case ch <- shardID:
+			default:
+			}
+		}
+		delete(s.pendingAllocs, workflowID)
+	}
+}
+
+// SetRoutingState updates the current routing state (steady or transition)
+func (s *Store) SetRoutingState(state *pb.RoutingState) {
+	s.mu.Lock()
+	defer s.mu.Unlock()
+	s.currentState = state
+}
+
+// GetRoutingState returns the current routing state
+func (s *Store) GetRoutingState() *pb.RoutingState {
+	s.mu.Lock()
+	defer s.mu.Unlock()
+	return s.currentState
+}
+
+// GetPendingAllocations returns workflow IDs that need allocation (non-blocking)
+func (s *Store) GetPendingAllocations() []string {
+	var pending []string
+	for {
+		select {
+		case req := <-s.allocRequests:
+			pending = append(pending, req.WorkflowID)
+		default:
+			return pending
+		}
+	}
+}
+
+// IsInTransition returns true if the store is in transition state
+func (s *Store) IsInTransition() bool {
+	s.mu.Lock()
+	defer s.mu.Unlock()
+	return !s.isInSteadyState()
 }
 
 func (s *Store) GetShardHealth() map[uint32]bool {

pkg/workflows/ring/store_test.go

@@ -1,8 +1,11 @@
 package ring
 
 import (
+	"context"
 	"testing"
+	"time"
 
+	"github.com/smartcontractkit/chainlink-common/pkg/workflows/ring/pb"
 	"github.com/stretchr/testify/require"
 )
 
@@ -17,10 +20,15 @@ func TestStore_DeterministicHashing(t *testing.T) {
 		2: true,
 	})
 
+	ctx := context.Background()
+
 	// Test determinism: same workflow always gets same shard
-	shard1 := store.GetShardForWorkflow("workflow-123")
-	shard2 := store.GetShardForWorkflow("workflow-123")
-	shard3 := store.GetShardForWorkflow("workflow-123")
+	shard1, err := store.GetShardForWorkflow(ctx, "workflow-123")
+	require.NoError(t, err)
+	shard2, err := store.GetShardForWorkflow(ctx, "workflow-123")
+	require.NoError(t, err)
+	shard3, err := store.GetShardForWorkflow(ctx, "workflow-123")
+	require.NoError(t, err)
 
 	require.Equal(t, shard1, shard2, "Same workflow should get same shard (call 2)")
 	require.Equal(t, shard2, shard3, "Same workflow should get same shard (call 3)")
@@ -39,12 +47,17 @@ func TestStore_ConsistentRingConsistency(t *testing.T) {
 	store2.SetAllShardHealth(healthyShards)
 	store3.SetAllShardHealth(healthyShards)
 
+	ctx := context.Background()
+
 	// All compute same assignments
 	workflows := []string{"workflow-A", "workflow-B", "workflow-C", "workflow-D"}
 	for _, wf := range workflows {
-		s1 := store1.GetShardForWorkflow(wf)
-		s2 := store2.GetShardForWorkflow(wf)
-		s3 := store3.GetShardForWorkflow(wf)
+		s1, err := store1.GetShardForWorkflow(ctx, wf)
+		require.NoError(t, err)
+		s2, err := store2.GetShardForWorkflow(ctx, wf)
+		require.NoError(t, err)
+		s3, err := store3.GetShardForWorkflow(ctx, wf)
+		require.NoError(t, err)
 
 		require.Equal(t, s1, s2, "All nodes should agree on %s assignment", wf)
 		require.Equal(t, s2, s3, "All nodes should agree on %s assignment", wf)
@@ -54,21 +67,26 @@ 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()
 
 	// Start with 3 healthy shards
 	store.SetAllShardHealth(map[uint32]bool{0: true, 1: true, 2: true})
 	assignments1 := make(map[string]uint32)
 	for i := 1; i <= 10; i++ {
 		wfID := "workflow-" + string(rune(i))
-		assignments1[wfID] = store.GetShardForWorkflow(wfID)
+		shard, err := store.GetShardForWorkflow(ctx, wfID)
+		require.NoError(t, err)
+		assignments1[wfID] = shard
 	}
 
 	// Shard 1 fails
 	store.SetShardHealth(1, false)
 	assignments2 := make(map[string]uint32)
 	for i := 1; i <= 10; i++ {
 		wfID := "workflow-" + string(rune(i))
-		assignments2[wfID] = store.GetShardForWorkflow(wfID)
+		shard, err := store.GetShardForWorkflow(ctx, wfID)
+		require.NoError(t, err)
+		assignments2[wfID] = shard
 	}
 
 	// Check that rebalancing occurred (some workflows moved)
@@ -95,20 +113,19 @@ func TestStore_GetHealthyShards(t *testing.T) {
 
 // TestStore_NilHashRingFallback verifies fallback when hash ring is uninitialized
 func TestStore_NilHashRingFallback(t *testing.T) {
-	store := &Store{
-		routingState:  make(map[string]uint32),
-		shardHealth:   make(map[uint32]bool),
-		healthyShards: make([]uint32, 0),
-	}
+	store := NewStore()
+	ctx := context.Background()
 
-	// Should not panic, should return 0 as fallback
-	shard := store.GetShardForWorkflow("workflow-123")
+	// Should not panic, should return 0 as fallback (no healthy shards set)
+	shard, err := store.GetShardForWorkflow(ctx, "workflow-123")
+	require.NoError(t, err)
 	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()
 
 	store.SetAllShardHealth(map[uint32]bool{
 		0: true,
@@ -120,7 +137,8 @@ func TestStore_DistributionAcrossShards(t *testing.T) {
 	distribution := make(map[uint32]int)
 	for i := 0; i < 100; i++ {
 		wfID := "workflow-" + string(rune(i))
-		shard := store.GetShardForWorkflow(wfID)
+		shard, err := store.GetShardForWorkflow(ctx, wfID)
+		require.NoError(t, err)
 		distribution[shard]++
 	}
 
@@ -137,3 +155,45 @@ 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})
+
+	// Put store in transition state
+	store.SetRoutingState(&pb.RoutingState{
+		State: &pb.RoutingState_Transition{
+			Transition: &pb.Transition{WantShards: 3},
+		},
+	})
+
+	ctx, cancel := context.WithTimeout(context.Background(), 100*time.Millisecond)
+	defer cancel()
+
+	// Start a goroutine that requests allocation (will block)
+	resultCh := make(chan uint32)
+	go func() {
+		shard, _ := store.GetShardForWorkflow(ctx, "workflow-X")
+		resultCh <- shard
+	}()
+
+	// Give goroutine time to enqueue request
+	time.Sleep(10 * time.Millisecond)
+
+	// Verify request is pending
+	pending := store.GetPendingAllocations()
+	require.Contains(t, pending, "workflow-X")
+
+	// Fulfill the allocation (simulates transmitter receiving OCR outcome)
+	store.SetShardForWorkflow("workflow-X", 2)
+
+	// Blocked goroutine should now receive result
+	select {
+	case shard := <-resultCh:
+		require.Equal(t, uint32(2), shard)
+	case <-time.After(50 * time.Millisecond):
+		t.Fatal("allocation was not fulfilled")
+	}
+}

pkg/workflows/ring/transmitter.go

@@ -37,6 +37,10 @@ func (t *Transmitter) Transmit(ctx context.Context, _ types.ConfigDigest, _ uint
 		return err
 	}
 
+	// Update routing state (steady vs transition)
+	t.store.SetRoutingState(outcome.State)
+
+	// Update workflow shard mappings (this also signals any waiting allocation requests)
 	for workflowID, route := range outcome.Routes {
 		t.store.SetShardForWorkflow(workflowID, route.Shard)
 		t.lggr.Debugw("Updated workflow shard mapping", "workflowID", workflowID, "shard", route.Shard)