Integrate the snapshotbatch with Juno

Author: MaksymMalicki

blockchain/blockchain.go

@@ -99,7 +99,8 @@ func (b *Blockchain) Network() *utils.Network {
 // If blockchain is empty zero felt is returned.
 func (b *Blockchain) StateCommitment() (felt.Felt, error) {
 	b.listener.OnRead("StateCommitment")
-	batch := b.database.NewIndexedBatch() // this is a hack because we don't need to write to the db
+	batch, closer := b.database.NewSnapshotBatch() // this is a hack because we don't need to write to the db
+	defer closer()
 	return core.NewState(batch).Commitment()
 }
 
@@ -116,8 +117,9 @@ func (b *Blockchain) Head() (*core.Block, error) {
 		return nil, err
 	}
 
-	txn := b.database.NewIndexedBatch()
-	return core.GetBlockByNumber(txn, curHeight)
+	snapshot := b.database.NewSnapshot()
+	defer snapshot.Close()
+	return core.GetBlockByNumber(snapshot, curHeight)
 }
 
 func (b *Blockchain) HeadsHeader() (*core.Header, error) {
@@ -141,8 +143,9 @@ func headsHeader(txn db.KeyValueReader) (*core.Header, error) {
 
 func (b *Blockchain) BlockByNumber(number uint64) (*core.Block, error) {
 	b.listener.OnRead("BlockByNumber")
-	txn := b.database.NewIndexedBatch()
-	return core.GetBlockByNumber(txn, number)
+	snapshot := b.database.NewSnapshot()
+	defer snapshot.Close()
+	return core.GetBlockByNumber(snapshot, number)
 }
 
 func (b *Blockchain) BlockHeaderByNumber(number uint64) (*core.Header, error) {
@@ -157,8 +160,9 @@ func (b *Blockchain) BlockByHash(hash *felt.Felt) (*core.Block, error) {
 		return nil, err
 	}
 
-	txn := b.database.NewIndexedBatch()
-	return core.GetBlockByNumber(txn, blockNum)
+	snapshot := b.database.NewSnapshot()
+	defer snapshot.Close()
+	return core.GetBlockByNumber(snapshot, blockNum)
 }
 
 func (b *Blockchain) BlockHeaderByHash(hash *felt.Felt) (*core.Header, error) {
@@ -244,48 +248,51 @@ func (b *Blockchain) Store(
 	stateUpdate *core.StateUpdate,
 	newClasses map[felt.Felt]core.ClassDefinition,
 ) error {
-	err := b.database.Update(func(txn db.IndexedBatch) error {
-		if err := verifyBlock(txn, block); err != nil {
-			return err
-		}
-
-		if err := core.NewState(txn).Update(block.Number, stateUpdate, newClasses, false); err != nil {
-			return err
-		}
-		if err := core.WriteBlockHeader(txn, block.Header); err != nil {
-			return err
-		}
+	txn, closer := b.database.NewSnapshotBatch()
+	defer closer()
+	if err := verifyBlock(txn, block); err != nil {
+		return err
+	}
 
-		for i, tx := range block.Transactions {
-			if err := core.WriteTxAndReceipt(txn, block.Number, uint64(i), tx,
-				block.Receipts[i]); err != nil {
-				return err
-			}
-		}
+	if err := core.NewState(txn).Update(block.Number, stateUpdate, newClasses, false); err != nil {
+		return err
+	}
+	if err := core.WriteBlockHeader(txn, block.Header); err != nil {
+		return err
+	}
 
-		if err := core.WriteStateUpdateByBlockNum(txn, block.Number, stateUpdate); err != nil {
+	for i, tx := range block.Transactions {
+		if err := core.WriteTxAndReceipt(txn, block.Number, uint64(i), tx,
+			block.Receipts[i]); err != nil {
 			return err
 		}
+	}
 
-		if err := core.WriteBlockCommitment(txn, block.Number, blockCommitments); err != nil {
-			return err
-		}
+	if err := core.WriteStateUpdateByBlockNum(txn, block.Number, stateUpdate); err != nil {
+		return err
+	}
 
-		if err := core.WriteL1HandlerMsgHashes(txn, block.Transactions); err != nil {
-			return err
-		}
+	if err := core.WriteBlockCommitment(txn, block.Number, blockCommitments); err != nil {
+		return err
+	}
 
-		err := storeCasmClassHashesV2ForBlock(txn, block.ProtocolVersion, newClasses, stateUpdate)
-		if err != nil {
-			return err
-		}
+	if err := core.WriteL1HandlerMsgHashes(txn, block.Transactions); err != nil {
+		return err
+	}
 
-		return core.WriteChainHeight(txn, block.Number)
-	})
+	err := storeCasmClassHashesV2ForBlock(txn, block.ProtocolVersion, newClasses, stateUpdate)
 	if err != nil {
 		return err
 	}
 
+	if err := core.WriteChainHeight(txn, block.Number); err != nil {
+		return err
+	}
+
+	if err := txn.Write(); err != nil {
+		return err
+	}
+
 	return b.runningFilter.Insert(
 		block.EventsBloom,
 		block.Number,
@@ -296,7 +303,7 @@ func (b *Blockchain) Store(
 // For versions < 0.14.1, it computes hashes from class definitions.
 // For versions >= 0.14.1, it uses pre-computed hashes from the state update.
 func storeCasmClassHashesV2ForBlock(
-	txn db.IndexedBatch,
+	txn db.KeyValueWriter,
 	protocolVersion string,
 	newClasses map[felt.Felt]core.ClassDefinition,
 	stateUpdate *core.StateUpdate,
@@ -315,7 +322,7 @@ func storeCasmClassHashesV2ForBlock(
 
 // computeAndStoreCasmClassHashesV2 computes and stores CASM class hashes V2 from class definitions.
 func computeAndStoreCasmClassHashesV2(
-	txn db.IndexedBatch,
+	txn db.KeyValueWriter,
 	declaredClasses map[felt.Felt]core.ClassDefinition,
 ) error {
 	for classHash, classDefinition := range declaredClasses {
@@ -336,7 +343,7 @@ func computeAndStoreCasmClassHashesV2(
 
 // storeCasmClassHashesV2 stores pre-computed CASM class hashes V2 from the state update.
 func storeCasmClassHashesV2(
-	txn db.IndexedBatch,
+	txn db.KeyValueWriter,
 	declaredV1Classes map[felt.Felt]*felt.Felt,
 ) error {
 	for classHash, casmClassHashV2 := range declaredV1Classes {
@@ -405,33 +412,32 @@ func (b *Blockchain) SanityCheckNewHeight(block *core.Block, stateUpdate *core.S
 
 type StateCloser = func() error
 
-var noopStateCloser = func() error { return nil } // TODO: remove this once we refactor the state
-
 // HeadState returns a StateReader that provides a stable view to the latest state
 func (b *Blockchain) HeadState() (core.StateReader, StateCloser, error) {
 	b.listener.OnRead("HeadState")
-	txn := b.database.NewIndexedBatch()
+	snapshot := b.database.NewSnapshot()
 
-	_, err := core.GetChainHeight(txn)
+	_, err := core.GetChainHeight(snapshot)
 	if err != nil {
 		return nil, nil, err
 	}
 
-	return core.NewState(txn), noopStateCloser, nil
+	return core.NewStateSnapshotReader(snapshot), snapshot.Close, nil
 }
 
 // StateAtBlockNumber returns a StateReader that provides
 // a stable view to the state at the given block number
 func (b *Blockchain) StateAtBlockNumber(blockNumber uint64) (core.StateReader, StateCloser, error) {
 	b.listener.OnRead("StateAtBlockNumber")
-	txn := b.database.NewIndexedBatch()
+	snapshot := b.database.NewSnapshot()
 
-	_, err := core.GetBlockHeaderByNumber(txn, blockNumber)
+	_, err := core.GetBlockHeaderByNumber(snapshot, blockNumber)
 	if err != nil {
 		return nil, nil, err
 	}
 
-	return core.NewDeprecatedStateHistory(core.NewState(txn), blockNumber), noopStateCloser, nil
+	stateReader := core.NewStateSnapshotReader(snapshot)
+	return core.NewDeprecatedStateHistory(stateReader, blockNumber), snapshot.Close, nil
 }
 
 // StateAtBlockHash returns a StateReader that provides
@@ -440,18 +446,19 @@ func (b *Blockchain) StateAtBlockHash(blockHash *felt.Felt) (core.StateReader, S
 	b.listener.OnRead("StateAtBlockHash")
 	if blockHash.IsZero() {
 		memDB := memory.New()
-		txn := memDB.NewIndexedBatch()
-		emptyState := core.NewState(txn)
-		return emptyState, noopStateCloser, nil
+		snapshot := memDB.NewSnapshot()
+		emptyState := core.NewStateSnapshotReader(snapshot)
+		return emptyState, snapshot.Close, nil
 	}
 
-	txn := b.database.NewIndexedBatch()
-	header, err := core.GetBlockHeaderByHash(txn, blockHash)
+	snapshot := b.database.NewSnapshot()
+	header, err := core.GetBlockHeaderByHash(snapshot, blockHash)
 	if err != nil {
 		return nil, nil, err
 	}
 
-	return core.NewDeprecatedStateHistory(core.NewState(txn), header.Number), noopStateCloser, nil
+	stateReader := core.NewStateSnapshotReader(snapshot)
+	return core.NewDeprecatedStateHistory(stateReader, header.Number), snapshot.Close, nil
 }
 
 // EventFilter returns an EventFilter object that is tied to a snapshot of the blockchain
@@ -480,25 +487,34 @@ func (b *Blockchain) EventFilter(
 
 // RevertHead reverts the head block
 func (b *Blockchain) RevertHead() error {
-	return b.database.Update(b.revertHead)
+	txn, closer := b.database.NewSnapshotBatch()
+	defer closer()
+	if err := b.revertHead(txn); err != nil {
+		return err
+	}
+	if err := txn.Write(); err != nil {
+		return err
+	}
+	return nil
 }
 
 // todo(rdr): return `core.StateDiff` by value
 func (b *Blockchain) GetReverseStateDiff() (*core.StateDiff, error) {
 	var reverseStateDiff *core.StateDiff
 
-	txn := b.database.NewIndexedBatch()
-	blockNum, err := core.GetChainHeight(txn)
+	snapshot := b.database.NewSnapshot()
+	defer snapshot.Close()
+	blockNum, err := core.GetChainHeight(snapshot)
 	if err != nil {
 		return nil, err
 	}
 
-	stateUpdate, err := core.GetStateUpdateByBlockNum(txn, blockNum)
+	stateUpdate, err := core.GetStateUpdateByBlockNum(snapshot, blockNum)
 	if err != nil {
 		return nil, err
 	}
 
-	state := core.NewState(txn)
+	state := core.NewStateSnapshotReader(snapshot)
 	reverseStateDiff, err = state.GetReverseStateDiff(blockNum, stateUpdate.StateDiff)
 	if err != nil {
 		return nil, err
@@ -507,7 +523,7 @@ func (b *Blockchain) GetReverseStateDiff() (*core.StateDiff, error) {
 	return reverseStateDiff, nil
 }
 
-func (b *Blockchain) revertHead(txn db.IndexedBatch) error {
+func (b *Blockchain) revertHead(txn db.SnapshotBatch) error {
 	blockNumber, err := core.GetChainHeight(txn)
 	if err != nil {
 		return err
@@ -579,8 +595,8 @@ func (b *Blockchain) Simulate(
 	sign utils.BlockSignFunc,
 ) (SimulateResult, error) {
 	// Simulate without commit
-	txn := b.database.NewIndexedBatch()
-	defer txn.Reset()
+	txn, closer := b.database.NewSnapshotBatch()
+	defer closer()
 
 	if err := b.updateStateRoots(txn, block, stateUpdate, newClasses); err != nil {
 		return SimulateResult{}, err
@@ -615,38 +631,41 @@ func (b *Blockchain) Finalise(
 	newClasses map[felt.Felt]core.ClassDefinition,
 	sign utils.BlockSignFunc,
 ) error {
-	err := b.database.Update(func(txn db.IndexedBatch) error {
-		if err := b.updateStateRoots(txn, block, stateUpdate, newClasses); err != nil {
-			return err
-		}
-		commitments, err := b.updateBlockHash(block, stateUpdate)
-		if err != nil {
-			return err
-		}
-		if err := b.signBlock(block, stateUpdate, sign); err != nil {
-			return err
-		}
-		if err := b.storeBlockData(txn, block, stateUpdate, commitments); err != nil {
-			return err
-		}
-
-		err = storeCasmClassHashesV2ForBlock(txn, block.ProtocolVersion, newClasses, stateUpdate)
-		if err != nil {
-			return err
-		}
+	txn, closer := b.database.NewSnapshotBatch()
+	defer closer()
+	if err := b.updateStateRoots(txn, block, stateUpdate, newClasses); err != nil {
+		return err
+	}
+	commitments, err := b.updateBlockHash(block, stateUpdate)
+	if err != nil {
+		return err
+	}
+	if err := b.signBlock(block, stateUpdate, sign); err != nil {
+		return err
+	}
+	if err := b.storeBlockData(txn, block, stateUpdate, commitments); err != nil {
+		return err
+	}
 
-		return core.WriteChainHeight(txn, block.Number)
-	})
+	err = storeCasmClassHashesV2ForBlock(txn, block.ProtocolVersion, newClasses, stateUpdate)
 	if err != nil {
 		return err
 	}
 
+	if err := core.WriteChainHeight(txn, block.Number); err != nil {
+		return err
+	}
+
+	if err := txn.Write(); err != nil {
+		return err
+	}
+
 	return b.runningFilter.Insert(block.EventsBloom, block.Number)
 }
 
 // updateStateRoots computes and updates state roots in the block and state update
 func (b *Blockchain) updateStateRoots(
-	txn db.IndexedBatch,
+	txn db.SnapshotBatch,
 	block *core.Block,
 	stateUpdate *core.StateUpdate,
 	newClasses map[felt.Felt]core.ClassDefinition,
@@ -715,7 +734,7 @@ func (b *Blockchain) signBlock(
 
 // storeBlockData persists all block-related data to the database
 func (b *Blockchain) storeBlockData(
-	txn db.IndexedBatch,
+	txn db.SnapshotBatch,
 	block *core.Block,
 	stateUpdate *core.StateUpdate,
 	commitments *core.BlockCommitments,