Repository: oasisprotocol/oasis-core (go)

Issue:
Clean stale nodes in checkpointer after aborting a restore
At the moment, `AbortRestore` just clears internal bureaucracy and leaves data from successfully restored chunks in the database. Those should all be cleared, probably by keeping a nearby list of restored hashes that were actually inserted in the db.

Pull request:
go/storage/mkvs: Add node tracking to chunk restore
Closes #3241 

Patch:
diff --git a/.changelog/3241.internal.md b/.changelog/3241.internal.md
new file mode 100644
index 000000000..b986ec66e
--- /dev/null
+++ b/.changelog/3241.internal.md
@@ -0,0 +1,1 @@
+go/storage/mkvs: Add node tracking to chunk restore
diff --git a/go/storage/mkvs/cache.go b/go/storage/mkvs/cache.go
index 9e8bc159c..c8c600872 100644
--- a/go/storage/mkvs/cache.go
+++ b/go/storage/mkvs/cache.go
@@ -419,7 +419,10 @@ func (c *cache) remoteSync(ctx context.Context, ptr *node.Pointer, fetcher readS
 	if c.persistEverythingFromSyncer {
 		// NOTE: This is a dummy batch, we assume that the node database backend is a
 		//       cache-only backend and does not care about correct values.
-		batch = c.db.NewBatch(c.syncRoot, c.syncRoot.Version, false)
+		batch, err = c.db.NewBatch(c.syncRoot, c.syncRoot.Version, false)
+		if err != nil {
+			return fmt.Errorf("mkvs: failed to create batch: %w", err)
+		}
 		dbSubtree = batch.MaybeStartSubtree(nil, 0, subtree)
 	}
 
diff --git a/go/storage/mkvs/checkpoint/checkpoint.go b/go/storage/mkvs/checkpoint/checkpoint.go
index 229a74fe5..2b4a16216 100644
--- a/go/storage/mkvs/checkpoint/checkpoint.go
+++ b/go/storage/mkvs/checkpoint/checkpoint.go
@@ -76,6 +76,8 @@ type Restorer interface {
 	StartRestore(ctx context.Context, checkpoint *Metadata) error
 
 	// AbortRestore aborts a checkpoint restore in progress.
+	//
+	// It is not an error to call this method when no checkpoint restore is in progress.
 	AbortRestore(ctx context.Context) error
 
 	// GetCurrentCheckpoint returns the checkpoint that is being restored. If no restoration is in
diff --git a/go/storage/mkvs/checkpoint/chunk.go b/go/storage/mkvs/checkpoint/chunk.go
index 3f6a7c921..db56792fc 100644
--- a/go/storage/mkvs/checkpoint/chunk.go
+++ b/go/storage/mkvs/checkpoint/chunk.go
@@ -136,7 +136,10 @@ func restoreChunk(ctx context.Context, ndb db.NodeDB, chunk *ChunkMetadata, r io
 	}
 	emptyRoot.Hash.Empty()
 
-	batch := ndb.NewBatch(emptyRoot, chunk.Root.Version, true)
+	batch, err := ndb.NewBatch(emptyRoot, chunk.Root.Version, true)
+	if err != nil {
+		return fmt.Errorf("chunk: failed to create batch: %w", err)
+	}
 	defer batch.Reset()
 
 	subtree := batch.MaybeStartSubtree(nil, 0, ptr)
diff --git a/go/storage/mkvs/checkpoint/restorer.go b/go/storage/mkvs/checkpoint/restorer.go
index e31a97fbd..e680a4f73 100644
--- a/go/storage/mkvs/checkpoint/restorer.go
+++ b/go/storage/mkvs/checkpoint/restorer.go
@@ -31,6 +31,10 @@ func (rs *restorer) StartRestore(ctx context.Context, checkpoint *Metadata) erro
 		return ErrRestoreAlreadyInProgress
 	}
 
+	if err := rs.ndb.StartMultipartInsert(checkpoint.Root.Version); err != nil {
+		return err
+	}
+
 	rs.currentCheckpoint = checkpoint
 	rs.pendingChunks = make(map[uint64]bool)
 	for idx := range checkpoint.Chunks {
@@ -44,16 +48,10 @@ func (rs *restorer) AbortRestore(ctx context.Context) error {
 	rs.Lock()
 	defer rs.Unlock()
 
-	if rs.currentCheckpoint == nil {
-		return ErrNoRestoreInProgress
-	}
-
 	rs.pendingChunks = nil
 	rs.currentCheckpoint = nil
 
-	// TODO: also properly remove leftover nodes from the database (#3241).
-
-	return nil
+	return rs.ndb.AbortMultipartInsert()
 }
 
 func (rs *restorer) GetCurrentCheckpoint() *Metadata {
@@ -95,12 +93,7 @@ func (rs *restorer) RestoreChunk(ctx context.Context, idx uint64, r io.Reader) (
 	case errors.Is(err, ErrChunkProofVerificationFailed):
 		// Chunk was as specified in the manifest but did not match the reported root. In this case
 		// we need to abort processing the given checkpoint.
-		rs.Lock()
-		defer rs.Unlock()
-
-		rs.pendingChunks = nil
-		rs.currentCheckpoint = nil
-
+		_ = rs.AbortRestore(ctx)
 		return false, err
 	default:
 		return false, err
@@ -116,7 +109,6 @@ func (rs *restorer) RestoreChunk(ctx context.Context, idx uint64, r io.Reader) (
 	if len(rs.pendingChunks) == 0 {
 		rs.pendingChunks = nil
 		rs.currentCheckpoint = nil
-
 		return true, nil
 	}
 
diff --git a/go/storage/mkvs/commit.go b/go/storage/mkvs/commit.go
index dc3d26a16..9c750d697 100644
--- a/go/storage/mkvs/commit.go
+++ b/go/storage/mkvs/commit.go
@@ -68,13 +68,17 @@ func (t *tree) commitWithHooks(
 	}
 
 	var batch db.Batch
+	var err error
 	switch opts.noPersist {
 	case false:
-		batch = t.cache.db.NewBatch(oldRoot, version, false)
+		batch, err = t.cache.db.NewBatch(oldRoot, version, false)
 	case true:
 		// Do not persist anything -- use a dummy batch.
 		nopDb, _ := db.NewNopNodeDB()
-		batch = nopDb.NewBatch(oldRoot, version, false)
+		batch, err = nopDb.NewBatch(oldRoot, version, false)
+	}
+	if err != nil {
+		return nil, hash.Hash{}, err
 	}
 	defer batch.Reset()
 
diff --git a/go/storage/mkvs/db/api/api.go b/go/storage/mkvs/db/api/api.go
index c86e27007..b1a282c0a 100644
--- a/go/storage/mkvs/db/api/api.go
+++ b/go/storage/mkvs/db/api/api.go
@@ -45,6 +45,12 @@ var (
 	ErrNotEarliest = errors.New(ModuleName, 11, "mkvs: version is not the earliest version")
 	// ErrReadOnly indicates that a write operation failed due to a read-only database.
 	ErrReadOnly = errors.New(ModuleName, 12, "mkvs: read-only database")
+	// ErrMultipartInProgress indicates that a multipart restore operation is already
+	// in progress.
+	ErrMultipartInProgress = errors.New(ModuleName, 13, "mkvs: multipart already in progress")
+	// ErrInvalidMultipartVersion indicates that a Finalize, NewBatch or Commit was called with a version
+	// that doesn't match the current multipart restore as set with StartMultipartRestore.
+	ErrInvalidMultipartVersion = errors.New(ModuleName, 14, "mkvs: operation called with different version than current multipart version")
 )
 
 // Config is the node database backend configuration.
@@ -73,7 +79,7 @@ type Config struct { // nolint: maligned
 
 // NodeDB is the persistence layer used for persisting the in-memory tree.
 type NodeDB interface {
-	// GetNode lookups up a node in the database.
+	// GetNode looks up a node in the database.
 	GetNode(root node.Root, ptr *node.Pointer) (node.Node, error)
 
 	// GetWriteLog retrieves a write log between two storage instances from the database.
@@ -88,13 +94,25 @@ type NodeDB interface {
 	// GetRootsForVersion returns a list of roots stored under the given version.
 	GetRootsForVersion(ctx context.Context, version uint64) ([]hash.Hash, error)
 
+	// StartMultipartInsert prepares the database for a batch insert job from multiple chunks.
+	// Batches from this call onwards will keep track of inserted nodes so that they can be
+	// deleted if the job fails for any reason.
+	StartMultipartInsert(version uint64) error
+
+	// AbortMultipartInsert cleans up the node insertion log that was kept since the last
+	// StartMultipartInsert operation. The log will be cleared and the associated nodes can
+	// be either removed (if the insertion failed) or left intact (if it was successful).
+	//
+	// It is not an error to call this method more than once.
+	AbortMultipartInsert() error
+
 	// NewBatch starts a new batch.
 	//
 	// The chunk argument specifies whether the given batch is being used to import a chunk of an
 	// existing root. Chunks may contain unresolved pointers (e.g., pointers that point to hashes
 	// which are not present in the database). Committing a chunk batch will prevent the version
 	// from being finalized.
-	NewBatch(oldRoot node.Root, version uint64, chunk bool) Batch
+	NewBatch(oldRoot node.Root, version uint64, chunk bool) (Batch, error)
 
 	// HasRoot checks whether the given root exists.
 	HasRoot(root node.Root) bool
@@ -213,6 +231,14 @@ func (d *nopNodeDB) HasRoot(root node.Root) bool {
 	return false
 }
 
+func (d *nopNodeDB) StartMultipartInsert(version uint64) error {
+	return nil
+}
+
+func (d *nopNodeDB) AbortMultipartInsert() error {
+	return nil
+}
+
 func (d *nopNodeDB) Finalize(ctx context.Context, version uint64, roots []hash.Hash) error {
 	return nil
 }
@@ -237,8 +263,8 @@ type nopBatch struct {
 	BaseBatch
 }
 
-func (d *nopNodeDB) NewBatch(oldRoot node.Root, version uint64, chunk bool) Batch {
-	return &nopBatch{}
+func (d *nopNodeDB) NewBatch(oldRoot node.Root, version uint64, chunk bool) (Batch, error) {
+	return &nopBatch{}, nil
 }
 
 func (b *nopBatch) MaybeStartSubtree(subtree Subtree, depth node.Depth, subtreeRoot *node.Pointer) Subtree {
diff --git a/go/storage/mkvs/db/badger/badger.go b/go/storage/mkvs/db/badger/badger.go
index 1ec27d076..241b1e6d8 100644
--- a/go/storage/mkvs/db/badger/badger.go
+++ b/go/storage/mkvs/db/badger/badger.go
@@ -3,6 +3,7 @@ package badger
 
 import (
 	"context"
+	"errors"
 	"fmt"
 	"sync"
 
@@ -20,7 +21,13 @@ import (
 	"github.com/oasisprotocol/oasis-core/go/storage/mkvs/writelog"
 )
 
-const dbVersion = 3
+const (
+	dbVersion = 3
+
+	// multipartVersionNone is the value used for the multipart version in metadata
+	// when no multipart restore is in progress.
+	multipartVersionNone uint64 = 0
+)
 
 var (
 	// nodeKeyFmt is the key format for nodes (node hash).
@@ -46,6 +53,13 @@ var (
 	//
 	// Value is CBOR-serialized metadata.
 	metadataKeyFmt = keyformat.New(0x04)
+	// multipartRestoreNodeLogKeyFmt is the key format for the nodes inserted during a chunk restore.
+	// Once a set of chunks is fully restored, these entries should be removed. If chunk restoration
+	// is interrupted for any reason, the nodes associated with these keys should be removed, along
+	// with these entries.
+	//
+	// Value is empty.
+	multipartRestoreNodeLogKeyFmt = keyformat.New(0x05, &hash.Hash{})
 )
 
 // New creates a new BadgerDB-backed node database.
@@ -87,6 +101,12 @@ func New(cfg *api.Config) (api.NodeDB, error) {
 		return nil, fmt.Errorf("mkvs/badger: failed to load metadata: %w", err)
 	}
 
+	// Cleanup any multipart restore remnants, since they can't be used anymore.
+	if err = db.cleanMultipartLocked(true); err != nil {
+		_ = db.db.Close()
+		return nil, fmt.Errorf("mkvs/badger: failed to clean leftovers from multipart restore: %w", err)
+	}
+
 	db.gc = cmnBadger.NewGCWorker(db.logger, db.db)
 
 	return db, nil
@@ -100,6 +120,8 @@ type badgerNodeDB struct { // nolint: maligned
 	readOnly         bool
 	discardWriteLogs bool
 
+	multipartVersion uint64
+
 	db *badger.DB
 	gc *cmnBadger.GCWorker
 
@@ -141,6 +163,7 @@ func (d *badgerNodeDB) load() error {
 				d.meta.value.Namespace,
 			)
 		}
+		return nil
 	case badger.ErrKeyNotFound:
 	default:
 		return err
@@ -163,6 +186,76 @@ func (d *badgerNodeDB) sanityCheckNamespace(ns common.Namespace) error {
 	return nil
 }
 
+// Assumes metaUpdateLock is held when called.
+func (d *badgerNodeDB) cleanMultipartLocked(removeNodes bool) error {
+	var version uint64
+
+	if d.multipartVersion != multipartVersionNone {
+		version = d.multipartVersion
+	} else {
+		version = d.meta.getMultipartVersion()
+	}
+	if version == multipartVersionNone {
+		// No multipart in progress, but it's not an error to call in a situation like this.
+		return nil
+	}
+
+	txn := d.db.NewTransactionAt(tsMetadata, false)
+	defer txn.Discard()
+
+	opts := badger.DefaultIteratorOptions
+	opts.Prefix = multipartRestoreNodeLogKeyFmt.Encode()
+	it := txn.NewIterator(opts)
+	defer it.Close()
+
+	metaBatch := d.db.NewWriteBatchAt(tsMetadata)
+	defer metaBatch.Cancel()
+	nodeBatch := d.db.NewWriteBatchAt(versionToTs(version))
+	defer nodeBatch.Cancel()
+
+	var logged bool
+	for it.Rewind(); it.Valid(); it.Next() {
+		key := it.Item().Key()
+		if removeNodes {
+			if !logged {
+				d.logger.Info("removing some nodes from a multipart restore")
+				logged = true
+			}
+			var hash hash.Hash
+			if !multipartRestoreNodeLogKeyFmt.Decode(key, &hash) {
+				panic("mkvs/badger: bad iterator")
+			}
+			if err := nodeBatch.Delete(nodeKeyFmt.Encode(&hash)); err != nil {
+				return err
+			}
+		}
+		if err := metaBatch.Delete(key); err != nil {
+			return err
+		}
+	}
+
+	// Flush both batches first. If anything fails, having corrupt
+	// multipart info in d.meta shouldn't hurt us next run.
+	if err := nodeBatch.Flush(); err != nil {
+		return err
+	}
+	if err := metaBatch.Flush(); err != nil {
+		return err
+	}
+
+	metaTx := d.db.NewTransactionAt(tsMetadata, true)
+	defer metaTx.Discard()
+	if err := d.meta.setMultipartVersion(metaTx, 0); err != nil {
+		return err
+	}
+	if err := metaTx.CommitAt(tsMetadata, nil); err != nil {
+		return err
+	}
+
+	d.multipartVersion = multipartVersionNone
+	return nil
+}
+
 func (d *badgerNodeDB) GetNode(root node.Root, ptr *node.Pointer) (node.Node, error) {
 	if ptr == nil || !ptr.IsClean() {
 		panic("mkvs/badger: attempted to get invalid pointer from node database")
@@ -414,6 +507,10 @@ func (d *badgerNodeDB) Finalize(ctx context.Context, version uint64, roots []has
 	d.metaUpdateLock.Lock()
 	defer d.metaUpdateLock.Unlock()
 
+	if d.multipartVersion != multipartVersionNone && d.multipartVersion != version {
+		return api.ErrInvalidMultipartVersion
+	}
+
 	// Version batch collects removals at the version timestamp.
 	versionBatch := d.db.NewWriteBatchAt(versionToTs(version))
 	defer versionBatch.Cancel()
@@ -563,6 +660,13 @@ func (d *badgerNodeDB) Finalize(ctx context.Context, version uint64, roots []has
 	if err := tx.CommitAt(tsMetadata, nil); err != nil {
 		return fmt.Errorf("mkvs/badger: failed to commit metadata: %w", err)
 	}
+
+	// Clean multipart metadata if there is any.
+	if d.multipartVersion != multipartVersionNone {
+		if err := d.cleanMultipartLocked(false); err != nil {
+			return err
+		}
+	}
 	return nil
 }
 
@@ -574,6 +678,10 @@ func (d *badgerNodeDB) Prune(ctx context.Context, version uint64) error {
 	d.metaUpdateLock.Lock()
 	defer d.metaUpdateLock.Unlock()
 
+	if d.multipartVersion != multipartVersionNone {
+		return api.ErrMultipartInProgress
+	}
+
 	// Make sure that the version that we try to prune has been finalized.
 	lastFinalizedVersion, exists := d.meta.getLastFinalizedVersion()
 	if !exists || lastFinalizedVersion < version {
@@ -672,7 +780,45 @@ func (d *badgerNodeDB) Prune(ctx context.Context, version uint64) error {
 	return nil
 }
 
-func (d *badgerNodeDB) NewBatch(oldRoot node.Root, version uint64, chunk bool) api.Batch {
+func (d *badgerNodeDB) StartMultipartInsert(version uint64) error {
+	d.metaUpdateLock.Lock()
+	defer d.metaUpdateLock.Unlock()
+
+	if version == multipartVersionNone {
+		return api.ErrInvalidMultipartVersion
+	}
+
+	if d.multipartVersion != multipartVersionNone {
+		if d.multipartVersion != version {
+			return api.ErrMultipartInProgress
+		}
+		// Multipart already initialized at the same version, so this was
+		// probably called e.g. as part of a further checkpoint restore.
+		return nil
+	}
+
+	tx := d.db.NewTransactionAt(tsMetadata, true)
+	defer tx.Discard()
+	if err := d.meta.setMultipartVersion(tx, version); err != nil {
+		return err
+	}
+	if err := tx.CommitAt(tsMetadata, nil); err != nil {
+		return err
+	}
+
+	d.multipartVersion = version
+
+	return nil
+}
+
+func (d *badgerNodeDB) AbortMultipartInsert() error {
+	d.metaUpdateLock.Lock()
+	defer d.metaUpdateLock.Unlock()
+
+	return d.cleanMultipartLocked(true)
+}
+
+func (d *badgerNodeDB) NewBatch(oldRoot node.Root, version uint64, chunk bool) (api.Batch, error) {
 	// WARNING: There is a maximum batch size and maximum batch entry count.
 	// Both of these things are derived from the MaxTableSize option.
 	//
@@ -680,12 +826,37 @@ func (d *badgerNodeDB) NewBatch(oldRoot node.Root, version uint64, chunk bool) a
 	// thing to do would be to either crank up MaxTableSize or maybe split
 	// the transaction out.
 
-	return &badgerBatch{
-		db:      d,
-		bat:     d.db.NewWriteBatchAt(versionToTs(version)),
-		oldRoot: oldRoot,
-		chunk:   chunk,
+	if d.readOnly {
+		return nil, api.ErrReadOnly
+	}
+
+	d.metaUpdateLock.Lock()
+	defer d.metaUpdateLock.Unlock()
+
+	if d.multipartVersion != multipartVersionNone && d.multipartVersion != version {
+		return nil, api.ErrInvalidMultipartVersion
 	}
+	if chunk != (d.multipartVersion != multipartVersionNone) {
+		return nil, api.ErrMultipartInProgress
+	}
+
+	var logBatch *badger.WriteBatch
+	var readTxn *badger.Txn
+	if d.multipartVersion != multipartVersionNone {
+		// The node log is at a different version than the nodes themselves,
+		// which is awkward.
+		logBatch = d.db.NewWriteBatchAt(tsMetadata)
+		readTxn = d.db.NewTransactionAt(versionToTs(version), false)
+	}
+
+	return &badgerBatch{
+		db:             d,
+		bat:            d.db.NewWriteBatchAt(versionToTs(version)),
+		multipartNodes: logBatch,
+		readTxn:        readTxn,
+		oldRoot:        oldRoot,
+		chunk:          chunk,
+	}, nil
 }
 
 func (d *badgerNodeDB) Size() (int64, error) {
@@ -712,8 +883,13 @@ func (d *badgerNodeDB) Close() {
 type badgerBatch struct {
 	api.BaseBatch
 
-	db  *badgerNodeDB
-	bat *badger.WriteBatch
+	db             *badgerNodeDB
+	bat            *badger.WriteBatch
+	multipartNodes *badger.WriteBatch
+
+	// readTx is the read transaction used to check for node existence during
+	// a multipart restore.
+	readTxn *badger.Txn
 
 	oldRoot node.Root
 	chunk   bool
@@ -758,14 +934,13 @@ func (ba *badgerBatch) RemoveNodes(nodes []node.Node) error {
 }
 
 func (ba *badgerBatch) Commit(root node.Root) error {
-	// XXX: Ideally this would fail at batch creation.
-	if ba.db.readOnly {
-		return api.ErrReadOnly
-	}
-
 	ba.db.metaUpdateLock.Lock()
 	defer ba.db.metaUpdateLock.Unlock()
 
+	if ba.db.multipartVersion != multipartVersionNone && ba.db.multipartVersion != root.Version {
+		return api.ErrInvalidMultipartVersion
+	}
+
 	if err := ba.db.sanityCheckNamespace(root.Namespace); err != nil {
 		return err
 	}
@@ -853,6 +1028,11 @@ func (ba *badgerBatch) Commit(root node.Root) error {
 	}
 
 	// Flush node updates.
+	if ba.multipartNodes != nil {
+		if err = ba.multipartNodes.Flush(); err != nil {
+			return fmt.Errorf("mkvs/badger: failed to flush node log batch: %w", err)
+		}
+	}
 	if err = ba.bat.Flush(); err != nil {
 		return fmt.Errorf("mkvs/badger: failed to flush batch: %w", err)
 	}
@@ -871,6 +1051,10 @@ func (ba *badgerBatch) Commit(root node.Root) error {
 
 func (ba *badgerBatch) Reset() {
 	ba.bat.Cancel()
+	if ba.multipartNodes != nil {
+		ba.multipartNodes.Cancel()
+		ba.readTxn.Discard()
+	}
 	ba.writeLog = nil
 	ba.annotations = nil
 	ba.updatedNodes = nil
@@ -888,7 +1072,15 @@ func (s *badgerSubtree) PutNode(depth node.Depth, ptr *node.Pointer) error {
 
 	h := ptr.Node.GetHash()
 	s.batch.updatedNodes = append(s.batch.updatedNodes, updatedNode{Hash: h})
-	if err = s.batch.bat.Set(nodeKeyFmt.Encode(&h), data); err != nil {
+	nodeKey := nodeKeyFmt.Encode(&h)
+	if s.batch.multipartNodes != nil {
+		if _, err = s.batch.readTxn.Get(nodeKey); err != nil && errors.Is(err, badger.ErrKeyNotFound) {
+			if err = s.batch.multipartNodes.Set(multipartRestoreNodeLogKeyFmt.Encode(&h), []byte{}); err != nil {
+				return err
+			}
+		}
+	}
+	if err = s.batch.bat.Set(nodeKey, data); err != nil {
 		return err
 	}
 	return nil
diff --git a/go/storage/mkvs/db/badger/metadata.go b/go/storage/mkvs/db/badger/metadata.go
index 172936c22..8f4186fc4 100644
--- a/go/storage/mkvs/db/badger/metadata.go
+++ b/go/storage/mkvs/db/badger/metadata.go
@@ -22,6 +22,8 @@ type serializedMetadata struct {
 	EarliestVersion uint64 `json:"earliest_version"`
 	// LastFinalizedVersion is the last finalized version.
 	LastFinalizedVersion *uint64 `json:"last_finalized_version"`
+	// MultipartVersion is the version for the in-progress multipart restore, or 0 if none was in progress.
+	MultipartVersion uint64 `json:"multipart_version"`
 }
 
 // metadata is the database metadata.
@@ -77,6 +79,21 @@ func (m *metadata) setLastFinalizedVersion(tx *badger.Txn, version uint64) error
 	return m.save(tx)
 }
 
+func (m *metadata) getMultipartVersion() uint64 {
+	m.Lock()
+	defer m.Unlock()
+
+	return m.value.MultipartVersion
+}
+
+func (m *metadata) setMultipartVersion(tx *badger.Txn, version uint64) error {
+	m.Lock()
+	defer m.Unlock()
+
+	m.value.MultipartVersion = version
+	return m.save(tx)
+}
+
 func (m *metadata) save(tx *badger.Txn) error {
 	return tx.Set(metadataKeyFmt.Encode(), cbor.Marshal(m.value))
 }
diff --git a/go/storage/mkvs/interop/cmd/protocol_server.go b/go/storage/mkvs/interop/cmd/protocol_server.go
index d2a855835..b182fba64 100644
--- a/go/storage/mkvs/interop/cmd/protocol_server.go
+++ b/go/storage/mkvs/interop/cmd/protocol_server.go
@@ -1,6 +1,11 @@
 package cmd
 
 import (
+	"fmt"
+	"io"
+	"os"
+	"path/filepath"
+
 	"github.com/spf13/cobra"
 	flag "github.com/spf13/pflag"
 	"github.com/spf13/viper"
@@ -38,7 +43,20 @@ func doProtoServer(cmd *cobra.Command, args []string) {
 
 	dataDir := viper.GetString(cfgServerDataDir)
 	if dataDir == "" {
-		logger.Error("no data directory specified")
+		fmt.Println("no data directory specified")
+		return
+	}
+
+	// Initialize logging.
+	logFile := filepath.Join(dataDir, "protocol_server.log")
+	w, err := os.OpenFile(logFile, os.O_CREATE|os.O_APPEND|os.O_WRONLY, 0o600)
+	if err != nil {
+		fmt.Printf("failed to open log file: %v\n", err)
+		return
+	}
+	logWriter := io.MultiWriter(os.Stdout, w)
+	if err = logging.Initialize(logWriter, logging.FmtJSON, logging.LevelWarn, nil); err != nil {
+		fmt.Printf("failed to initialize logging: %v\n", err)
 		return
 	}
 
diff --git a/go/worker/storage/committee/checkpoint_sync.go b/go/worker/storage/committee/checkpoint_sync.go
index 21d73d668..1e0a3fe32 100644
--- a/go/worker/storage/committee/checkpoint_sync.go
+++ b/go/worker/storage/committee/checkpoint_sync.go
@@ -13,6 +13,7 @@ import (
 
 	"github.com/cenkalti/backoff/v4"
 
+	"github.com/oasisprotocol/oasis-core/go/common/crypto/hash"
 	"github.com/oasisprotocol/oasis-core/go/runtime/committee"
 	schedulerApi "github.com/oasisprotocol/oasis-core/go/scheduler/api"
 	storageApi "github.com/oasisprotocol/oasis-core/go/storage/api"
@@ -207,20 +208,9 @@ func (n *Node) handleCheckpoint(check *checkpoint.Metadata, committeeClient comm
 	defer cancel()
 
 	err = n.localStorage.Checkpointer().StartRestore(n.ctx, check)
-	switch {
-	case err == nil:
-		break
-	case errors.Is(err, checkpoint.ErrRestoreAlreadyInProgress):
-		n.logger.Info("a restore is already in progress, trying to abort")
-		err = n.localStorage.Checkpointer().AbortRestore(n.ctx)
-		if err != nil {
-			return checkpointStatusBail, fmt.Errorf("can't abort previous checkpoint restore: %w", err)
-		}
-		err = n.localStorage.Checkpointer().StartRestore(n.ctx, check)
-		if err != nil {
-			return checkpointStatusBail, fmt.Errorf("can't start checkpoint restore after abort: %w", err)
-		}
-	default:
+	if err != nil {
+		// Any previous restores were already aborted by the driver up the call stack, so
+		// things should have been going smoothly here; bail.
 		return checkpointStatusBail, fmt.Errorf("can't start checkpoint restore: %w", err)
 	}
 
@@ -385,20 +375,12 @@ func (n *Node) checkCheckpointUsable(cp *checkpoint.Metadata, remainingMask outs
 	return maskNone
 }
 
-func (n *Node) updateLastSynced(check *checkpoint.Metadata, mask outstandingMask) {
-	n.syncedLock.Lock()
-	defer n.syncedLock.Unlock()
+func (n *Node) syncCheckpoints() (*blockSummary, error) {
+	// Store roots and round info for checkpoints that finished syncing.
+	// Round and namespace info will get overwritten as rounds are skipped
+	// for errors, driven by remainingRoots.
+	var syncState blockSummary
 
-	n.syncedState.LastBlock.Round = check.Root.Version
-	n.syncedState.LastBlock.Namespace = check.Root.Namespace
-	if mask == maskIO {
-		n.syncedState.LastBlock.IORoot = check.Root
-	} else if mask == maskState {
-		n.syncedState.LastBlock.StateRoot = check.Root
-	}
-}
-
-func (n *Node) syncCheckpoints() error {
 	// Start following the storage committee.
 	committeeWatcher, err := committee.NewWatcher(
 		n.ctx,
@@ -410,41 +392,73 @@ func (n *Node) syncCheckpoints() error {
 		committee.WithFilter(committee.IgnoreNodeFilter(n.commonNode.Identity.NodeSigner.Public())),
 	)
 	if err != nil {
-		return fmt.Errorf("can't establish storage committee watcher: %w", err)
+		return nil, fmt.Errorf("can't establish storage committee watcher: %w", err)
 	}
 	committeeClient, err := committee.NewClient(n.ctx, committeeWatcher.Nodes(), committee.WithClientAuthentication(n.commonNode.Identity))
 	if err != nil {
-		return fmt.Errorf("can't create committee client: %w", err)
+		return nil, fmt.Errorf("can't create committee client: %w", err)
 	}
 
 	descriptor, err := n.commonNode.Runtime.RegistryDescriptor(n.ctx)
 	if err != nil {
-		return fmt.Errorf("can't get runtime descriptor: %w", err)
+		return nil, fmt.Errorf("can't get runtime descriptor: %w", err)
 	}
 
 	// Fetch metadata from the current committee.
 	metadata, err := n.getCheckpointList(committeeClient)
 	if err != nil {
-		return fmt.Errorf("can't get checkpoint list from storage committee: %w", err)
+		return nil, fmt.Errorf("can't get checkpoint list from storage committee: %w", err)
 	}
 
 	// Try all the checkpoints now, from most recent backwards.
-	remainingRoots := maskAll
+	var prevVersion uint64
 	var mask outstandingMask
+	var doneRoots []hash.Hash
+	remainingRoots := maskAll
 	for _, check := range metadata {
 		mask = n.checkCheckpointUsable(check, remainingRoots)
 		if mask == maskNone {
 			continue
 		}
 
+		if check.Root.Version != prevVersion {
+			// Kill any previous restores that might be active. This should kill
+			// the restorer's state as well as the underlying DB multipart bookkeeping.
+			if err := n.localStorage.Checkpointer().AbortRestore(n.ctx); err != nil {
+				return nil, fmt.Errorf("error aborting previous restore for checkpoint sync: %w", err)
+			}
+			remainingRoots = maskAll
+			prevVersion = check.Root.Version
+			doneRoots = []hash.Hash{}
+		}
+
 		status, err := n.handleCheckpoint(check, committeeClient, descriptor.Storage.GroupSize)
 		switch status {
 		case checkpointStatusDone:
 			n.logger.Info("successfully restored from checkpoint", "root", check.Root, "mask", mask)
-			n.updateLastSynced(check, mask)
+
+			syncState.Namespace = check.Root.Namespace
+			syncState.Round = check.Root.Version
+			switch mask {
+			case maskIO:
+				syncState.IORoot = check.Root
+			case maskState:
+				syncState.StateRoot = check.Root
+			}
+
+			doneRoots = append(doneRoots, check.Root.Hash)
 			remainingRoots &= ^mask
 			if remainingRoots == maskNone {
-				return nil
+				if err = n.localStorage.NodeDB().Finalize(n.ctx, prevVersion, doneRoots); err != nil {
+					n.logger.Error("can't finalize version after all checkpoints restored",
+						"err", err,
+						"version", prevVersion,
+						"roots", doneRoots,
+					)
+					// Likely a local problem, so just bail.
+					return nil, fmt.Errorf("can't finalize version after checkpoints restored: %w", err)
+				}
+				return &syncState, nil
 			}
 			continue
 		case checkpointStatusNext:
@@ -452,9 +466,9 @@ func (n *Node) syncCheckpoints() error {
 			continue
 		case checkpointStatusBail:
 			n.logger.Error("error trying to restore from checkpoint, unrecoverable", "root", check.Root, "err", err)
-			return fmt.Errorf("error restoring from checkpoints: %w", err)
+			return nil, fmt.Errorf("error restoring from checkpoints: %w", err)
 		}
 	}
 
-	return ErrNoUsableCheckpoints
+	return nil, ErrNoUsableCheckpoints
 }
diff --git a/go/worker/storage/committee/node.go b/go/worker/storage/committee/node.go
index 9f7eb8655..36e6aa5b3 100644
--- a/go/worker/storage/committee/node.go
+++ b/go/worker/storage/committee/node.go
@@ -604,6 +604,19 @@ func (n *Node) initGenesis(rt *registryApi.Runtime) error {
 	return nil
 }
 
+func (n *Node) flushSyncedState(summary *blockSummary) uint64 {
+	n.syncedLock.Lock()
+	defer n.syncedLock.Unlock()
+
+	n.syncedState.LastBlock = *summary
+	rtID := n.commonNode.Runtime.ID()
+	if err := n.stateStore.PutCBOR(rtID[:], &n.syncedState); err != nil {
+		n.logger.Error("can't store watcher state to database", "err", err)
+	}
+
+	return n.syncedState.LastBlock.Round
+}
+
 func (n *Node) worker() { // nolint: gocyclo
 	defer close(n.quitCh)
 	defer close(n.diffCh)
@@ -686,17 +699,20 @@ func (n *Node) worker() { // nolint: gocyclo
 
 	// Try to perform initial sync from state and io checkpoints.
 	if !n.checkpointSyncDisabled {
-		err = n.syncCheckpoints()
+		var summary *blockSummary
+		summary, err = n.syncCheckpoints()
 		if err != nil {
 			// Try syncing again. The main reason for this is the sync failing due to a
 			// checkpoint pruning race condition (where nodes list a checkpoint which is
 			// then deleted just before we request its chunks). One retry is enough.
 			n.logger.Info("first checkpoint sync failed, trying once more", "err", err)
-			err = n.syncCheckpoints()
+			summary, err = n.syncCheckpoints()
 		}
 		if err != nil {
 			n.logger.Info("checkpoint sync failed", "err", err)
 		} else {
+			cachedLastRound = n.flushSyncedState(summary)
+			lastFullyAppliedRound = cachedLastRound
 			n.logger.Info("checkpoint sync succeeded",
 				logging.LogEvent, LogEventCheckpointSyncSuccess,
 			)
@@ -893,18 +909,7 @@ mainLoop:
 		case finalized := <-n.finalizeCh:
 			// No further sync or out of order handling needed here, since
 			// only one finalize at a time is triggered (for round cachedLastRound+1)
-			n.syncedLock.Lock()
-			n.syncedState.LastBlock.Round = finalized.Round
-			n.syncedState.LastBlock.IORoot = finalized.IORoot
-			n.syncedState.LastBlock.StateRoot = finalized.StateRoot
-			rtID := n.commonNode.Runtime.ID()
-			err = n.stateStore.PutCBOR(rtID[:], &n.syncedState)
-			n.syncedLock.Unlock()
-			cachedLastRound = finalized.Round
-			if err != nil {
-				n.logger.Error("can't store watcher state to database", "err", err)
-			}
-
+			cachedLastRound = n.flushSyncedState(finalized)
 			storageWorkerLastFullRound.With(n.getMetricLabels()).Set(float64(finalized.Round))
 
 			// Notify the checkpointer that there is a new finalized round.
diff --git a/runtime/src/storage/mkvs/interop/rpc.rs b/runtime/src/storage/mkvs/interop/rpc.rs
index 5ca200193..cfaa19a68 100644
--- a/runtime/src/storage/mkvs/interop/rpc.rs
+++ b/runtime/src/storage/mkvs/interop/rpc.rs
@@ -1,3 +1,5 @@
+use std::time::Duration;
+
 use grpcio::{CallOption, Channel, Client, Error, Marshaller, Method, MethodType, Result};
 use serde::{de::DeserializeOwned, Deserialize, Serialize};
 use serde_cbor::Value;
@@ -61,6 +63,9 @@ const METHOD_SYNC_ITERATE: Method<sync::IterateRequest, sync::ProofResponse> = M
     },
 };
 
+// Calls should still have a timeout to handle the case where the interop server exits prematurely.
+const CALL_TIMEOUT: Duration = Duration::from_secs(30);
+
 #[derive(Clone, Debug, Serialize, Deserialize)]
 pub struct ApplyRequest {
     pub namespace: Namespace,
@@ -94,7 +99,9 @@ impl StorageClient {
             .unary_call(
                 &METHOD_APPLY,
                 &request,
-                CallOption::default().wait_for_ready(true),
+                CallOption::default()
+                    .wait_for_ready(true)
+                    .timeout(CALL_TIMEOUT),
             )
             .map(|_| ())
     }
@@ -103,7 +110,9 @@ impl StorageClient {
         self.client.unary_call(
             &METHOD_SYNC_GET,
             &request,
-            CallOption::default().wait_for_ready(true),
+            CallOption::default()
+                .wait_for_ready(true)
+                .timeout(CALL_TIMEOUT),
         )
     }
 
@@ -114,7 +123,9 @@ impl StorageClient {
         self.client.unary_call(
             &METHOD_SYNC_GET_PREFIXES,
             &request,
-            CallOption::default().wait_for_ready(true),
+            CallOption::default()
+                .wait_for_ready(true)
+                .timeout(CALL_TIMEOUT),
         )
     }
 
@@ -122,7 +133,9 @@ impl StorageClient {
         self.client.unary_call(
             &METHOD_SYNC_ITERATE,
             &request,
-            CallOption::default().wait_for_ready(true),
+            CallOption::default()
+                .wait_for_ready(true)
+                .timeout(CALL_TIMEOUT),
         )
     }
 }


Test patch:
diff --git a/go/storage/mkvs/checkpoint/checkpoint_test.go b/go/storage/mkvs/checkpoint/checkpoint_test.go
index 2e3383b0d..7e67b56b7 100644
--- a/go/storage/mkvs/checkpoint/checkpoint_test.go
+++ b/go/storage/mkvs/checkpoint/checkpoint_test.go
@@ -50,7 +50,7 @@ func TestFileCheckpointCreator(t *testing.T) {
 	require.NoError(err, "Commit")
 	root := node.Root{
 		Namespace: testNs,
-		Version:   0,
+		Version:   1,
 		Hash:      rootHash,
 	}
 
diff --git a/go/storage/mkvs/db/badger/badger_test.go b/go/storage/mkvs/db/badger/badger_test.go
new file mode 100644
index 000000000..70077a083
--- /dev/null
+++ b/go/storage/mkvs/db/badger/badger_test.go
@@ -0,0 +1,327 @@
+package badger
+
+import (
+	"bytes"
+	"context"
+	"fmt"
+	"io/ioutil"
+	"os"
+	"strconv"
+	"testing"
+
+	"github.com/dgraph-io/badger/v2"
+	"github.com/stretchr/testify/require"
+
+	"github.com/oasisprotocol/oasis-core/go/common"
+	"github.com/oasisprotocol/oasis-core/go/common/crypto/hash"
+	"github.com/oasisprotocol/oasis-core/go/storage/mkvs"
+	"github.com/oasisprotocol/oasis-core/go/storage/mkvs/checkpoint"
+	"github.com/oasisprotocol/oasis-core/go/storage/mkvs/db/api"
+	"github.com/oasisprotocol/oasis-core/go/storage/mkvs/node"
+	"github.com/oasisprotocol/oasis-core/go/storage/mkvs/writelog"
+)
+
+var (
+	nodePrefix = nodeKeyFmt.Encode()
+
+	logPrefix = multipartRestoreNodeLogKeyFmt.Encode()
+
+	testNs = common.NewTestNamespaceFromSeed([]byte("badger node db test ns"), 0)
+
+	dbCfg = &api.Config{
+		Namespace:    testNs,
+		MaxCacheSize: 16 * 1024 * 1024,
+		NoFsync:      true,
+		MemoryOnly:   true,
+	}
+
+	testValues = [][]byte{
+		[]byte("colorless green ideas sleep furiously"),
+		[]byte("excepting understandable chairs piously"),
+		[]byte("at the prickle for rainbow hoovering"),
+	}
+)
+
+type keySet map[string]struct{}
+
+type test struct {
+	require  *require.Assertions
+	ctx      context.Context
+	dir      string
+	badgerdb *badgerNodeDB
+	ckMeta   *checkpoint.Metadata
+	ckNodes  keySet
+}
+
+func fillDB(ctx context.Context, require *require.Assertions, values [][]byte, version uint64, ndb api.NodeDB) node.Root {
+	emptyRoot := node.Root{
+		Namespace: testNs,
+		Version:   version,
+	}
+	emptyRoot.Hash.Empty()
+
+	tree := mkvs.NewWithRoot(nil, ndb, emptyRoot)
+	require.NotNil(tree, "NewWithRoot()")
+
+	var wl writelog.WriteLog
+	for i, val := range values {
+		wl = append(wl, writelog.LogEntry{Key: []byte(strconv.Itoa(i)), Value: val})
+	}
+
+	err := tree.ApplyWriteLog(ctx, writelog.NewStaticIterator(wl))
+	require.NoError(err, "ApplyWriteLog()")
+
+	_, hash, err := tree.Commit(ctx, testNs, 2)
+	require.NoError(err, "Commit()")
+
+	return node.Root{
+		Namespace: testNs,
+		Version:   version + 1,
+		Hash:      hash,
+	}
+}
+
+func createCheckpoint(ctx context.Context, require *require.Assertions, dir string, values [][]byte, version uint64) (*checkpoint.Metadata, keySet) {
+	ndb, err := New(dbCfg)
+	require.NoError(err, "New()")
+	defer ndb.Close()
+	badgerdb := ndb.(*badgerNodeDB)
+	fc, err := checkpoint.NewFileCreator(dir, ndb)
+	require.NoError(err, "NewFileCreator()")
+
+	ckRoot := fillDB(ctx, require, values, version, ndb)
+	ckMeta, err := fc.CreateCheckpoint(ctx, ckRoot, 1024*1024)
+	require.NoError(err, "CreateCheckpoint()")
+
+	nodeKeys := keySet{}
+	err = badgerdb.db.View(func(txn *badger.Txn) error {
+		it := txn.NewIterator(badger.DefaultIteratorOptions)
+		defer it.Close()
+		for it.Rewind(); it.Valid(); it.Next() {
+			if bytes.HasPrefix(it.Item().Key(), nodePrefix) {
+				nodeKeys[string(it.Item().Key())] = struct{}{}
+			}
+		}
+		return nil
+	})
+	require.NoError(err, "createCheckpoint()")
+
+	return ckMeta, nodeKeys
+}
+
+func verifyNodes(require *require.Assertions, badgerdb *badgerNodeDB, keySet keySet) {
+	notVisited := map[string]struct{}{}
+	for k := range keySet {
+		notVisited[k] = struct{}{}
+	}
+	err := badgerdb.db.View(func(txn *badger.Txn) error {
+		it := txn.NewIterator(badger.DefaultIteratorOptions)
+		defer it.Close()
+		for it.Rewind(); it.Valid(); it.Next() {
+			key := it.Item().Key()
+			if !bytes.HasPrefix(key, nodePrefix) {
+				continue
+			}
+			_, ok := keySet[string(key)]
+			require.Equal(true, ok, "unexpected node in db")
+			delete(notVisited, string(key))
+		}
+		return nil
+	})
+	require.NoError(err, "verifyNodes()")
+	require.Equal(0, len(notVisited), "some nodes not visited")
+}
+
+func checkNoLogKeys(require *require.Assertions, badgerdb *badgerNodeDB) {
+	err := badgerdb.db.View(func(txn *badger.Txn) error {
+		it := txn.NewIterator(badger.DefaultIteratorOptions)
+		defer it.Close()
+		for it.Rewind(); it.Valid(); it.Next() {
+			key := it.Item().Key()
+			require.False(bytes.HasPrefix(key, logPrefix), "checkLogKeys()/iteration")
+		}
+		return nil
+	})
+	require.NoError(err, "checkNoLogKeys()")
+}
+
+func restoreCheckpoint(ctx *test, ckMeta *checkpoint.Metadata, ckNodes keySet) checkpoint.Restorer {
+	fc, err := checkpoint.NewFileCreator(ctx.dir, ctx.badgerdb)
+	ctx.require.NoError(err, "NewFileCreator() - 2")
+
+	restorer, err := checkpoint.NewRestorer(ctx.badgerdb)
+	ctx.require.NoError(err, "NewRestorer()")
+
+	err = restorer.StartRestore(ctx.ctx, ckMeta)
+	ctx.require.NoError(err, "StartRestore()")
+	for i := range ckMeta.Chunks {
+		idx := uint64(i)
+		chunkMeta, err := ckMeta.GetChunkMetadata(idx)
+		ctx.require.NoError(err, fmt.Sprintf("GetChunkMetadata(%d)", idx))
+		func() {
+			r, w, err := os.Pipe()
+			ctx.require.NoError(err, "Pipe()")
+			errCh := make(chan error)
+			go func() {
+				_, errr := restorer.RestoreChunk(ctx.ctx, idx, r)
+				errCh <- errr
+			}()
+			err = fc.GetCheckpointChunk(ctx.ctx, chunkMeta, w)
+			w.Close()
+			errRestore := <-errCh
+			ctx.require.NoError(err, "GetCheckpointChunk()")
+			ctx.require.NoError(errRestore, "RestoreChunk()")
+		}()
+	}
+
+	verifyNodes(ctx.require, ctx.badgerdb, ckNodes)
+
+	return restorer
+}
+
+func TestMultipartRestore(t *testing.T) {
+	ctx := context.Background()
+	wrap := func(testFunc func(ctx *test), initialValues [][]byte) func(*testing.T) {
+		return func(t *testing.T) {
+			require := require.New(t)
+
+			dir, err := ioutil.TempDir("", "oasis-storage-database-test")
+			require.NoError(err, "TempDir()")
+			defer os.RemoveAll(dir)
+
+			ckMeta, ckNodes := createCheckpoint(ctx, require, dir, initialValues, 1)
+
+			ndb, err := New(dbCfg)
+			require.NoError(err, "New() - 2")
+			defer ndb.Close()
+			badgerdb := ndb.(*badgerNodeDB)
+
+			testCtx := &test{
+				require:  require,
+				ctx:      ctx,
+				dir:      dir,
+				badgerdb: badgerdb,
+				ckMeta:   ckMeta,
+				ckNodes:  ckNodes,
+			}
+			testFunc(testCtx)
+		}
+	}
+
+	t.Run("Abort", wrap(testAbort, testValues))
+	t.Run("Finalize", wrap(testFinalize, testValues))
+	t.Run("ExistingNodes", wrap(testExistingNodes, testValues[:1]))
+}
+
+func testAbort(ctx *test) {
+	// Abort a restore, check nodes again.
+	// There should be no leftover nodes, and the log keys should be gone too.
+	restorer := restoreCheckpoint(ctx, ctx.ckMeta, ctx.ckNodes)
+	err := restorer.AbortRestore(ctx.ctx)
+	ctx.require.NoError(err, "AbortRestore()")
+
+	verifyNodes(ctx.require, ctx.badgerdb, keySet{})
+	checkNoLogKeys(ctx.require, ctx.badgerdb)
+}
+
+func testFinalize(ctx *test) {
+	// Finalize a restore, check nodes again.
+	// This time, all the restored nodes should be present, but the
+	// log keys should be gone.
+	restoreCheckpoint(ctx, ctx.ckMeta, ctx.ckNodes)
+	err := ctx.badgerdb.Finalize(ctx.ctx, ctx.ckMeta.Root.Version, []hash.Hash{ctx.ckMeta.Root.Hash})
+	ctx.require.NoError(err, "Finalize()")
+
+	verifyNodes(ctx.require, ctx.badgerdb, ctx.ckNodes)
+	checkNoLogKeys(ctx.require, ctx.badgerdb)
+}
+
+func testExistingNodes(ctx *test) {
+	// Create two checkpoints, so we have two sets of nodes.
+	// The first checkpoint will be the base for a fresh database and must include
+	// a node from the second checkpoint, which will be used for multipart restore.
+	// The pre-existing node should then not be deleted after aborting the second
+	// checkpoint.
+
+	// Create the checkpoint to be used as the overriding restore.
+	ckMeta2, ckNodes2 := createCheckpoint(ctx.ctx, ctx.require, ctx.dir, testValues, 2)
+	var overlap bool
+	for node1 := range ctx.ckNodes {
+		if _, ok := ckNodes2[node1]; ok {
+			overlap = true
+			break
+		}
+	}
+	ctx.require.Equal(true, overlap, "pointless test when no nodes would overlap")
+
+	// Restore first checkpoint. The database is empty.
+	restoreCheckpoint(ctx, ctx.ckMeta, ctx.ckNodes)
+	err := ctx.badgerdb.Finalize(ctx.ctx, ctx.ckMeta.Root.Version, []hash.Hash{ctx.ckMeta.Root.Hash})
+	ctx.require.NoError(err, "Finalize()")
+	verifyNodes(ctx.require, ctx.badgerdb, ctx.ckNodes)
+
+	// Restore the second checkpoint. One of the nodes from it already exists. After aborting,
+	// exactly the nodes from the first checkpoint should remain.
+	restorer := restoreCheckpoint(ctx, ckMeta2, ckNodes2)
+	err = restorer.AbortRestore(ctx.ctx)
+	ctx.require.NoError(err, "AbortRestore()")
+	verifyNodes(ctx.require, ctx.badgerdb, ctx.ckNodes)
+}
+
+func TestVersionChecks(t *testing.T) {
+	require := require.New(t)
+	ndb, err := New(dbCfg)
+	require.NoError(err, "New()")
+	defer ndb.Close()
+	badgerdb := ndb.(*badgerNodeDB)
+
+	err = badgerdb.StartMultipartInsert(0)
+	require.Error(err, "StartMultipartInsert(0)")
+
+	err = badgerdb.StartMultipartInsert(42)
+	require.NoError(err, "StartMultipartInsert(42)")
+	err = badgerdb.StartMultipartInsert(44)
+	require.Error(err, "StartMultipartInsert(44)")
+
+	root := node.Root{}
+	_, err = badgerdb.NewBatch(root, 0, false) // Normal chunks not allowed during multipart.
+	require.Error(err, "NewBatch(.., 0, false)")
+	_, err = badgerdb.NewBatch(root, 13, true)
+	require.Error(err, "NewBatch(.., 13, true)")
+	batch, err := badgerdb.NewBatch(root, 42, true)
+	require.NoError(err, "NewBatch(.., 42, true)")
+	defer batch.Reset()
+
+	err = batch.Commit(root)
+	require.Error(err, "Commit(Root{0})")
+}
+
+func TestReadOnlyBatch(t *testing.T) {
+	require := require.New(t)
+
+	// No way to initialize a readonly-database, so it needs to be created rw first.
+	// This means we need persistence.
+	dir, err := ioutil.TempDir("", "oasis-storage-database-test")
+	require.NoError(err, "TempDir()")
+	defer os.RemoveAll(dir)
+
+	readonlyCfg := *dbCfg
+	readonlyCfg.MemoryOnly = false
+	readonlyCfg.ReadOnly = false
+	readonlyCfg.DB = dir
+
+	func() {
+		ndb, errRw := New(&readonlyCfg)
+		require.NoError(errRw, "New() - 1")
+		defer ndb.Close()
+	}()
+
+	readonlyCfg.ReadOnly = true
+	ndb, err := New(&readonlyCfg)
+	require.NoError(err, "New() - 2")
+	defer ndb.Close()
+	badgerdb := ndb.(*badgerNodeDB)
+
+	_, err = badgerdb.NewBatch(node.Root{}, 13, false)
+	require.Error(err, "NewBatch()")
+}
diff --git a/go/storage/mkvs/tree_test.go b/go/storage/mkvs/tree_test.go
index 8e28f0203..78a30422d 100644
--- a/go/storage/mkvs/tree_test.go
+++ b/go/storage/mkvs/tree_test.go
@@ -944,7 +944,8 @@ func testOnCommitHooks(t *testing.T, ndb db.NodeDB, factory NodeDBFactory) {
 		Hash:      emptyRoot,
 	}
 
-	batch := ndb.NewBatch(root, root.Version, false)
+	batch, err := ndb.NewBatch(root, root.Version, false)
+	require.NoError(t, err, "NewBatch")
 	defer batch.Reset()
 
 	var calls []int
@@ -961,7 +962,7 @@ func testOnCommitHooks(t *testing.T, ndb db.NodeDB, factory NodeDBFactory) {
 
 	require.True(t, len(calls) == 0, "OnCommit hooks should not fire before commit")
 
-	err := batch.Commit(root)
+	err = batch.Commit(root)
 	require.NoError(t, err, "Commit")
 	require.EqualValues(t, calls, []int{1, 2, 3}, "OnCommit hooks should fire in order")
 }