cre-1601: improved comments; improved naming; simplified logic; observations validation; ring performance improvements

Author: mchain0

pkg/workflows/ring/factory.go

@@ -2,6 +2,7 @@ package ring
 
 import (
 	"context"
+	"errors"
 
 	"github.com/smartcontractkit/chainlink-common/pkg/logger"
 	"github.com/smartcontractkit/chainlink-common/pkg/services"
@@ -27,8 +28,10 @@ type Factory struct {
 	services.StateMachine
 }
 
-// NewFactory creates a factory for the shard orchestrator consensus plugin
 func NewFactory(s *Store, arbiterScaler pb.ArbiterScalerClient, lggr logger.Logger, cfg *ConsensusConfig) (*Factory, error) {
+	if arbiterScaler == nil {
+		return nil, errors.New("arbiterScaler is required")
+	}
 	if cfg == nil {
 		cfg = &ConsensusConfig{
 			BatchSize: defaultBatchSize,
@@ -38,14 +41,14 @@ func NewFactory(s *Store, arbiterScaler pb.ArbiterScalerClient, lggr logger.Logg
 		store:         s,
 		arbiterScaler: arbiterScaler,
 		config:        cfg,
-		lggr:          logger.Named(lggr, "ShardOrchestratorFactory"),
+		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)
 	pluginInfo := ocr3types.ReportingPluginInfo{
-		Name: "Shard Orchestrator Consensus Plugin",
+		Name: "RingPlugin",
 		Limits: ocr3types.ReportingPluginLimits{
 			MaxQueryLength:       defaultMaxPhaseOutputBytes,
 			MaxObservationLength: defaultMaxPhaseOutputBytes,

pkg/workflows/ring/factory_test.go

@@ -12,40 +12,47 @@ import (
 func TestFactory_NewFactory(t *testing.T) {
 	lggr := logger.Test(t)
 	store := NewStore()
+	arbiter := &mockArbiter{}
 
 	t.Run("with_nil_config", func(t *testing.T) {
-		f, err := NewFactory(store, nil, lggr, nil)
+		f, err := NewFactory(store, arbiter, lggr, nil)
 		require.NoError(t, err)
 		require.NotNil(t, f)
 	})
 
 	t.Run("with_custom_config", func(t *testing.T) {
 		cfg := &ConsensusConfig{BatchSize: 50}
-		f, err := NewFactory(store, nil, lggr, cfg)
+		f, err := NewFactory(store, arbiter, lggr, cfg)
 		require.NoError(t, err)
 		require.NotNil(t, f)
 	})
+
+	t.Run("nil_arbiter_returns_error", func(t *testing.T) {
+		_, err := NewFactory(store, nil, lggr, nil)
+		require.Error(t, err)
+		require.Contains(t, err.Error(), "arbiterScaler is required")
+	})
 }
 
 func TestFactory_NewReportingPlugin(t *testing.T) {
 	lggr := logger.Test(t)
 	store := NewStore()
-	f, err := NewFactory(store, nil, lggr, nil)
+	f, err := NewFactory(store, &mockArbiter{}, lggr, nil)
 	require.NoError(t, err)
 
 	config := ocr3types.ReportingPluginConfig{N: 4, F: 1}
 	plugin, info, err := f.NewReportingPlugin(context.Background(), config)
 	require.NoError(t, err)
 	require.NotNil(t, plugin)
 	require.NotEmpty(t, info.Name)
-	require.Equal(t, "Shard Orchestrator Consensus Plugin", info.Name)
+	require.Equal(t, "RingPlugin", info.Name)
 	require.Equal(t, defaultMaxReportCount, info.Limits.MaxReportCount)
 }
 
 func TestFactory_Lifecycle(t *testing.T) {
 	lggr := logger.Test(t)
 	store := NewStore()
-	f, err := NewFactory(store, nil, lggr, nil)
+	f, err := NewFactory(store, &mockArbiter{}, lggr, nil)
 	require.NoError(t, err)
 
 	err = f.Start(context.Background())

pkg/workflows/ring/plugin.go

@@ -44,8 +44,10 @@ const (
 	DefaultTimeToSync = 5 * time.Minute
 )
 
-// NewPlugin creates a consensus reporting plugin for shard orchestration
 func NewPlugin(store *Store, arbiterScaler pb.ArbiterScalerClient, config ocr3types.ReportingPluginConfig, lggr logger.Logger, cfg *ConsensusConfig) (*Plugin, error) {
+	if arbiterScaler == nil {
+		return nil, errors.New("RingOCR arbiterScaler is required")
+	}
 	if cfg == nil {
 		cfg = &ConsensusConfig{
 			BatchSize:  DefaultBatchSize,
@@ -86,20 +88,14 @@ func (p *Plugin) Observation(ctx context.Context, _ ocr3types.OutcomeContext, _
 	var wantShards uint32
 	shardStatus := make(map[uint32]*pb.ShardStatus)
 
-	if p.arbiterScaler != nil {
-		status, err := p.arbiterScaler.Status(ctx, &emptypb.Empty{})
-		if err != nil {
-			p.lggr.Warnw("failed to get arbiter scaler status", "error", err)
-		} else {
-			wantShards = status.WantShards
-			shardStatus = status.Status
-		}
+	status, err := p.arbiterScaler.Status(ctx, &emptypb.Empty{})
+	if err != nil {
+		p.lggr.Warnw("RingOCR failed to get arbiter scaler status", "error", err)
+		wantShards = 0
+		shardStatus = make(map[uint32]*pb.ShardStatus)
 	} else {
-		// Fallback to store if no arbiter scaler configured
-		shardHealth := p.store.GetShardHealth()
-		for shardID, healthy := range shardHealth {
-			shardStatus[shardID] = &pb.ShardStatus{IsHealthy: healthy}
-		}
+		wantShards = status.WantShards
+		shardStatus = status.Status
 	}
 
 	allWorkflowIDs := make([]string, 0)
@@ -108,9 +104,11 @@ func (p *Plugin) Observation(ctx context.Context, _ ocr3types.OutcomeContext, _
 	}
 
 	pendingAllocs := p.store.GetPendingAllocations()
-	allWorkflowIDs = append(allWorkflowIDs, pendingAllocs...)
+	p.lggr.Infow("RingOCR Observation pending allocations", "pendingAllocs", pendingAllocs)
 
+	allWorkflowIDs = append(allWorkflowIDs, pendingAllocs...)
 	allWorkflowIDs = uniqueSorted(allWorkflowIDs)
+	p.lggr.Infow("RingOCR Observation all workflow IDs unique", "allWorkflowIDs", allWorkflowIDs, "wantShards", wantShards)
 
 	observation := &pb.Observation{
 		ShardStatus: shardStatus,
@@ -122,23 +120,30 @@ func (p *Plugin) Observation(ctx context.Context, _ ocr3types.OutcomeContext, _
 	return proto.MarshalOptions{Deterministic: true}.Marshal(observation)
 }
 
-//coverage:ignore
-func (p *Plugin) ValidateObservation(_ context.Context, _ ocr3types.OutcomeContext, _ types.Query, _ types.AttributedObservation) error {
+func (p *Plugin) ValidateObservation(_ context.Context, _ ocr3types.OutcomeContext, _ types.Query, ao types.AttributedObservation) error {
+	observation := &pb.Observation{}
+	if err := proto.Unmarshal(ao.Observation, observation); err != nil {
+		return err
+	}
+	if observation.Now == nil {
+		return errors.New("observation missing timestamp")
+	}
+	if observation.WantShards == 0 {
+		return errors.New("observation missing WantShards")
+	}
 	return nil
 }
 
 func (p *Plugin) ObservationQuorum(_ context.Context, _ ocr3types.OutcomeContext, _ types.Query, aos []types.AttributedObservation) (quorumReached bool, err error) {
 	return quorumhelper.ObservationCountReachesObservationQuorum(quorumhelper.QuorumTwoFPlusOne, p.config.N, p.config.F, aos), nil
 }
 
-func (p *Plugin) collectShardInfo(aos []types.AttributedObservation) (shardHealth map[uint32]int, workflows []string, timestamps []time.Time, wantShardVotes []uint32) {
+func (p *Plugin) collectShardInfo(aos []types.AttributedObservation) (shardHealth map[uint32]int, workflows []string, timestamps []time.Time, wantShardVotes map[commontypes.OracleID]uint32) {
 	shardHealth = make(map[uint32]int)
+	wantShardVotes = make(map[commontypes.OracleID]uint32)
 	for _, ao := range aos {
 		observation := &pb.Observation{}
-		if err := proto.Unmarshal(ao.Observation, observation); err != nil {
-			p.lggr.Warnf("failed to unmarshal observation: %v", err)
-			continue
-		}
+		_ = proto.Unmarshal(ao.Observation, observation) // validated in ValidateObservation
 
 		for shardID, status := range observation.ShardStatus {
 			if status != nil && status.IsHealthy {
@@ -147,14 +152,9 @@ func (p *Plugin) collectShardInfo(aos []types.AttributedObservation) (shardHealt
 		}
 
 		workflows = append(workflows, observation.WorkflowIds...)
+		timestamps = append(timestamps, observation.Now.AsTime())
 
-		if observation.Now != nil {
-			timestamps = append(timestamps, observation.Now.AsTime())
-		}
-
-		if observation.WantShards > 0 {
-			wantShardVotes = append(wantShardVotes, observation.WantShards)
-		}
+		wantShardVotes[ao.Observer] = observation.WantShards
 	}
 	return shardHealth, workflows, timestamps, wantShardVotes
 }
@@ -183,6 +183,7 @@ func (p *Plugin) Outcome(_ context.Context, outctx ocr3types.OutcomeContext, _ t
 	}
 
 	currentShardHealth, allWorkflows, nows, wantShardVotes := p.collectShardInfo(aos)
+	p.lggr.Infow("RingOCR Outcome collect shard info", "currentShardHealth", currentShardHealth, "wantShardVotes", wantShardVotes)
 
 	// Need at least F+1 timestamps; fewer means >F faulty nodes and we can't trust this round
 	if len(nows) < p.config.F+1 {
@@ -193,18 +194,13 @@ func (p *Plugin) Outcome(_ context.Context, outctx ocr3types.OutcomeContext, _ t
 	// Use the median timestamp to determine the current time
 	now := nows[len(nows)/2]
 
-	// Use median for wantShards consensus; fall back to current state if insufficient votes
-	var wantShards uint32
-	if len(wantShardVotes) >= p.config.F+1 {
-		slices.Sort(wantShardVotes)
-		wantShards = wantShardVotes[len(wantShardVotes)/2]
-	} else if rs := prior.State.GetRoutableShards(); rs > 0 {
-		wantShards = rs
-	} else if tr := prior.State.GetTransition(); tr != nil {
-		wantShards = tr.WantShards
-	} else {
-		wantShards = 1 // ultimate fallback
+	// Use median for wantShards consensus (all validated observations have WantShards > 0)
+	votes := make([]uint32, 0, len(wantShardVotes))
+	for _, v := range wantShardVotes {
+		votes = append(votes, v)
 	}
+	slices.Sort(votes)
+	wantShards := votes[len(votes)/2]
 
 	allWorkflows = uniqueSorted(allWorkflows)
 
@@ -217,20 +213,23 @@ func (p *Plugin) Outcome(_ context.Context, outctx ocr3types.OutcomeContext, _ t
 
 	// Deterministic hashing ensures all nodes agree on workflow-to-shard assignments
 	// without coordination, preventing protocol failures from inconsistent routing
+	ring := newShardRing(healthyShards)
 	routes := make(map[string]*pb.WorkflowRoute)
 	for _, wfID := range allWorkflows {
-		assignedShard := getShardForWorkflow(wfID, healthyShards)
-		routes[wfID] = &pb.WorkflowRoute{
-			Shard: assignedShard,
+		shard, err := locateShard(ring, wfID)
+		if err != nil {
+			p.lggr.Warnw("RingOCR failed to locate shard for workflow", "workflowID", wfID, "error", err)
+			shard = 0 // fallback to shard 0 when no healthy shards
 		}
+		routes[wfID] = &pb.WorkflowRoute{Shard: shard}
 	}
 
 	outcome := &pb.Outcome{
 		State:  nextState,
 		Routes: routes,
 	}
 
-	p.lggr.Infow("Consensus Outcome", "healthyShards", len(healthyShards), "totalObservations", len(aos), "workflowCount", len(routes))
+	p.lggr.Infow("RingOCR Outcome", "healthyShards", len(healthyShards), "totalObservations", len(aos), "workflowCount", len(routes))
 
 	return proto.MarshalOptions{Deterministic: true}.Marshal(outcome)
 }

pkg/workflows/ring/plugin_test.go

@@ -7,7 +7,9 @@ import (
 
 	"github.com/smartcontractkit/libocr/commontypes"
 	"github.com/stretchr/testify/require"
+	"google.golang.org/grpc"
 	"google.golang.org/protobuf/proto"
+	"google.golang.org/protobuf/types/known/emptypb"
 	"google.golang.org/protobuf/types/known/timestamppb"
 
 	"github.com/smartcontractkit/chainlink-common/pkg/logger"
@@ -16,6 +18,21 @@ import (
 	"github.com/smartcontractkit/libocr/offchainreporting2plus/ocr3types"
 )
 
+type mockArbiter struct {
+	status *pb.ReplicaStatus
+}
+
+func (m *mockArbiter) Status(ctx context.Context, in *emptypb.Empty, opts ...grpc.CallOption) (*pb.ReplicaStatus, error) {
+	if m.status != nil {
+		return m.status, nil
+	}
+	return &pb.ReplicaStatus{}, nil
+}
+
+func (m *mockArbiter) ConsensusWantShards(ctx context.Context, req *pb.ConsensusWantShardsRequest, opts ...grpc.CallOption) (*emptypb.Empty, error) {
+	return &emptypb.Empty{}, nil
+}
+
 var twoHealthyShards = []map[uint32]*pb.ShardStatus{
 	{0: {IsHealthy: true}, 1: {IsHealthy: true}},
 	{0: {IsHealthy: true}, 1: {IsHealthy: true}},
@@ -44,7 +61,7 @@ func TestPlugin_Outcome(t *testing.T) {
 			MaxDurationShouldTransmitAcceptedReport: 0,
 		}
 
-		plugin, err := NewPlugin(store, nil, config, lggr, nil)
+		plugin, err := NewPlugin(store, &mockArbiter{}, config, lggr, nil)
 		require.NoError(t, err)
 
 		ctx := t.Context()
@@ -154,7 +171,7 @@ func TestPlugin_StateTransitions(t *testing.T) {
 	}
 
 	// Use short time to sync for testing
-	plugin, err := NewPlugin(store, nil, config, lggr, &ConsensusConfig{
+	plugin, err := NewPlugin(store, &mockArbiter{}, config, lggr, &ConsensusConfig{
 		BatchSize:  100,
 		TimeToSync: 1 * time.Second,
 	})
@@ -371,6 +388,16 @@ func makeObservationsWithWantShards(t *testing.T, shardStatuses []map[uint32]*pb
 	return aos
 }
 
+func TestPlugin_NewPlugin_NilArbiter(t *testing.T) {
+	lggr := logger.Test(t)
+	store := NewStore()
+	config := ocr3types.ReportingPluginConfig{N: 4, F: 1}
+
+	_, err := NewPlugin(store, nil, config, lggr, nil)
+	require.Error(t, err)
+	require.Contains(t, err.Error(), "RingOCR arbiterScaler is required")
+}
+
 func TestPlugin_getHealthyShards(t *testing.T) {
 	tests := []struct {
 		name  string
@@ -407,13 +434,14 @@ func TestPlugin_NoHealthyShardsFallbackToShardZero(t *testing.T) {
 		N: 4, F: 1,
 	}
 
-	plugin, err := NewPlugin(store, nil, config, lggr, &ConsensusConfig{
+	arbiter := &mockArbiter{}
+	plugin, err := NewPlugin(store, arbiter, config, lggr, &ConsensusConfig{
 		BatchSize:  100,
 		TimeToSync: 1 * time.Second,
 	})
 	require.NoError(t, err)
 
-	transmitter := NewTransmitter(lggr, store, nil, "test-account")
+	transmitter := NewTransmitter(lggr, store, arbiter, "test-account")
 
 	ctx, cancel := context.WithTimeout(context.Background(), 200*time.Millisecond)
 	defer cancel()
@@ -507,7 +535,7 @@ func TestPlugin_ObservationQuorum(t *testing.T) {
 	lggr := logger.Test(t)
 	store := NewStore()
 	config := ocr3types.ReportingPluginConfig{N: 4, F: 1}
-	plugin, err := NewPlugin(store, nil, config, lggr, nil)
+	plugin, err := NewPlugin(store, &mockArbiter{}, config, lggr, nil)
 	require.NoError(t, err)
 
 	ctx := context.Background()