cre-1601: Transmitter notifies Arbiter

Author: mchain0

pkg/workflows/ring/transmitter.go

@@ -15,26 +15,26 @@ var _ ocr3types.ContractTransmitter[[]byte] = (*Transmitter)(nil)
 
 // Transmitter handles transmission of shard orchestration outcomes
 type Transmitter struct {
-	lggr        logger.Logger
-	store       *Store
-	fromAccount types.Account
+	lggr          logger.Logger
+	store         *Store
+	arbiterScaler pb.ArbiterScalerClient
+	fromAccount   types.Account
 }
 
-func NewTransmitter(lggr logger.Logger, store *Store, fromAccount types.Account) *Transmitter {
-	return &Transmitter{lggr: lggr, store: store, fromAccount: fromAccount}
+func NewTransmitter(lggr logger.Logger, store *Store, arbiterScaler pb.ArbiterScalerClient, fromAccount types.Account) *Transmitter {
+	return &Transmitter{lggr: lggr, store: store, arbiterScaler: arbiterScaler, fromAccount: fromAccount}
 }
 
-func (t *Transmitter) Transmit(_ context.Context, _ types.ConfigDigest, _ uint64, r ocr3types.ReportWithInfo[[]byte], _ []types.AttributedOnchainSignature) error {
+func (t *Transmitter) Transmit(ctx context.Context, _ types.ConfigDigest, _ uint64, r ocr3types.ReportWithInfo[[]byte], _ []types.AttributedOnchainSignature) error {
 	outcome := &pb.Outcome{}
 	if err := proto.Unmarshal(r.Report, outcome); err != nil {
 		t.lggr.Errorf("failed to unmarshal report")
 		return err
 	}
 
-	if outcome.State != nil {
-		if routableShards, ok := outcome.State.State.(*pb.RoutingState_RoutableShards); ok {
-			t.lggr.Infow("Transmitting shard routing", "routableShards", routableShards.RoutableShards)
-		}
+	if err := t.notifyArbiter(ctx, outcome.State); err != nil {
+		t.lggr.Errorf("failed to notify arbiter", "err", err)
+		return err
 	}
 
 	for workflowID, route := range outcome.Routes {
@@ -45,6 +45,30 @@ func (t *Transmitter) Transmit(_ context.Context, _ types.ConfigDigest, _ uint64
 	return nil
 }
 
+func (t *Transmitter) notifyArbiter(ctx context.Context, state *pb.RoutingState) error {
+	if state == nil {
+		return nil
+	}
+
+	var nShards uint32
+	switch s := state.State.(type) {
+	case *pb.RoutingState_RoutableShards:
+		nShards = s.RoutableShards
+		t.lggr.Infow("Transmitting shard routing", "routableShards", nShards)
+	case *pb.RoutingState_Transition:
+		nShards = s.Transition.WantShards
+		t.lggr.Infow("Transmitting shard routing (in transition)", "wantShards", nShards)
+	}
+
+	if t.arbiterScaler != nil && nShards > 0 {
+		if _, err := t.arbiterScaler.ConsensusWantShards(ctx, &pb.ConsensusWantShardsRequest{NShards: nShards}); err != nil {
+			return err
+		}
+	}
+
+	return nil
+}
+
 func (t *Transmitter) FromAccount(ctx context.Context) (types.Account, error) {
 	return t.fromAccount, nil
 }