Added memory diff to memory mutation (ledgerwatch#1068)

kv/memdb/memory_mutation.go

@@ -17,6 +17,7 @@ import (
 	"bytes"
 	"context"
 
+	"github.com/ledgerwatch/erigon-lib/common"
 	"github.com/ledgerwatch/erigon-lib/kv/iter"
 	"github.com/ledgerwatch/erigon-lib/kv/order"
 	"github.com/ledgerwatch/log/v3"
@@ -378,6 +379,72 @@ func (m *MemoryMutation) Flush(tx kv.RwTx) error {
 	return nil
 }
 
+func (m *MemoryMutation) Diff() (*MemoryDiff, error) {
+	memDiff := &MemoryDiff{
+		diff:           make(map[table][]entry),
+		deletedEntries: make(map[string][]string),
+	}
+	// Obtain buckets touched.
+	buckets, err := m.memTx.ListBuckets()
+	if err != nil {
+		return nil, err
+	}
+	// Obliterate buckets who are to be deleted
+	for bucket := range m.clearedTables {
+		memDiff.clearedTableNames = append(memDiff.clearedTableNames, bucket)
+	}
+	// Obliterate entries who are to be deleted
+	for bucket, keys := range m.deletedEntries {
+		for key := range keys {
+			memDiff.deletedEntries[bucket] = append(memDiff.deletedEntries[bucket], key)
+		}
+	}
+	// Iterate over each bucket and apply changes accordingly.
+	for _, bucket := range buckets {
+		if isTablePurelyDupsort(bucket) {
+			cbucket, err := m.memTx.CursorDupSort(bucket)
+			if err != nil {
+				return nil, err
+			}
+			defer cbucket.Close()
+
+			t := table{
+				name:    bucket,
+				dupsort: true,
+			}
+			for k, v, err := cbucket.First(); k != nil; k, v, err = cbucket.Next() {
+				if err != nil {
+					return nil, err
+				}
+				memDiff.diff[t] = append(memDiff.diff[t], entry{
+					k: common.Copy(k),
+					v: common.Copy(v),
+				})
+			}
+		} else {
+			cbucket, err := m.memTx.Cursor(bucket)
+			if err != nil {
+				return nil, err
+			}
+			defer cbucket.Close()
+			t := table{
+				name:    bucket,
+				dupsort: false,
+			}
+			for k, v, err := cbucket.First(); k != nil; k, v, err = cbucket.Next() {
+				if err != nil {
+					return nil, err
+				}
+				memDiff.diff[t] = append(memDiff.diff[t], entry{
+					k: common.Copy(k),
+					v: common.Copy(v),
+				})
+			}
+		}
+	}
+	return memDiff, nil
+}
+
 // Check if a bucket is dupsorted and has dupsort conversion off
 func isTablePurelyDupsort(bucket string) bool {
 	config, ok := kv.ChaindataTablesCfg[bucket]

kv/memdb/memory_mutation_diff.go

@@ -0,0 +1,58 @@
+package memdb
+
+import "github.com/ledgerwatch/erigon-lib/kv"
+
+type entry struct {
+	k []byte
+	v []byte
+}
+
+type MemoryDiff struct {
+	diff              map[table][]entry // god.
+	deletedEntries    map[string][]string
+	clearedTableNames []string
+}
+
+type table struct {
+	name    string
+	dupsort bool
+}
+
+func (m *MemoryDiff) Flush(tx kv.RwTx) error {
+	// Obliterate buckets who are to be deleted
+	for _, bucket := range m.clearedTableNames {
+		if err := tx.ClearBucket(bucket); err != nil {
+			return err
+		}
+	}
+	// Obliterate entries who are to be deleted
+	for bucket, keys := range m.deletedEntries {
+		for _, key := range keys {
+			if err := tx.Delete(bucket, []byte(key)); err != nil {
+				return err
+			}
+		}
+	}
+	// Iterate over each bucket and apply changes accordingly.
+	for bucketInfo, bucketDiff := range m.diff {
+		if bucketInfo.dupsort {
+			dbCursor, err := tx.RwCursorDupSort(bucketInfo.name)
+			if err != nil {
+				return err
+			}
+			defer dbCursor.Close()
+			for _, entry := range bucketDiff {
+				if err := dbCursor.Put(entry.k, entry.v); err != nil {
+					return err
+				}
+			}
+		} else {
+			for _, entry := range bucketDiff {
+				if err := tx.Put(bucketInfo.name, entry.k, entry.v); err != nil {
+					return err
+				}
+			}
+		}
+	}
+	return nil
+}