From fd620dd78c8e9dc1f90983611a8ccaf41d5e2151 Mon Sep 17 00:00:00 2001
From: Alexey Feldgendler <alexey@feldgendler.ru>
Date: Mon, 6 Apr 2020 15:37:26 +0200
Subject: [PATCH] Txn snapshot (#66)

* Using a fork of go-immutable-radix temporarily until PR #26 is accepted

* Adding Snapshot method to Txn to create a read-only snapshot of the transaction

* Removed temporary replacement; updated go-immutable-radix to v1.2.0
---
 go.mod      |  2 +-
 go.sum      |  4 ++--
 txn.go      | 23 ++++++++++++++++++++
 txn_test.go | 62 +++++++++++++++++++++++++++++++++++++++++++++++++++++
 4 files changed, 88 insertions(+), 3 deletions(-)

diff --git a/go.mod b/go.mod
index 34bfd82..f7fa64f 100644
--- a/go.mod
+++ b/go.mod
@@ -2,4 +2,4 @@ module github.com/hashicorp/go-memdb
 
 go 1.12
 
-require github.com/hashicorp/go-immutable-radix v1.1.0
+require github.com/hashicorp/go-immutable-radix v1.2.0
diff --git a/go.sum b/go.sum
index 1a21d60..21e7403 100644
--- a/go.sum
+++ b/go.sum
@@ -1,5 +1,5 @@
-github.com/hashicorp/go-immutable-radix v1.1.0 h1:vN9wG1D6KG6YHRTWr8512cxGOVgTMEfgEdSj/hr8MPc=
-github.com/hashicorp/go-immutable-radix v1.1.0/go.mod h1:0y9vanUI8NX6FsYoO3zeMjhV/C5i9g4Q3DwcSNZ4P60=
+github.com/hashicorp/go-immutable-radix v1.2.0 h1:l6UW37iCXwZkZoAbEYnptSHVE/cQ5bOTPYG5W3vf9+8=
+github.com/hashicorp/go-immutable-radix v1.2.0/go.mod h1:0y9vanUI8NX6FsYoO3zeMjhV/C5i9g4Q3DwcSNZ4P60=
 github.com/hashicorp/go-uuid v1.0.0 h1:RS8zrF7PhGwyNPOtxSClXXj9HA8feRnJzgnI1RJCSnM=
 github.com/hashicorp/go-uuid v1.0.0/go.mod h1:6SBZvOh/SIDV7/2o3Jml5SYk/TvGqwFJ/bN7x4byOro=
 github.com/hashicorp/golang-lru v0.5.0 h1:CL2msUPvZTLb5O648aiLNJw3hnBxN2+1Jq8rCOH9wdo=
diff --git a/txn.go b/txn.go
index 0c00d79..6466845 100644
--- a/txn.go
+++ b/txn.go
@@ -804,3 +804,26 @@ func (r *radixIterator) Next() interface{} {
 	}
 	return value
 }
+
+// Snapshot creates a snapshot of the current state of the transaction.
+// Returns a new read-only transaction or nil if the transaction is already
+// aborted or committed.
+func (txn *Txn) Snapshot() *Txn {
+	if txn.rootTxn == nil {
+		return nil
+	}
+
+	snapshot := &Txn{
+		db:      txn.db,
+		rootTxn: txn.rootTxn.Clone(),
+	}
+
+	// Commit sub-transactions into the snapshot
+	for key, subTxn := range txn.modified {
+		path := indexPath(key.Table, key.Index)
+		final := subTxn.CommitOnly()
+		snapshot.rootTxn.Insert(path, final)
+	}
+
+	return snapshot
+}
diff --git a/txn_test.go b/txn_test.go
index 8a4a065..0a7f483 100644
--- a/txn_test.go
+++ b/txn_test.go
@@ -1245,6 +1245,68 @@ func TestTxn_LowerBound(t *testing.T) {
 	}
 }
 
+func TestTxn_Snapshot(t *testing.T) {
+	db := testDB(t)
+	txn := db.Txn(true)
+
+	err := txn.Insert("main", &TestObject{
+		ID:  "one",
+		Foo: "abc",
+		Qux: []string{"abc1", "abc2"},
+	})
+	if err != nil {
+		t.Fatalf("err: %v", err)
+	}
+
+	snapshot := txn.Snapshot()
+
+	err = txn.Insert("main", &TestObject{
+		ID:  "two",
+		Foo: "def",
+		Qux: []string{"def1", "def2"},
+	})
+	if err != nil {
+		t.Fatalf("err: %v", err)
+	}
+
+	txn.Commit()
+
+	raw, err := snapshot.First("main", "id", "one")
+	if err != nil {
+		t.Fatalf("err: %v", err)
+	}
+	if raw == nil || raw.(*TestObject).ID != "one" {
+		t.Fatalf("TestObject one not found")
+	}
+
+	raw, err = snapshot.First("main", "id", "two")
+	if err != nil {
+		t.Fatalf("err: %v", err)
+	}
+	if raw != nil {
+		t.Fatalf("TestObject two found")
+	}
+
+	txn = db.Txn(false)
+	snapshot = txn.Snapshot()
+
+	raw, err = snapshot.First("main", "id", "one")
+	if err != nil {
+		t.Fatalf("err: %v", err)
+	}
+	if raw == nil || raw.(*TestObject).ID != "one" {
+		t.Fatalf("TestObject one not found")
+	}
+
+	raw, err = snapshot.First("main", "id", "two")
+	if err != nil {
+		t.Fatalf("err: %v", err)
+	}
+	if raw == nil || raw.(*TestObject).ID != "two" {
+		t.Fatalf("TestObject two not found")
+	}
+}
+
 func TestStringFieldIndexerEmptyPointerFromArgs(t *testing.T) {
 	t.Run("does not error with AllowMissing", func(t *testing.T) {
 		schema := &DBSchema{