[Merged by Bors] - Simplify switch to different poet services

activation/nipost.go

@@ -86,8 +86,7 @@ func NewNIPostBuilder(
 	opts ...NIPostBuilderOption,
 ) (*NIPostBuilder, error) {
 	b := &NIPostBuilder{
-		localDB: db,
-
+		localDB:     db,
 		postService: postService,
 		logger:      lg,
 		poetCfg:     poetCfg,
@@ -208,9 +207,9 @@ func (nb *NIPostBuilder) BuildNIPost(
 	//  WE ARE HERE            PROOF BECOMES         ATX PUBLICATION
 	//                           AVAILABLE               DEADLINE
 
-	poetRoundStart := nb.layerClock.LayerToTime((postChallenge.PublishEpoch - 1).FirstLayer()).
+	curPoetRoundStart := nb.layerClock.LayerToTime((postChallenge.PublishEpoch - 1).FirstLayer()).
 		Add(nb.poetCfg.PhaseShift)
-	poetRoundEnd := nb.layerClock.LayerToTime(postChallenge.PublishEpoch.FirstLayer()).
+	curPoetRoundEnd := nb.layerClock.LayerToTime(postChallenge.PublishEpoch.FirstLayer()).
 		Add(nb.poetCfg.PhaseShift).
 		Add(-nb.poetCfg.CycleGap)
 
@@ -223,42 +222,21 @@ func (nb *NIPostBuilder) BuildNIPost(
 	poetProofDeadline := publishEpochEnd.Add(-nb.poetCfg.CycleGap)
 
 	logger.Info("building nipost",
-		zap.Time("poet round start", poetRoundStart),
-		zap.Time("poet round end", poetRoundEnd),
+		zap.Time("poet round start", curPoetRoundStart),
+		zap.Time("poet round end", curPoetRoundEnd),
 		zap.Time("publish epoch end", publishEpochEnd),
 		zap.Uint32("publish epoch", postChallenge.PublishEpoch.Uint32()),
 	)
 
 	// Phase 0: Submit challenge to PoET services.
-	count, err := nipost.PoetRegistrationCount(nb.localDB, signer.NodeID())
+	// Deadline: start of PoET round: we will not accept registrations after that
+	submitedRegistrations, err := nb.submitPoetChallenges(
+		ctx,
+		signer,
+		poetProofDeadline,
+		curPoetRoundStart, challenge.Bytes())
 	if err != nil {
-		return nil, fmt.Errorf("failed to get poet registration count: %w", err)
-	}
-	if count == 0 {
-		now := time.Now()
-		// Deadline: start of PoET round for publish epoch. PoET won't accept registrations after that.
-		if poetRoundStart.Before(now) {
-			return nil, fmt.Errorf(
-				"%w: poet round has already started at %s (now: %s)",
-				ErrATXChallengeExpired,
-				poetRoundStart,
-				now,
-			)
-		}
-
-		submitCtx, cancel := context.WithDeadline(ctx, poetRoundStart)
-		defer cancel()
-		err := nb.submitPoetChallenges(submitCtx, signer, poetProofDeadline, challenge.Bytes())
-		if err != nil {
-			return nil, fmt.Errorf("submitting to poets: %w", err)
-		}
-		count, err := nipost.PoetRegistrationCount(nb.localDB, signer.NodeID())
-		if err != nil {
-			return nil, fmt.Errorf("failed to get poet registration count: %w", err)
-		}
-		if count == 0 {
-			return nil, &PoetSvcUnstableError{msg: "failed to submit challenge to any PoET", source: submitCtx.Err()}
-		}
+		return nil, fmt.Errorf("submitting to poets: %w", err)
 	}
 
 	// Phase 1: query PoET services for proofs
@@ -280,8 +258,8 @@ func (nb *NIPostBuilder) BuildNIPost(
 			)
 		}
 
-		events.EmitPoetWaitProof(signer.NodeID(), postChallenge.PublishEpoch, poetRoundEnd)
-		poetProofRef, membership, err = nb.getBestProof(ctx, signer.NodeID(), challenge, postChallenge.PublishEpoch)
+		events.EmitPoetWaitProof(signer.NodeID(), postChallenge.PublishEpoch, curPoetRoundEnd)
+		poetProofRef, membership, err = nb.getBestProof(ctx, signer.NodeID(), challenge, submitedRegistrations)
 		if err != nil {
 			return nil, &PoetSvcUnstableError{msg: "getBestProof failed", source: err}
 		}
@@ -388,23 +366,74 @@ func (nb *NIPostBuilder) submitPoetChallenge(
 	})
 }
 
-// Submit the challenge to all registered PoETs.
+// Submit the challenge to registered PoETs with given addresses.
 func (nb *NIPostBuilder) submitPoetChallenges(
 	ctx context.Context,
 	signer *signing.EdSigner,
-	deadline time.Time,
+	poetProofDeadline time.Time,
+	curPoetRoundStartDeadline time.Time,
 	challenge []byte,
-) error {
+) ([]nipost.PoETRegistration, error) {
+	// check if some registrations missing
+	existingRegistrations, err := nipost.PoetRegistrationsByNodeId(nb.localDB, signer.NodeID())
+	if err != nil {
+		return nil, fmt.Errorf("failed to get poet registrations from db: %w", err)
+	}
+
+	var i int
+	addresses := make([]string, len(nb.poetProvers))
+	for _, client := range nb.poetProvers {
+		addresses[i] = client.Address()
+		i++
+	}
+
+	existingRegistrationsMap := make(map[string]struct{})
+	for _, er := range existingRegistrations {
+		existingRegistrationsMap[er.Address] = struct{}{}
+	}
+
+	missingRegistrationsPoETAddresses := make(map[string]struct{})
+	for _, addr := range addresses {
+		if _, ok := existingRegistrationsMap[addr]; !ok {
+			missingRegistrationsPoETAddresses[addr] = struct{}{}
+		}
+	}
+
+	if len(missingRegistrationsPoETAddresses) == 0 {
+		return existingRegistrations, nil
+	}
+
+	now := time.Now()
+	if curPoetRoundStartDeadline.Before(now) {
+		if len(missingRegistrationsPoETAddresses) == len(addresses) {
+			return nil, fmt.Errorf(
+				"%w: poet round has already started at %s (now: %s)",
+				ErrATXChallengeExpired, // no existing registration, drop current registration challenge
+				curPoetRoundStartDeadline,
+				now,
+			)
+		}
+		return existingRegistrations, nil
+	}
+
+	// send registrations to missing addresses
 	signature := signer.Sign(signing.POET, challenge)
 	prefix := bytes.Join([][]byte{signer.Prefix(), {byte(signing.POET)}}, nil)
 	nodeID := signer.NodeID()
-	g, ctx := errgroup.WithContext(ctx)
-	errChan := make(chan error, len(nb.poetProvers))
+
+	submitCtx, cancel := context.WithDeadline(ctx, curPoetRoundStartDeadline)
+	defer cancel()
+
+	g, ctx := errgroup.WithContext(submitCtx)
+	errChan := make(chan error, len(missingRegistrationsPoETAddresses))
+
 	for _, client := range nb.poetProvers {
-		g.Go(func() error {
-			errChan <- nb.submitPoetChallenge(ctx, nodeID, deadline, client, prefix, challenge, signature)
-			return nil
-		})
+		if _, ok := missingRegistrationsPoETAddresses[client.Address()]; ok {
+			g.Go(func() error {
+				errChan <- nb.submitPoetChallenge(ctx, nodeID, poetProofDeadline, client, prefix, challenge, signature)
+				return nil
+			})
+		}
 	}
 	g.Wait()
 	close(errChan)
@@ -423,18 +452,19 @@ func (nb *NIPostBuilder) submitPoetChallenges(
 	}
 	if allInvalid {
 		nb.logger.Warn("all poet submits were too late. ATX challenge expires", log.ZShortStringer("smesherID", nodeID))
-		return ErrATXChallengeExpired
+		return nil, ErrATXChallengeExpired
 	}
-	return nil
-}
 
-func (nb *NIPostBuilder) getPoetService(ctx context.Context, address string) PoetService {
-	for _, service := range nb.poetProvers {
-		if address == service.Address() {
-			return service
-		}
+	existingRegistrations, err = nipost.PoetRegistrationsByNodeIdAndAddresses(nb.localDB, signer.NodeID(), addresses)
+	if err != nil {
+		return nil, fmt.Errorf("failed to get poet registration count: %w", err)
 	}
-	return nil
+
+	if len(existingRegistrations) == 0 {
+		return nil, &PoetSvcUnstableError{msg: "failed to submit challenge to any PoET", source: ctx.Err()}
+	}
+
+	return existingRegistrations, nil
 }
 
 // membersContainChallenge verifies that the challenge is included in proof's members.
@@ -451,16 +481,12 @@ func (nb *NIPostBuilder) getBestProof(
 	ctx context.Context,
 	nodeID types.NodeID,
 	challenge types.Hash32,
-	publishEpoch types.EpochID,
+	registrations []nipost.PoETRegistration,
 ) (types.PoetProofRef, *types.MerkleProof, error) {
 	type poetProof struct {
 		poet       *types.PoetProof
 		membership *types.MerkleProof
 	}
-	registrations, err := nipost.PoetRegistrations(nb.localDB, nodeID)
-	if err != nil {
-		return types.PoetProofRef{}, nil, fmt.Errorf("getting poet registrations: %w", err)
-	}
 	proofs := make(chan *poetProof, len(registrations))
 
 	var eg errgroup.Group
@@ -471,11 +497,13 @@ func (nb *NIPostBuilder) getBestProof(
 			zap.String("poet_address", r.Address),
 			zap.String("round", r.RoundID),
 		)
-		client := nb.getPoetService(ctx, r.Address)
-		if client == nil {
+
+		client, ok := nb.poetProvers[r.Address]
+		if !ok {
 			logger.Warn("poet client not found")
 			continue
 		}
+
 		round := r.RoundID
 		waitDeadline := proofDeadline(r.RoundEnd, nb.poetCfg.CycleGap)
 		eg.Go(func() error {