Merge pull request #5008 from onflow/leo/storehouse-extending-block-s…

…napshot

[Storehouse] Add Extending block snapshot

engine/execution/storehouse.go

@@ -1,6 +1,7 @@
 package execution
 
 import (
+	"github.com/onflow/flow-go/fvm/storage/snapshot"
 	"github.com/onflow/flow-go/model/flow"
 	"github.com/onflow/flow-go/module/finalizedreader"
 	"github.com/onflow/flow-go/storage"
@@ -105,3 +106,9 @@ type WALReader interface {
 	// It returns EOF when there are no more entries.
 	Next() (height uint64, registers flow.RegisterEntries, err error)
 }
+
+type ExtendableStorageSnapshot interface {
+	snapshot.StorageSnapshot
+	Extend(newCommit flow.StateCommitment, updatedRegisters map[flow.RegisterID]flow.RegisterValue) ExtendableStorageSnapshot
+	Commitment() flow.StateCommitment
+}

engine/execution/storehouse/executing_block_snapshot.go

@@ -0,0 +1,70 @@
+package storehouse
+
+import (
+	"github.com/onflow/flow-go/engine/execution"
+	"github.com/onflow/flow-go/fvm/storage/snapshot"
+	"github.com/onflow/flow-go/model/flow"
+)
+
+var _ execution.ExtendableStorageSnapshot = (*ExecutingBlockSnapshot)(nil)
+
+// ExecutingBlockSnapshot is a snapshot of the storage at an executed collection.
+// It starts with a storage snapshot at the end of previous block,
+// The register updates at the executed collection at baseHeight + 1 are cached in
+// a map, such that retrieving register values at the snapshot will first check
+// the cache, and then the storage.
+type ExecutingBlockSnapshot struct {
+	// the snapshot at the end of previous block
+	previous snapshot.StorageSnapshot
+
+	commitment      flow.StateCommitment
+	registerUpdates map[flow.RegisterID]flow.RegisterValue
+}
+
+// create a new storage snapshot for an executed collection
+// at the base block at height h - 1
+func NewExecutingBlockSnapshot(
+	previous snapshot.StorageSnapshot,
+	// the statecommitment of a block at height h
+	commitment flow.StateCommitment,
+) *ExecutingBlockSnapshot {
+	return &ExecutingBlockSnapshot{
+		previous:        previous,
+		commitment:      commitment,
+		registerUpdates: make(map[flow.RegisterID]flow.RegisterValue),
+	}
+}
+
+// Get returns the register value at the snapshot.
+func (s *ExecutingBlockSnapshot) Get(id flow.RegisterID) (flow.RegisterValue, error) {
+	// get from latest updates first
+	value, ok := s.getFromUpdates(id)
+	if ok {
+		return value, nil
+	}
+
+	// get from BlockEndStateSnapshot at previous block
+	value, err := s.previous.Get(id)
+	return value, err
+}
+
+func (s *ExecutingBlockSnapshot) getFromUpdates(id flow.RegisterID) (flow.RegisterValue, bool) {
+	value, ok := s.registerUpdates[id]
+	return value, ok
+}
+
+// Extend returns a new storage snapshot at the same block but but for a different state commitment,
+// which contains the given registerUpdates
+// Usually it's used to create a new storage snapshot at the next executed collection.
+// The registerUpdates contains the register updates at the executed collection.
+func (s *ExecutingBlockSnapshot) Extend(newCommit flow.StateCommitment, updates map[flow.RegisterID]flow.RegisterValue) execution.ExtendableStorageSnapshot {
+	return &ExecutingBlockSnapshot{
+		previous:        s,
+		commitment:      newCommit,
+		registerUpdates: updates,
+	}
+}
+
+func (s *ExecutingBlockSnapshot) Commitment() flow.StateCommitment {
+	return s.commitment
+}

engine/execution/storehouse/executing_block_snapshot_test.go

@@ -0,0 +1,92 @@
+package storehouse_test
+
+import (
+	"testing"
+
+	"github.com/stretchr/testify/require"
+
+	"github.com/onflow/flow-go/engine/execution/storehouse"
+	"github.com/onflow/flow-go/fvm/storage/snapshot"
+	"github.com/onflow/flow-go/model/flow"
+	"github.com/onflow/flow-go/utils/unittest"
+)
+
+func TestExtendingBlockSnapshot(t *testing.T) {
+	t.Run("Get register", func(t *testing.T) {
+		reg1 := makeReg("key1", "val1")
+		base := snapshot.MapStorageSnapshot{
+			reg1.Key: reg1.Value,
+		}
+		baseCommit := unittest.StateCommitmentFixture()
+		snap := storehouse.NewExecutingBlockSnapshot(base, baseCommit)
+
+		// should get value
+		value, err := snap.Get(reg1.Key)
+		require.NoError(t, err)
+		require.Equal(t, reg1.Value, value)
+
+		// should get nil for unknown register
+		unknown := makeReg("unknown", "unknownV")
+		value, err = snap.Get(unknown.Key)
+		require.NoError(t, err)
+		require.Equal(t, []byte(nil), value)
+	})
+
+	t.Run("Extend snapshot", func(t *testing.T) {
+		reg1 := makeReg("key1", "val1")
+		reg2 := makeReg("key2", "val2")
+		base := snapshot.MapStorageSnapshot{
+			reg1.Key: reg1.Value,
+			reg2.Key: reg2.Value,
+		}
+		// snap1: { key1: val1, key2: val2 }
+		snap1 := storehouse.NewExecutingBlockSnapshot(base, unittest.StateCommitmentFixture())
+
+		updatedReg2 := makeReg("key2", "val22")
+		reg3 := makeReg("key3", "val3")
+		// snap2: { key1: val1, key2: val22, key3: val3 }
+		snap2 := snap1.Extend(unittest.StateCommitmentFixture(), map[flow.RegisterID]flow.RegisterValue{
+			updatedReg2.Key: updatedReg2.Value,
+			reg3.Key:        reg3.Value,
+		})
+
+		// should get un-changed value
+		value, err := snap2.Get(reg1.Key)
+		require.NoError(t, err)
+		require.Equal(t, []byte("val1"), value)
+
+		value, err = snap2.Get(reg2.Key)
+		require.NoError(t, err)
+		require.Equal(t, []byte("val22"), value)
+
+		value, err = snap2.Get(reg3.Key)
+		require.NoError(t, err)
+		require.Equal(t, []byte("val3"), value)
+
+		// should get nil for unknown register
+		unknown := makeReg("unknown", "unknownV")
+		value, err = snap2.Get(unknown.Key)
+		require.NoError(t, err)
+		require.Equal(t, []byte(nil), value)
+
+		// create snap3 with reg3 updated
+		// snap3: { key1: val1, key2: val22, key3: val33 }
+		updatedReg3 := makeReg("key3", "val33")
+		snap3 := snap2.Extend(unittest.StateCommitmentFixture(), map[flow.RegisterID]flow.RegisterValue{
+			updatedReg3.Key: updatedReg3.Value,
+		})
+
+		// verify all keys
+		value, err = snap3.Get(reg1.Key)
+		require.NoError(t, err)
+		require.Equal(t, []byte("val1"), value)
+
+		value, err = snap3.Get(reg2.Key)
+		require.NoError(t, err)
+		require.Equal(t, []byte("val22"), value)
+
+		value, err = snap3.Get(reg3.Key)
+		require.NoError(t, err)
+		require.Equal(t, []byte("val33"), value)
+	})
+}

utils/unittest/fixtures.go

@@ -2764,3 +2764,13 @@ func RegisterEntryFixture() flow.RegisterEntry {
 		Value: val,
 	}
 }
+
+func MakeOwnerReg(key string, value string) flow.RegisterEntry {
+	return flow.RegisterEntry{
+		Key: flow.RegisterID{
+			Owner: "owner",
+			Key:   key,
+		},
+		Value: []byte(value),
+	}
+}