fix: improve performance of ApplyPowerTableDiffs in ImportSnapshotToDatastore (#1047)

* fix: improve performance of ApplyPowerTableDiffs in ImportSnapshotToDatastore

* avoid sorting in CertStore.Put

* validate power tables

* Resolve AI comments

Author: hanabi1224

certs/certs.go

@@ -236,15 +236,38 @@ func MakePowerTableDiff(oldPowerTable, newPowerTable gpbft.PowerEntries) PowerTa
 	return diff
 }
 
+func PowerTableArrayToMap(pt gpbft.PowerEntries) map[gpbft.ActorID]gpbft.PowerEntry {
+	ptm := make(map[gpbft.ActorID]gpbft.PowerEntry, len(pt))
+	for _, pe := range pt {
+		ptm[pe.ID] = pe
+	}
+	return ptm
+}
+
+func PowerTableMapToArray(ptm map[gpbft.ActorID]gpbft.PowerEntry) gpbft.PowerEntries {
+	pt := make(gpbft.PowerEntries, 0, len(ptm))
+	for _, pe := range ptm {
+		pt = append(pt, pe)
+	}
+	sort.Sort(pt)
+	return pt
+}
+
 // Apply a set of power table diffs to the passed power table.
 //
 // - The delta must be sorted by participant ID, ascending.
 // - The returned power table is sorted by power, descending.
 func ApplyPowerTableDiffs(prevPowerTable gpbft.PowerEntries, diffs ...PowerTableDiff) (gpbft.PowerEntries, error) {
-	powerTableMap := make(map[gpbft.ActorID]gpbft.PowerEntry, len(prevPowerTable))
-	for _, pe := range prevPowerTable {
-		powerTableMap[pe.ID] = pe
+	powerTableMap := PowerTableArrayToMap(prevPowerTable)
+	powerTableMap, err := ApplyPowerTableDiffsToMap(powerTableMap, diffs...)
+	if err != nil {
+		return nil, err
 	}
+	newPowerTable := PowerTableMapToArray(powerTableMap)
+	return newPowerTable, nil
+}
+
+func ApplyPowerTableDiffsToMap(powerTableMap map[gpbft.ActorID]gpbft.PowerEntry, diffs ...PowerTableDiff) (map[gpbft.ActorID]gpbft.PowerEntry, error) {
 	for j, diff := range diffs {
 		var lastActorId gpbft.ActorID
 		for i, d := range diff {
@@ -298,17 +321,10 @@ func ApplyPowerTableDiffs(prevPowerTable gpbft.PowerEntries, diffs ...PowerTable
 			default: // if the power becomes negative, something went wrong
 				return nil, fmt.Errorf("diff %d resulted in negative power for participant %d", j, pe.ID)
 			}
-
 		}
 	}
 
-	newPowerTable := make(gpbft.PowerEntries, 0, len(powerTableMap))
-	for _, pe := range powerTableMap {
-		newPowerTable = append(newPowerTable, pe)
-	}
-
-	sort.Sort(newPowerTable)
-	return newPowerTable, nil
+	return powerTableMap, nil
 }
 
 // MakePowerTableCID returns the DagCBOR-blake2b256 CID of the given power entries. This method does

certstore/snapshot.go

@@ -19,7 +19,10 @@ import (
 	"golang.org/x/crypto/blake2b"
 )
 
-var ErrUnknownLatestCertificate = errors.New("latest certificate is not known")
+var (
+	ErrUnknownLatestCertificate = errors.New("latest certificate is not known")
+	ErrNoCertificateExtracted   = errors.New("no certificate is found in the snapshot")
+)
 
 // ExportLatestSnapshot exports an F3 snapshot that includes the finality certificate chain until the current `latestCertificate`.
 //
@@ -104,34 +107,79 @@ func importSnapshotToDatastoreWithTestingPowerTableFrequency(ctx context.Context
 	dsb := autobatch.NewAutoBatching(ds, 1000)
 	defer dsb.Flush(ctx)
 	cs, err := OpenOrCreateStore(ctx, dsb, header.FirstInstance, header.InitialPowerTable)
-	if testingPowerTableFrequency > 0 {
-		cs.powerTableFrequency = testingPowerTableFrequency
-	}
 	if err != nil {
 		return err
 	}
-	pt := header.InitialPowerTable
-	for {
+	if testingPowerTableFrequency > 0 {
+		cs.powerTableFrequency = testingPowerTableFrequency
+	}
+	var latestCert *certs.FinalityCertificate
+	ptm := certs.PowerTableArrayToMap(header.InitialPowerTable)
+	for i := header.FirstInstance; ; i += 1 {
 		certBytes, err := readSnapshotBlockBytes(snapshot)
 		if err == io.EOF {
 			break
 		} else if err != nil {
 			return fmt.Errorf("failed to decode finality certificate: %w", err)
 		}
+
 		var cert certs.FinalityCertificate
-		cert.UnmarshalCBOR(bytes.NewReader(certBytes))
-		if err = cs.Put(ctx, &cert); err != nil {
+		if err = cert.UnmarshalCBOR(bytes.NewReader(certBytes)); err != nil {
 			return err
 		}
-		if pt, err = certs.ApplyPowerTableDiffs(pt, cert.PowerTableDelta); err != nil {
+		latestCert = &cert
+
+		if i != cert.GPBFTInstance {
+			return fmt.Errorf("the certificate of instance %d is missing", i)
+		}
+
+		if i > header.LatestInstance {
+			return fmt.Errorf("certificate of instance %d is found, expected latest instance %d", i, header.LatestInstance)
+		}
+
+		if err := cs.ds.Put(ctx, cs.keyForCert(cert.GPBFTInstance), certBytes); err != nil {
 			return err
 		}
+
+		if ptm, err = certs.ApplyPowerTableDiffsToMap(ptm, cert.PowerTableDelta); err != nil {
+			return err
+		}
+
 		if (cert.GPBFTInstance+1)%cs.powerTableFrequency == 0 {
+			pt := certs.PowerTableMapToArray(ptm)
+			if err = checkPowerTable(pt, cert.SupplementalData.PowerTable); err != nil {
+				return err
+			}
 			if err := cs.putPowerTable(ctx, cert.GPBFTInstance+1, pt); err != nil {
 				return err
 			}
 		}
 	}
+
+	if latestCert == nil {
+		return ErrNoCertificateExtracted
+	}
+
+	if latestCert.GPBFTInstance != header.LatestInstance {
+		return fmt.Errorf("extracted latest instance %d, but %d is expected", latestCert.GPBFTInstance, header.LatestInstance)
+	}
+
+	pt := certs.PowerTableMapToArray(ptm)
+	if err = checkPowerTable(pt, latestCert.SupplementalData.PowerTable); err != nil {
+		return err
+	}
+
+	return cs.writeInstanceNumber(ctx, certStoreLatestKey, header.LatestInstance)
+}
+
+func checkPowerTable(pt gpbft.PowerEntries, expectedCid cid.Cid) error {
+	ptCid, err := certs.MakePowerTableCID(pt)
+	if err != nil {
+		return err
+	}
+	if ptCid != expectedCid {
+		return fmt.Errorf("new power table differs from expected power table: %s != %s", ptCid, expectedCid)
+	}
 	return nil
 }