Split validator databases

validator/accounts/account.go

@@ -225,17 +225,26 @@ func HandleEmptyKeystoreFlags(cliCtx *cli.Context, confirmPassword bool) (string
 func Merge(ctx context.Context, sourceDirectories []string, targetDirectory string) (err error) {
 	var sourceStores []*db.Store
 	defer func() {
+		errorMessage := "failed to close one or more source databases"
+		failedToClose := false
 		for _, store := range sourceStores {
 			if deferErr := store.Close(); deferErr != nil {
-				err = errors.Wrapf(deferErr, "Failed to close the database in %s", store.DatabasePath())
+				failedToClose = true
+			}
+		}
+		if failedToClose {
+			if err != nil {
+				err = errors.Wrapf(err, errorMessage)
+			} else {
+				err = errors.New(errorMessage)
 			}
 		}
 	}()
 
 	for _, dir := range sourceDirectories {
 		store, err := db.GetKVStore(dir)
 		if err != nil {
-			return errors.Wrapf(err, "Failed to prepare the database in %s for merging", dir)
+			return errors.Wrapf(err, "failed to prepare the database in %s for merging", dir)
 		}
 		if store == nil {
 			continue
@@ -247,11 +256,34 @@ func Merge(ctx context.Context, sourceDirectories []string, targetDirectory stri
 		return errors.New("no validator databases found in source directories")
 	}
 
-	if err := db.Merge(ctx, sourceStores, targetDirectory); err != nil {
-		return errors.Wrapf(err, "Failed to merge validator databases into %s", targetDirectory)
+	return db.Merge(ctx, sourceStores, targetDirectory)
+}
+
+// Split splits data from one validator database in sourceDirectory into several validator databases.
+// Each validator database is created in its own subdirectory inside targetDirectory.
+func Split(ctx context.Context, sourceDirectory string, targetDirectory string) (err error) {
+	var sourceStore *db.Store
+	sourceStore, err = db.GetKVStore(sourceDirectory)
+	if err != nil {
+		return errors.Wrap(err, "failed to prepare the source database for splitting")
+	}
+	if sourceStore == nil {
+		return errors.New("no database found in source directory")
 	}
+	defer func() {
+		if sourceStore != nil {
+			errorMessage := "failed to close the source database"
+			if deferErr := sourceStore.Close(); deferErr != nil {
+				if err != nil {
+					err = errors.Wrap(err, errorMessage)
+				} else {
+					err = errors.Wrap(deferErr, errorMessage)
+				}
+			}
+		}
+	}()
 
-	return nil
+	return db.Split(ctx, sourceStore, targetDirectory)
 }
 
 // ChangePassword changes the password for all keys located in a keystore.
@@ -276,14 +308,14 @@ func ChangePassword(keystorePath string, oldPassword string, newPassword string)
 func changePasswordForKeyType(keystorePath string, filePrefix string, oldPassword string, newPassword string) error {
 	keys, err := DecryptKeysFromKeystore(keystorePath, filePrefix, oldPassword)
 	if err != nil {
-		return errors.Wrap(err, "Failed to decrypt keys")
+		return errors.Wrap(err, "failed to decrypt keys")
 	}
 
 	keyStore := keystore.NewKeystore(keystorePath)
 	for _, key := range keys {
 		keyFileName := keystorePath + filePrefix + hex.EncodeToString(key.PublicKey.Marshal())[:12]
 		if err := keyStore.StoreKey(keyFileName, key, newPassword); err != nil {
-			return errors.Wrapf(err, "Failed to encrypt key %s with the new password", keyFileName)
+			return errors.Wrapf(err, "failed to encrypt key %s with the new password", keyFileName)
 		}
 	}
 
@@ -305,7 +337,7 @@ func homeDir() string {
 func ExtractPublicKeysFromKeyStore(keystorePath string, passphrase string) ([][]byte, error) {
 	decryptedKeys, err := DecryptKeysFromKeystore(keystorePath, params.BeaconConfig().ValidatorPrivkeyFileName, passphrase)
 	if err != nil {
-		return nil, errors.Wrapf(err, "Could not decrypt keys from keystore in path %s", keystorePath)
+		return nil, errors.Wrapf(err, "could not decrypt keys from keystore in path %s", keystorePath)
 	}
 
 	i := 0

validator/accounts/account_test.go

@@ -243,6 +243,55 @@ func TestMerge_FailsWhenNoDatabaseExistsInAllSourceDirectories(t *testing.T) {
 	}
 }
 
+func TestSplit(t *testing.T) {
+	pubKeys := [][48]byte{{1}, {2}}
+	sourceStore := db.SetupDB(t, pubKeys)
+
+	proposalEpoch := uint64(0)
+	proposalHistory1 := bitfield.Bitlist{0x01, 0x00, 0x00, 0x00, 0x01}
+	if err := sourceStore.SaveProposalHistoryForEpoch(context.Background(), pubKeys[0][:], proposalEpoch, proposalHistory1); err != nil {
+		t.Fatal("Saving proposal history failed")
+	}
+	proposalHistory2 := bitfield.Bitlist{0x02, 0x00, 0x00, 0x00, 0x01}
+	if err := sourceStore.SaveProposalHistoryForEpoch(context.Background(), pubKeys[1][:], proposalEpoch, proposalHistory2); err != nil {
+		t.Fatal("Saving proposal history failed")
+	}
+
+	attestationHistoryMap1 := make(map[uint64]uint64)
+	attestationHistoryMap1[0] = 0
+	pubKeyAttestationHistory1 := &slashpb.AttestationHistory{
+		TargetToSource:     attestationHistoryMap1,
+		LatestEpochWritten: 0,
+	}
+	attestationHistoryMap2 := make(map[uint64]uint64)
+	attestationHistoryMap2[0] = 1
+	pubKeyAttestationHistory2 := &slashpb.AttestationHistory{
+		TargetToSource:     attestationHistoryMap2,
+		LatestEpochWritten: 0,
+	}
+	dbAttestationHistory := make(map[[48]byte]*slashpb.AttestationHistory)
+	dbAttestationHistory[pubKeys[0]] = pubKeyAttestationHistory1
+	dbAttestationHistory[pubKeys[1]] = pubKeyAttestationHistory2
+	if err := sourceStore.SaveAttestationHistoryForPubKeys(context.Background(), dbAttestationHistory); err != nil {
+		t.Fatalf("Saving attestation history failed %v", err)
+	}
+
+	if err := sourceStore.Close(); err != nil {
+		t.Fatalf("Closing source store failed: %v", err)
+	}
+
+	targetDirectory := testutil.TempDir() + "/target"
+	t.Cleanup(func() {
+		if err := os.RemoveAll(targetDirectory); err != nil {
+			t.Errorf("Could not remove target directory : %v", err)
+		}
+	})
+
+	if err := Split(context.Background(), sourceStore.DatabasePath(), targetDirectory); err != nil {
+		t.Fatalf("Splitting failed: %v", err)
+	}
+}
+
 func prepareSourcesForMerging(firstStorePubKey [48]byte, firstStore *db.Store, secondStorePubKey [48]byte, secondStore *db.Store) (*sourceStoresHistory, error) {
 	proposalEpoch := uint64(0)
 	proposalHistory1 := bitfield.Bitlist{0x01, 0x00, 0x00, 0x00, 0x01}

validator/db/BUILD.bazel

@@ -43,5 +43,6 @@ go_test(
         "//shared/testutil:go_default_library",
         "@com_github_pkg_errors//:go_default_library",
         "@com_github_prysmaticlabs_go_bitfield//:go_default_library",
+        "@io_etcd_go_bbolt//:go_default_library",
     ],
 )