Pull apart writeCurrentStakers method - Part 1

vms/platformvm/state/state.go

@@ -1735,14 +1735,21 @@ func (s *state) initValidatorSets() error {
 }
 
 func (s *state) write(updateValidators bool, height uint64) error {
+	diffs, err := s.calculateDiffs()
+	if err != nil {
+		return err
+	}
+
 	codecVersion := v1
 	if !s.cfg.IsDurangoActivated(s.GetTimestamp()) {
 		codecVersion = v0
 	}
 
 	return utils.Err(
 		s.writeBlocks(),
-		s.writeCurrentStakers(updateValidators, height, codecVersion),
+		s.writeCurrentStakers(updateValidators, codecVersion),
+		s.writeWeightDiffs(height, diffs.weightDiffs),
+		s.writeBlsKeyDiffs(height, diffs.blsKeyDiffs),
 		s.writePendingStakers(),
 		s.WriteValidatorMetadata(s.currentValidatorList, s.currentSubnetValidatorList, codecVersion), // Must be called after writeCurrentStakers
 		s.writeTXs(),
@@ -1757,6 +1764,165 @@ func (s *state) write(updateValidators bool, height uint64) error {
 	)
 }
 
+func (s *state) writeWeightDiffs(height uint64, weightDiffs map[subnetNodeKey]*ValidatorWeightDiff) error {
+	for k, weightDiff := range weightDiffs {
+		if weightDiff.Amount == 0 {
+			continue
+		}
+
+		err := s.flatValidatorWeightDiffsDB.Put(
+			marshalDiffKey(k.subnetID, height, k.nodeID),
+			marshalWeightDiff(weightDiff),
+		)
+		if err != nil {
+			return err
+		}
+
+		// TODO: Remove this once we no longer support version rollbacks.
+		prefixStruct := heightWithSubnet{
+			Height:   height,
+			SubnetID: k.subnetID,
+		}
+		prefixBytes, err := block.GenesisCodec.Marshal(block.Version, prefixStruct)
+		if err != nil {
+			return fmt.Errorf("failed to create prefix bytes: %w", err)
+		}
+		rawNestedWeightDiffDB := prefixdb.New(prefixBytes, s.nestedValidatorWeightDiffsDB)
+		nestedWeightDiffDB := linkeddb.NewDefault(rawNestedWeightDiffDB)
+
+		weightDiffBytes, err := block.GenesisCodec.Marshal(block.Version, weightDiff)
+		if err != nil {
+			return fmt.Errorf("failed to serialize validator weight diff: %w", err)
+		}
+		if err := nestedWeightDiffDB.Put(k.nodeID[:], weightDiffBytes); err != nil {
+			return err
+		}
+	}
+	return nil
+}
+
+func (s *state) writeBlsKeyDiffs(height uint64, blsKeyDiffs map[ids.NodeID]*bls.PublicKey) error {
+	heightBytes := database.PackUInt64(height)
+	rawNestedPublicKeyDiffDB := prefixdb.New(heightBytes, s.nestedValidatorPublicKeyDiffsDB)
+	nestedPKDiffDB := linkeddb.NewDefault(rawNestedPublicKeyDiffDB)
+
+	for nodeID, blsKey := range blsKeyDiffs {
+		key := marshalDiffKey(constants.PrimaryNetworkID, height, nodeID)
+		var blsKeyBytes []byte
+		if blsKey != nil {
+			// Note: in flatValidatorPublicKeyDiffsDB we store the
+			// uncompressed public key here as it is
+			// significantly more efficient to parse when applying diffs.
+			blsKeyBytes = bls.SerializePublicKey(blsKey)
+		}
+		if err := s.flatValidatorPublicKeyDiffsDB.Put(key, blsKeyBytes); err != nil {
+			return fmt.Errorf("failed to add bls key diffs: %w", err)
+		}
+
+		// TODO: Remove this once we no longer support version rollbacks.
+		if blsKey != nil {
+			// Note: We store the compressed public key here.
+			pkBytes := bls.PublicKeyToBytes(blsKey)
+			if err := nestedPKDiffDB.Put(nodeID[:], pkBytes); err != nil {
+				return err
+			}
+		}
+	}
+	return nil
+}
+
+type subnetNodeKey struct {
+	subnetID ids.ID
+	nodeID   ids.NodeID
+}
+
+type diffs struct {
+	weightDiffs map[subnetNodeKey]*ValidatorWeightDiff
+	blsKeyDiffs map[ids.NodeID]*bls.PublicKey
+}
+
+func (s *state) calculateDiffs() (diffs, error) {
+	var (
+		outputWeights = make(map[subnetNodeKey]*ValidatorWeightDiff)
+		outputBlsKey  = make(map[ids.NodeID]*bls.PublicKey)
+	)
+
+	for subnetID, subnetValidatorDiffs := range s.currentStakers.validatorDiffs {
+		// for now, let writeCurrentStakers consume s.currentStakers.validatorDiffs
+		for nodeID, validatorDiff := range subnetValidatorDiffs {
+			// Note: validatorDiff.validator is not guaranteed to be non-nil here.
+			// Access it only if validatorDiff.validatorStatus is added or deleted
+
+			weightKey := subnetNodeKey{
+				subnetID: subnetID,
+				nodeID:   nodeID,
+			}
+
+			// make sure there is an entry for delegators even in case
+			// there are no validators modified.
+			outputWeights[weightKey] = &ValidatorWeightDiff{
+				Decrease: validatorDiff.validatorStatus == deleted,
+			}
+
+			switch validatorDiff.validatorStatus {
+			case added:
+				var (
+					weight = validatorDiff.validator.Weight
+					blsKey = validatorDiff.validator.PublicKey
+				)
+
+				outputWeights[weightKey].Amount = weight
+
+				if blsKey != nil {
+					// Record that the public key for the validator is being
+					// added. This means the prior value for the public key was
+					// nil.
+					outputBlsKey[nodeID] = nil
+				}
+
+			case deleted:
+				var (
+					weight = validatorDiff.validator.Weight
+					blkKey = validatorDiff.validator.PublicKey
+				)
+				outputWeights[weightKey].Amount = weight
+				if blkKey != nil {
+					// Record that the public key for the validator is being
+					// removed. This means we must record the prior value of the
+					// public key.
+					outputBlsKey[nodeID] = blkKey
+				}
+
+			default:
+				// nothing to do
+			}
+
+			addedDelegatorIterator := NewTreeIterator(validatorDiff.addedDelegators)
+			// We don't defer iterator.Release here to avoid pinning every iterator in
+			// memory till the outer loop is done.
+			for addedDelegatorIterator.Next() {
+				staker := addedDelegatorIterator.Value()
+				if err := outputWeights[weightKey].Add(false /*negative*/, staker.Weight); err != nil {
+					addedDelegatorIterator.Release()
+					return diffs{}, fmt.Errorf("failed to increase node weight diff: %w", err)
+				}
+			}
+			addedDelegatorIterator.Release()
+
+			for _, staker := range validatorDiff.deletedDelegators {
+				if err := outputWeights[weightKey].Add(true /*negative*/, staker.Weight); err != nil {
+					return diffs{}, fmt.Errorf("failed to decrease node weight diff: %w", err)
+				}
+			}
+		}
+	}
+
+	return diffs{
+		weightDiffs: outputWeights,
+		blsKeyDiffs: outputBlsKey,
+	}, nil
+}
+
 func (s *state) Close() error {
 	return utils.Err(
 		s.pendingSubnetValidatorBaseDB.Close(),
@@ -1970,11 +2136,7 @@ func (s *state) GetBlockIDAtHeight(height uint64) (ids.ID, error) {
 	return blkID, nil
 }
 
-func (s *state) writeCurrentStakers(updateValidators bool, height uint64, codecVersion uint16) error {
-	heightBytes := database.PackUInt64(height)
-	rawNestedPublicKeyDiffDB := prefixdb.New(heightBytes, s.nestedValidatorPublicKeyDiffsDB)
-	nestedPKDiffDB := linkeddb.NewDefault(rawNestedPublicKeyDiffDB)
-
+func (s *state) writeCurrentStakers(updateValidators bool, codecVersion uint16) error {
 	for subnetID, validatorDiffs := range s.currentStakers.validatorDiffs {
 		delete(s.currentStakers.validatorDiffs, subnetID)
 
@@ -1986,17 +2148,6 @@ func (s *state) writeCurrentStakers(updateValidators bool, height uint64, codecV
 			delegatorDB = s.currentDelegatorList
 		}
 
-		prefixStruct := heightWithSubnet{
-			Height:   height,
-			SubnetID: subnetID,
-		}
-		prefixBytes, err := block.GenesisCodec.Marshal(block.Version, prefixStruct)
-		if err != nil {
-			return fmt.Errorf("failed to create prefix bytes: %w", err)
-		}
-		rawNestedWeightDiffDB := prefixdb.New(prefixBytes, s.nestedValidatorWeightDiffsDB)
-		nestedWeightDiffDB := linkeddb.NewDefault(rawNestedWeightDiffDB)
-
 		// Record the change in weight and/or public key for each validator.
 		for nodeID, validatorDiff := range validatorDiffs {
 			// Copy [nodeID] so it doesn't get overwritten next iteration.
@@ -2010,21 +2161,6 @@ func (s *state) writeCurrentStakers(updateValidators bool, height uint64, codecV
 				staker := validatorDiff.validator
 				weightDiff.Amount = staker.Weight
 
-				// Invariant: Only the Primary Network contains non-nil public
-				// keys.
-				if staker.PublicKey != nil {
-					// Record that the public key for the validator is being
-					// added. This means the prior value for the public key was
-					// nil.
-					err := s.flatValidatorPublicKeyDiffsDB.Put(
-						marshalDiffKey(constants.PrimaryNetworkID, height, nodeID),
-						nil,
-					)
-					if err != nil {
-						return err
-					}
-				}
-
 				// The validator is being added.
 				//
 				// Invariant: It's impossible for a delegator to have been
@@ -2055,34 +2191,6 @@ func (s *state) writeCurrentStakers(updateValidators bool, height uint64, codecV
 				staker := validatorDiff.validator
 				weightDiff.Amount = staker.Weight
 
-				// Invariant: Only the Primary Network contains non-nil public
-				// keys.
-				if staker.PublicKey != nil {
-					// Record that the public key for the validator is being
-					// removed. This means we must record the prior value of the
-					// public key.
-					//
-					// Note: We store the uncompressed public key here as it is
-					// significantly more efficient to parse when applying
-					// diffs.
-					err := s.flatValidatorPublicKeyDiffsDB.Put(
-						marshalDiffKey(constants.PrimaryNetworkID, height, nodeID),
-						bls.SerializePublicKey(staker.PublicKey),
-					)
-					if err != nil {
-						return err
-					}
-
-					// TODO: Remove this once we no longer support version
-					// rollbacks.
-					//
-					// Note: We store the compressed public key here.
-					pkBytes := bls.PublicKeyToBytes(staker.PublicKey)
-					if err := nestedPKDiffDB.Put(nodeID.Bytes(), pkBytes); err != nil {
-						return err
-					}
-				}
-
 				if err := validatorDB.Delete(staker.TxID[:]); err != nil {
 					return fmt.Errorf("failed to delete current staker: %w", err)
 				}
@@ -2100,48 +2208,28 @@ func (s *state) writeCurrentStakers(updateValidators bool, height uint64, codecV
 				return err
 			}
 
-			if weightDiff.Amount == 0 {
-				// No weight change to record; go to next validator.
-				continue
-			}
-
-			err = s.flatValidatorWeightDiffsDB.Put(
-				marshalDiffKey(subnetID, height, nodeID),
-				marshalWeightDiff(weightDiff),
-			)
-			if err != nil {
-				return err
-			}
-
-			// TODO: Remove this once we no longer support version rollbacks.
-			weightDiffBytes, err := block.GenesisCodec.Marshal(block.Version, weightDiff)
-			if err != nil {
-				return fmt.Errorf("failed to serialize validator weight diff: %w", err)
-			}
-			if err := nestedWeightDiffDB.Put(nodeID.Bytes(), weightDiffBytes); err != nil {
-				return err
-			}
-
 			// TODO: Move the validator set management out of the state package
 			if !updateValidators {
 				continue
 			}
 
-			if weightDiff.Decrease {
+			switch {
+			case weightDiff.Amount == 0:
+				// No weight change to record; go to next validator.
+				continue
+			case weightDiff.Decrease:
 				err = s.validators.RemoveWeight(subnetID, nodeID, weightDiff.Amount)
-			} else {
-				if validatorDiff.validatorStatus == added {
-					staker := validatorDiff.validator
-					err = s.validators.AddStaker(
-						subnetID,
-						nodeID,
-						staker.PublicKey,
-						staker.TxID,
-						weightDiff.Amount,
-					)
-				} else {
-					err = s.validators.AddWeight(subnetID, nodeID, weightDiff.Amount)
-				}
+			case validatorDiff.validatorStatus == added:
+				staker := validatorDiff.validator
+				err = s.validators.AddStaker(
+					subnetID,
+					nodeID,
+					staker.PublicKey,
+					staker.TxID,
+					weightDiff.Amount,
+				)
+			default:
+				err = s.validators.AddWeight(subnetID, nodeID, weightDiff.Amount)
 			}
 			if err != nil {
 				return fmt.Errorf("failed to update validator weight: %w", err)