Make DB Backup More Efficient (#8543)

* checkpoint

* add test

* Update beacon-chain/db/kv/backup_test.go

Co-authored-by: Raul Jordan <raul@prysmaticlabs.com>

beacon-chain/db/kv/backup.go

@@ -47,27 +47,73 @@ func (s *Store) Backup(ctx context.Context, outputDir string) error {
 	copyDB, err := bolt.Open(
 		backupPath,
 		params.BeaconIoConfig().ReadWritePermissions,
-		&bolt.Options{Timeout: params.BeaconIoConfig().BoltTimeout},
+		&bolt.Options{NoFreelistSync: true, NoSync: true, Timeout: params.BeaconIoConfig().BoltTimeout, FreelistType: bolt.FreelistMapType},
 	)
 	if err != nil {
 		return err
 	}
+	copyDB.AllocSize = boltAllocSize
+
 	defer func() {
 		if err := copyDB.Close(); err != nil {
 			log.WithError(err).Error("Failed to close backup database")
 		}
 	}()
-
-	return s.db.View(func(tx *bolt.Tx) error {
+	// Prefetch all keys of buckets, and inner keys in a
+	// bucket to use less memory usage when backing up.
+	bucketKeys := [][]byte{}
+	bucketMap := make(map[string][][]byte)
+	err = s.db.View(func(tx *bolt.Tx) error {
 		return tx.ForEach(func(name []byte, b *bolt.Bucket) error {
-			log.Debugf("Copying bucket %s\n", name)
-			return copyDB.Update(func(tx2 *bolt.Tx) error {
-				b2, err := tx2.CreateBucketIfNotExists(name)
-				if err != nil {
-					return err
+			newName := make([]byte, len(name))
+			copy(newName, name)
+			bucketKeys = append(bucketKeys, newName)
+			innerKeys := [][]byte{}
+			err := b.ForEach(func(k, v []byte) error {
+				if k == nil {
+					return nil
 				}
-				return b.ForEach(b2.Put)
+				nKey := make([]byte, len(k))
+				copy(nKey, k)
+				innerKeys = append(innerKeys, nKey)
+				return nil
 			})
+			if err != nil {
+				return err
+			}
+			bucketMap[string(newName)] = innerKeys
+			return nil
 		})
 	})
+	if err != nil {
+		return err
+	}
+	// Utilize much smaller writes, compared to
+	// writing for a whole bucket in a single transaction. Also
+	// prevent long-running read transactions, as Bolt doesn't
+	// handle those well.
+	for _, k := range bucketKeys {
+		log.Debugf("Copying bucket %s\n", k)
+		innerKeys := bucketMap[string(k)]
+		for _, ik := range innerKeys {
+			err = s.db.View(func(tx *bolt.Tx) error {
+				bkt := tx.Bucket(k)
+				return copyDB.Update(func(tx2 *bolt.Tx) error {
+					b2, err := tx2.CreateBucketIfNotExists(k)
+					if err != nil {
+						return err
+					}
+					return b2.Put(ik, bkt.Get(ik))
+				})
+			})
+			if err != nil {
+				return err
+			}
+		}
+	}
+	// Re-enable sync to allow bolt to fsync
+	// again.
+	copyDB.NoSync = false
+	copyDB.NoFreelistSync = false
+	return nil
 }

beacon-chain/db/kv/backup_test.go

@@ -7,6 +7,7 @@ import (
 	"path/filepath"
 	"testing"
 
+	types "github.com/prysmaticlabs/eth2-types"
 	"github.com/prysmaticlabs/prysm/shared/testutil"
 	"github.com/prysmaticlabs/prysm/shared/testutil/require"
 )
@@ -48,3 +49,58 @@ func TestStore_Backup(t *testing.T) {
 	})
 	require.Equal(t, true, backedDB.HasState(ctx, root))
 }
+
+func TestStore_BackupMultipleBuckets(t *testing.T) {
+	db, err := NewKVStore(context.Background(), t.TempDir(), &Config{})
+	require.NoError(t, err, "Failed to instantiate DB")
+	ctx := context.Background()
+
+	startSlot := types.Slot(5000)
+
+	for i := startSlot; i < 5200; i++ {
+		head := testutil.NewBeaconBlock()
+		head.Block.Slot = i
+		require.NoError(t, db.SaveBlock(ctx, head))
+		root, err := head.Block.HashTreeRoot()
+		require.NoError(t, err)
+		st, err := testutil.NewBeaconState()
+		require.NoError(t, st.SetSlot(i))
+		require.NoError(t, err)
+		require.NoError(t, db.SaveState(ctx, st, root))
+		require.NoError(t, db.SaveHeadBlockRoot(ctx, root))
+	}
+
+	require.NoError(t, db.Backup(ctx, ""))
+
+	backupsPath := filepath.Join(db.databasePath, backupsDirectoryName)
+	files, err := ioutil.ReadDir(backupsPath)
+	require.NoError(t, err)
+	require.NotEqual(t, 0, len(files), "No backups created")
+	require.NoError(t, db.Close(), "Failed to close database")
+
+	oldFilePath := filepath.Join(backupsPath, files[0].Name())
+	newFilePath := filepath.Join(backupsPath, DatabaseFileName)
+	// We rename the file to match the database file name
+	// our NewKVStore function expects when opening a database.
+	require.NoError(t, os.Rename(oldFilePath, newFilePath))
+
+	backedDB, err := NewKVStore(ctx, backupsPath, &Config{})
+	require.NoError(t, err, "Failed to instantiate DB")
+	t.Cleanup(func() {
+		require.NoError(t, backedDB.Close(), "Failed to close database")
+	})
+	for i := startSlot; i < 5200; i++ {
+		head := testutil.NewBeaconBlock()
+		head.Block.Slot = i
+		root, err := head.Block.HashTreeRoot()
+		require.NoError(t, err)
+		nBlock, err := backedDB.Block(ctx, root)
+		require.NoError(t, err)
+		require.NotNil(t, nBlock)
+		require.Equal(t, nBlock.Block.Slot, i)
+		nState, err := backedDB.State(ctx, root)
+		require.NoError(t, err)
+		require.NotNil(t, nState)
+		require.Equal(t, nState.Slot(), i)
+	}
+}