[CRE-1601] shard-orchestrator implementation (#1747)

* feat: implement shard-orchestrator logic

* chore: use state constants instead of strings

* test: add pluging <-> shardorchestrator integration tests

* tests: add server tests

* feat: add client to communicate with shard 0

* chore: remove unused timestamps

* fix: remove reduntant check

---------

Co-authored-by: mchain0 <maciej.wisniewski@smartcontract.com>

Author: 2 people

pkg/workflows/ring/factory.go

@@ -4,11 +4,13 @@ import (
 	"context"
 	"errors"
 
+	"github.com/smartcontractkit/libocr/offchainreporting2plus/ocr3types"
+
 	"github.com/smartcontractkit/chainlink-common/pkg/logger"
 	"github.com/smartcontractkit/chainlink-common/pkg/services"
 	"github.com/smartcontractkit/chainlink-common/pkg/types/core"
 	"github.com/smartcontractkit/chainlink-common/pkg/workflows/ring/pb"
-	"github.com/smartcontractkit/libocr/offchainreporting2plus/ocr3types"
+	"github.com/smartcontractkit/chainlink-common/pkg/workflows/shardorchestrator"
 )
 
 const (
@@ -20,15 +22,16 @@ const (
 var _ core.OCR3ReportingPluginFactory = &Factory{}
 
 type Factory struct {
-	store         *Store
-	arbiterScaler pb.ArbiterScalerClient
-	config        *ConsensusConfig
-	lggr          logger.Logger
+	ringStore              *Store
+	shardOrchestratorStore *shardorchestrator.Store
+	arbiterScaler          pb.ArbiterScalerClient
+	config                 *ConsensusConfig
+	lggr                   logger.Logger
 
 	services.StateMachine
 }
 
-func NewFactory(s *Store, arbiterScaler pb.ArbiterScalerClient, lggr logger.Logger, cfg *ConsensusConfig) (*Factory, error) {
+func NewFactory(s *Store, shardOrchestratorStore *shardorchestrator.Store, arbiterScaler pb.ArbiterScalerClient, lggr logger.Logger, cfg *ConsensusConfig) (*Factory, error) {
 	if arbiterScaler == nil {
 		return nil, errors.New("arbiterScaler is required")
 	}
@@ -38,15 +41,16 @@ func NewFactory(s *Store, arbiterScaler pb.ArbiterScalerClient, lggr logger.Logg
 		}
 	}
 	return &Factory{
-		store:         s,
-		arbiterScaler: arbiterScaler,
-		config:        cfg,
-		lggr:          logger.Named(lggr, "RingPluginFactory"),
+		ringStore:              s,
+		shardOrchestratorStore: shardOrchestratorStore,
+		arbiterScaler:          arbiterScaler,
+		config:                 cfg,
+		lggr:                   logger.Named(lggr, "RingPluginFactory"),
 	}, nil
 }
 
 func (o *Factory) NewReportingPlugin(_ context.Context, config ocr3types.ReportingPluginConfig) (ocr3types.ReportingPlugin[[]byte], ocr3types.ReportingPluginInfo, error) {
-	plugin, err := NewPlugin(o.store, o.arbiterScaler, config, o.lggr, o.config)
+	plugin, err := NewPlugin(o.ringStore, o.arbiterScaler, config, o.lggr, o.config)
 	pluginInfo := ocr3types.ReportingPluginInfo{
 		Name: "RingPlugin",
 		Limits: ocr3types.ReportingPluginLimits{

pkg/workflows/ring/factory_test.go

@@ -4,15 +4,18 @@ import (
 	"context"
 	"testing"
 
-	"github.com/smartcontractkit/chainlink-common/pkg/logger"
-	"github.com/smartcontractkit/chainlink-common/pkg/workflows/ring/pb"
 	"github.com/smartcontractkit/libocr/offchainreporting2plus/ocr3types"
 	"github.com/stretchr/testify/require"
+
+	"github.com/smartcontractkit/chainlink-common/pkg/logger"
+	"github.com/smartcontractkit/chainlink-common/pkg/workflows/ring/pb"
+	"github.com/smartcontractkit/chainlink-common/pkg/workflows/shardorchestrator"
 )
 
 func TestFactory_NewFactory(t *testing.T) {
 	lggr := logger.Test(t)
 	store := NewStore()
+	shardOrchestratorStore := shardorchestrator.NewStore(lggr)
 	arbiter := &mockArbiter{}
 
 	tests := []struct {
@@ -45,7 +48,7 @@ func TestFactory_NewFactory(t *testing.T) {
 
 	for _, tt := range tests {
 		t.Run(tt.name, func(t *testing.T) {
-			f, err := NewFactory(store, tt.arbiter, lggr, tt.config)
+			f, err := NewFactory(store, shardOrchestratorStore, tt.arbiter, lggr, tt.config)
 			if tt.wantErr {
 				require.Error(t, err)
 				require.Contains(t, err.Error(), tt.errSubstr)
@@ -60,7 +63,7 @@ func TestFactory_NewFactory(t *testing.T) {
 func TestFactory_NewReportingPlugin(t *testing.T) {
 	lggr := logger.Test(t)
 	store := NewStore()
-	f, err := NewFactory(store, &mockArbiter{}, lggr, nil)
+	f, err := NewFactory(store, nil, &mockArbiter{}, lggr, nil)
 	require.NoError(t, err)
 
 	config := ocr3types.ReportingPluginConfig{N: 4, F: 1}
@@ -75,7 +78,7 @@ func TestFactory_NewReportingPlugin(t *testing.T) {
 func TestFactory_Lifecycle(t *testing.T) {
 	lggr := logger.Test(t)
 	store := NewStore()
-	f, err := NewFactory(store, &mockArbiter{}, lggr, nil)
+	f, err := NewFactory(store, nil, &mockArbiter{}, lggr, nil)
 	require.NoError(t, err)
 
 	err = f.Start(context.Background())

pkg/workflows/ring/pb/generate.go

@@ -1,6 +1,5 @@
 //go:generate protoc --go_out=. --go_opt=paths=source_relative shared.proto
 //go:generate protoc --go_out=. --go_opt=paths=source_relative --go-grpc_out=. --go-grpc_opt=paths=source_relative arbiter.proto
-//go:generate protoc --go_out=. --go_opt=paths=source_relative --go-grpc_out=. --go-grpc_opt=paths=source_relative shard_orchestrator.proto
 //go:generate protoc --go_out=. --go_opt=paths=source_relative consensus.proto
 
 package pb

pkg/workflows/ring/plugin_test.go

@@ -12,10 +12,12 @@ import (
 	"google.golang.org/protobuf/types/known/emptypb"
 	"google.golang.org/protobuf/types/known/timestamppb"
 
-	"github.com/smartcontractkit/chainlink-common/pkg/logger"
-	"github.com/smartcontractkit/chainlink-common/pkg/workflows/ring/pb"
 	"github.com/smartcontractkit/libocr/offchainreporting2/types"
 	"github.com/smartcontractkit/libocr/offchainreporting2plus/ocr3types"
+
+	"github.com/smartcontractkit/chainlink-common/pkg/logger"
+	"github.com/smartcontractkit/chainlink-common/pkg/workflows/ring/pb"
+	"github.com/smartcontractkit/chainlink-common/pkg/workflows/shardorchestrator"
 )
 
 type mockArbiter struct {
@@ -443,7 +445,7 @@ func TestPlugin_NoHealthyShardsFallbackToShardZero(t *testing.T) {
 	})
 	require.NoError(t, err)
 
-	transmitter := NewTransmitter(lggr, store, arbiter, "test-account")
+	transmitter := NewTransmitter(lggr, store, nil, arbiter, "test-account")
 
 	ctx, cancel := context.WithTimeout(context.Background(), 200*time.Millisecond)
 	defer cancel()
@@ -591,3 +593,158 @@ func TestPlugin_ObservationQuorum(t *testing.T) {
 		require.True(t, quorum)
 	})
 }
+
+func TestPlugin_ShardOrchestratorIntegration(t *testing.T) {
+	lggr := logger.Test(t)
+
+	// Create both stores
+	ringStore := NewStore()
+	orchestratorStore := shardorchestrator.NewStore(lggr)
+
+	// Initialize ring store with healthy shards
+	ringStore.SetAllShardHealth(map[uint32]bool{0: true, 1: true, 2: true})
+
+	config := ocr3types.ReportingPluginConfig{
+		N: 4, F: 1,
+	}
+
+	arbiter := &mockArbiter{}
+	plugin, err := NewPlugin(ringStore, arbiter, config, lggr, &ConsensusConfig{
+		BatchSize:  100,
+		TimeToSync: 1 * time.Second,
+	})
+	require.NoError(t, err)
+
+	// Create transmitter with both stores
+	transmitter := NewTransmitter(lggr, ringStore, orchestratorStore, arbiter, "test-account")
+
+	ctx := context.Background()
+	now := time.Now()
+
+	t.Run("initial_workflow_assignments", func(t *testing.T) {
+		// Create observations with workflows
+		workflows := []string{"wf-A", "wf-B", "wf-C"}
+		aos := makeObservationsWithWantShards(t, []map[uint32]*pb.ShardStatus{
+			{0: {IsHealthy: true}, 1: {IsHealthy: true}, 2: {IsHealthy: true}},
+			{0: {IsHealthy: true}, 1: {IsHealthy: true}, 2: {IsHealthy: true}},
+			{0: {IsHealthy: true}, 1: {IsHealthy: true}, 2: {IsHealthy: true}},
+		}, workflows, now, 3)
+
+		outcomeCtx := ocr3types.OutcomeContext{
+			SeqNr:           1,
+			PreviousOutcome: nil,
+		}
+
+		// Generate outcome
+		outcome, err := plugin.Outcome(ctx, outcomeCtx, nil, aos)
+		require.NoError(t, err)
+
+		// Generate report and transmit
+		reports, err := plugin.Reports(ctx, 1, outcome)
+		require.NoError(t, err)
+		require.Len(t, reports, 1)
+
+		err = transmitter.Transmit(ctx, types.ConfigDigest{}, 1, reports[0].ReportWithInfo, nil)
+		require.NoError(t, err)
+
+		// Verify ring store was updated
+		for _, wf := range workflows {
+			shard, err := ringStore.GetShardForWorkflow(ctx, wf)
+			require.NoError(t, err)
+			require.LessOrEqual(t, shard, uint32(2), "workflow should be assigned to valid shard")
+			t.Logf("Ring store: %s → shard %d", wf, shard)
+		}
+
+		// Verify orchestrator store was updated with correct state
+		for _, wf := range workflows {
+			mapping, err := orchestratorStore.GetWorkflowMapping(ctx, wf)
+			require.NoError(t, err)
+			require.Equal(t, wf, mapping.WorkflowID)
+			require.LessOrEqual(t, mapping.NewShardID, uint32(2))
+			require.Equal(t, uint32(0), mapping.OldShardID, "initial assignment should have oldShardID=0")
+			require.Equal(t, shardorchestrator.StateSteady, mapping.TransitionState, "initial assignment should be steady")
+			t.Logf("Orchestrator store: %s → shard %d (state: %s)", wf, mapping.NewShardID, mapping.TransitionState.String())
+		}
+
+		// Verify version tracking
+		version := orchestratorStore.GetMappingVersion()
+		require.Equal(t, uint64(1), version, "version should increment after first update")
+	})
+
+	t.Run("workflow_transition_detected", func(t *testing.T) {
+		// First, establish a baseline with workflows distributed across 3 shards
+		// Use wantShards=3 to ensure workflows actually get assigned to shard 2
+		baselineAos := makeObservationsWithWantShards(t, []map[uint32]*pb.ShardStatus{
+			{0: {IsHealthy: true}, 1: {IsHealthy: true}, 2: {IsHealthy: true}},
+			{0: {IsHealthy: true}, 1: {IsHealthy: true}, 2: {IsHealthy: true}},
+			{0: {IsHealthy: true}, 1: {IsHealthy: true}, 2: {IsHealthy: true}},
+		}, []string{"wf-A", "wf-B", "wf-C", "wf-D", "wf-E"}, now, 3)
+
+		baselineOutcome, err := plugin.Outcome(ctx, ocr3types.OutcomeContext{SeqNr: 2}, nil, baselineAos)
+		require.NoError(t, err)
+
+		baselineReports, err := plugin.Reports(ctx, 2, baselineOutcome)
+		require.NoError(t, err)
+
+		err = transmitter.Transmit(ctx, types.ConfigDigest{}, 2, baselineReports[0].ReportWithInfo, nil)
+		require.NoError(t, err)
+
+		// Parse baseline to see which workflows were on shard 2
+		baselineProto := &pb.Outcome{}
+		err = proto.Unmarshal(baselineOutcome, baselineProto)
+		require.NoError(t, err)
+
+		workflowsOnShard2 := []string{}
+		for wfID, route := range baselineProto.Routes {
+			if route.Shard == 2 {
+				workflowsOnShard2 = append(workflowsOnShard2, wfID)
+			}
+			t.Logf("Baseline: %s on shard %d", wfID, route.Shard)
+		}
+		require.NotEmpty(t, workflowsOnShard2, "at least one workflow should be on shard 2 for this test")
+
+		// Now scale down to 2 shards - workflows on shard 2 MUST move
+		transitionAos := makeObservationsWithWantShards(t, []map[uint32]*pb.ShardStatus{
+			{0: {IsHealthy: true}, 1: {IsHealthy: true}},
+			{0: {IsHealthy: true}, 1: {IsHealthy: true}},
+			{0: {IsHealthy: true}, 1: {IsHealthy: true}},
+		}, []string{"wf-A", "wf-B", "wf-C", "wf-D", "wf-E"}, now, 2)
+
+		outcomeCtx := ocr3types.OutcomeContext{
+			SeqNr:           3,
+			PreviousOutcome: baselineOutcome,
+		}
+
+		outcome, err := plugin.Outcome(ctx, outcomeCtx, nil, transitionAos)
+		require.NoError(t, err)
+
+		reports, err := plugin.Reports(ctx, 3, outcome)
+		require.NoError(t, err)
+
+		err = transmitter.Transmit(ctx, types.ConfigDigest{}, 3, reports[0].ReportWithInfo, nil)
+		require.NoError(t, err)
+
+		// Verify orchestrator store shows transition state for workflows that moved from shard 2
+		outcomeProto := &pb.Outcome{}
+		err = proto.Unmarshal(outcome, outcomeProto)
+		require.NoError(t, err)
+
+		// Workflows that were on shard 2 must have moved and should show TransitionState
+		for _, wfID := range workflowsOnShard2 {
+			mapping, err := orchestratorStore.GetWorkflowMapping(ctx, wfID)
+			require.NoError(t, err)
+
+			newRoute := outcomeProto.Routes[wfID]
+			require.NotEqual(t, uint32(2), newRoute.Shard, "workflow should have moved from shard 2")
+			require.Equal(t, shardorchestrator.StateTransitioning, mapping.TransitionState,
+				"workflow %s moved from shard 2 to shard %d, should be transitioning", wfID, newRoute.Shard)
+			require.Equal(t, uint32(2), mapping.OldShardID, "should track old shard")
+			require.Equal(t, newRoute.Shard, mapping.NewShardID, "should track new shard")
+			t.Logf("Workflow %s transitioned: shard 2 → %d", wfID, newRoute.Shard)
+		}
+
+		// Verify version incremented
+		version := orchestratorStore.GetMappingVersion()
+		require.Equal(t, uint64(3), version, "version should increment after update")
+	})
+}

pkg/workflows/ring/state.go

@@ -7,8 +7,25 @@ import (
 	"google.golang.org/protobuf/types/known/timestamppb"
 
 	"github.com/smartcontractkit/chainlink-common/pkg/workflows/ring/pb"
+	"github.com/smartcontractkit/chainlink-common/pkg/workflows/shardorchestrator"
 )
 
+// TransitionStateFromBool converts a proto bool (in_transition) to TransitionState
+func TransitionStateFromBool(inTransition bool) shardorchestrator.TransitionState {
+	if inTransition {
+		return shardorchestrator.StateTransitioning
+	}
+	return shardorchestrator.StateSteady
+}
+
+// TransitionStateFromRoutingState returns the TransitionState based on RoutingState
+func TransitionStateFromRoutingState(state *pb.RoutingState) shardorchestrator.TransitionState {
+	if IsInSteadyState(state) {
+		return shardorchestrator.StateSteady
+	}
+	return shardorchestrator.StateTransitioning
+}
+
 func IsInSteadyState(state *pb.RoutingState) bool {
 	if state == nil {
 		return false