Validator subscribe subnet to next epoch (#5312)

* Alert to subscribe to next epoch * Fixed tests * Comments * Fixed tests * Update validator/client/validator.go Co-authored-by: Ivan Martinez <[email protected]> Co-authored-by: prylabs-bulldozer[bot] <58059840+prylabs-bulldozer[bot]@users.noreply.github.com>
prysmaticlabs · Apr 6, 2020 · 963d11a · 963d11a
1 parent 96af0e8
commit 963d11a
Show file tree

Hide file tree

Showing 2 changed files with 42 additions and 23 deletions.
diff --git a/validator/client/validator.go b/validator/client/validator.go
@@ -311,12 +311,49 @@ func (v *validator) UpdateDuties(ctx context.Context, slot uint64) error {
 					lFields["proposerSlot"] = duty.ProposerSlot
 				}
 				lFields["attesterSlot"] = duty.AttesterSlot
+
+				aggregator, err := v.isAggregator(ctx, duty.Committee, duty.AttesterSlot, bytesutil.ToBytes48(duty.PublicKey))
+				if err != nil {
+					return errors.Wrap(err, "could not check if a validator is an aggregator")
+				}
+				if _, err := v.validatorClient.SubscribeCommitteeSubnet(ctx, &ethpb.CommitteeSubnetSubscribeRequest{
+					Slot:         duty.AttesterSlot,
+					CommitteeId:  duty.CommitteeIndex,
+					IsAggregator: aggregator,
+				}); err != nil {
+					return err
+				}
 			}
 
 			log.WithFields(lFields).Info("New assignment")
 		}
 	}
 
+	// Notify beacon node to subscribe to the attester and aggregator subnets for the next epoch.
+	req.Epoch++
+	dutiesNextEpoch, err := v.validatorClient.GetDuties(ctx, req)
+	if err != nil {
+		log.Error(err)
+		return err
+	}
+	if slot%params.BeaconConfig().SlotsPerEpoch == 0 || firstDutiesReceived {
+		for _, duty := range dutiesNextEpoch.Duties {
+			if duty.Status == ethpb.ValidatorStatus_ACTIVE {
+				aggregator, err := v.isAggregator(ctx, duty.Committee, duty.AttesterSlot, bytesutil.ToBytes48(duty.PublicKey))
+				if err != nil {
+					return errors.Wrap(err, "could not check if a validator is an aggregator")
+				}
+				if _, err := v.validatorClient.SubscribeCommitteeSubnet(ctx, &ethpb.CommitteeSubnetSubscribeRequest{
+					Slot:         duty.AttesterSlot,
+					CommitteeId:  duty.CommitteeIndex,
+					IsAggregator: aggregator,
+				}); err != nil {
+					return err
+				}
+			}
+		}
+	}
+
 	return nil
 }
 
@@ -344,14 +381,6 @@ func (v *validator) RolesAt(ctx context.Context, slot uint64) (map[[48]byte][]va
 			if aggregator {
 				roles = append(roles, roleAggregator)
 			}
-
-			if _, err := v.validatorClient.SubscribeCommitteeSubnet(ctx, &ethpb.CommitteeSubnetSubscribeRequest{
-				Slot:         slot,
-				CommitteeId:  duty.CommitteeIndex,
-				IsAggregator: aggregator,
-			}); err != nil {
-				return nil, err
-			}
 		}
 		if len(roles) == 0 {
 			roles = append(roles, roleUnknown)

diff --git a/validator/client/validator_test.go b/validator/client/validator_test.go
@@ -530,6 +530,11 @@ func TestUpdateDuties_OK(t *testing.T) {
 		gomock.Any(),
 	).Return(resp, nil)
 
+	client.EXPECT().GetDuties(
+		gomock.Any(),
+		gomock.Any(),
+	).Return(resp, nil)
+
 	if err := v.UpdateDuties(context.Background(), slot); err != nil {
 		t.Fatalf("Could not update assignments: %v", err)
 	}
@@ -609,16 +614,6 @@ func TestRolesAt_OK(t *testing.T) {
 		gomock.Any(), // epoch
 	).Return(&ethpb.DomainResponse{}, nil /*err*/)
 
-	m.validatorClient.EXPECT().SubscribeCommitteeSubnet(
-		gomock.Any(), // ctx
-		gomock.Any(), // epoch
-	).Return(nil, nil /*err*/)
-
-	m.validatorClient.EXPECT().SubscribeCommitteeSubnet(
-		gomock.Any(), // ctx
-		gomock.Any(), // epoch
-	).Return(nil, nil /*err*/)
-
 	roleMap, err := v.RolesAt(context.Background(), 1)
 	if err != nil {
 		t.Fatal(err)
@@ -681,11 +676,6 @@ func TestRolesAt_DoesNotAssignProposer_Slot0(t *testing.T) {
 		gomock.Any(), // epoch
 	).Return(&ethpb.DomainResponse{}, nil /*err*/)
 
-	m.validatorClient.EXPECT().SubscribeCommitteeSubnet(
-		gomock.Any(), // ctx
-		gomock.Any(), // epoch
-	).Return(nil, nil /*err*/)
-
 	roleMap, err := v.RolesAt(context.Background(), 0)
 	if err != nil {
 		t.Fatal(err)