From e6de6fdae9cc13a803438c3cad653646b80579b6 Mon Sep 17 00:00:00 2001
From: Jernej Kos <jernej@kos.mx>
Date: Mon, 23 Mar 2020 12:59:46 +0100
Subject: [PATCH] storage/mkvs: Don't try to sync dirty keys

---
 .changelog/2775.bugfix.2.md                   |  1 +
 go/storage/mkvs/urkel/lookup.go               |  7 +++++++
 go/storage/mkvs/urkel/remove.go               |  9 ++++++++-
 go/storage/mkvs/urkel/testdata/case-5.json    | 18 ++++++++++++++++++
 go/storage/mkvs/urkel/urkel_test.go           |  5 +++++
 runtime/src/storage/mkvs/urkel/tree/lookup.rs |  5 +++++
 runtime/src/storage/mkvs/urkel/tree/remove.rs |  5 +++++
 7 files changed, 49 insertions(+), 1 deletion(-)
 create mode 100644 .changelog/2775.bugfix.2.md
 create mode 100644 go/storage/mkvs/urkel/testdata/case-5.json

diff --git a/.changelog/2775.bugfix.2.md b/.changelog/2775.bugfix.2.md
new file mode 100644
index 00000000000..fc63f0f8fc3
--- /dev/null
+++ b/.changelog/2775.bugfix.2.md
@@ -0,0 +1 @@
+storage/mkvs: Don't try to sync dirty keys
diff --git a/go/storage/mkvs/urkel/lookup.go b/go/storage/mkvs/urkel/lookup.go
index c7195938d8c..6154d7e2cfe 100644
--- a/go/storage/mkvs/urkel/lookup.go
+++ b/go/storage/mkvs/urkel/lookup.go
@@ -17,6 +17,13 @@ func (t *tree) Get(ctx context.Context, key []byte) ([]byte, error) {
 		return nil, ErrClosed
 	}
 
+	// If the key has been modified locally, no need to perform any lookups.
+	if !t.withoutWriteLog {
+		if entry := t.pendingWriteLog[node.ToMapKey(key)]; entry != nil {
+			return entry.value, nil
+		}
+	}
+
 	// Remember where the path from root to target node ends (will end).
 	t.cache.markPosition()
 
diff --git a/go/storage/mkvs/urkel/remove.go b/go/storage/mkvs/urkel/remove.go
index d27f6d2264e..15d09821d1f 100644
--- a/go/storage/mkvs/urkel/remove.go
+++ b/go/storage/mkvs/urkel/remove.go
@@ -16,6 +16,14 @@ func (t *tree) RemoveExisting(ctx context.Context, key []byte) ([]byte, error) {
 		return nil, ErrClosed
 	}
 
+	// If the key has already been removed locally, don't try to remove it again.
+	var entry *pendingEntry
+	if !t.withoutWriteLog {
+		if entry = t.pendingWriteLog[node.ToMapKey(key)]; entry != nil && entry.value == nil {
+			return nil, nil
+		}
+	}
+
 	// Remember where the path from root to target node ends (will end).
 	t.cache.markPosition()
 
@@ -26,7 +34,6 @@ func (t *tree) RemoveExisting(ctx context.Context, key []byte) ([]byte, error) {
 
 	// Update the pending write log.
 	if !t.withoutWriteLog {
-		entry := t.pendingWriteLog[node.ToMapKey(key)]
 		if entry == nil {
 			t.pendingWriteLog[node.ToMapKey(key)] = &pendingEntry{key, nil, changed, nil}
 		} else {
diff --git a/go/storage/mkvs/urkel/testdata/case-5.json b/go/storage/mkvs/urkel/testdata/case-5.json
new file mode 100644
index 00000000000..8182c474c80
--- /dev/null
+++ b/go/storage/mkvs/urkel/testdata/case-5.json
@@ -0,0 +1,18 @@
+[
+    {
+        "op": "Insert",
+        "key": "AwAA"
+    },
+    {
+        "op": "Insert",
+        "key": "AwAe"
+    },
+    {
+        "op": "Insert",
+        "key": "AAAe"
+    },
+    {
+        "op": "Remove",
+        "key": "AAAe"
+    }
+]
\ No newline at end of file
diff --git a/go/storage/mkvs/urkel/urkel_test.go b/go/storage/mkvs/urkel/urkel_test.go
index 3b191c52937..6bcfae0faad 100644
--- a/go/storage/mkvs/urkel/urkel_test.go
+++ b/go/storage/mkvs/urkel/urkel_test.go
@@ -1872,6 +1872,10 @@ func testSpecialCase4(t *testing.T, ndb db.NodeDB, factory NodeDBFactory) {
 	testSpecialCaseFromJSON(t, ndb, "case-4.json")
 }
 
+func testSpecialCase5(t *testing.T, ndb db.NodeDB, factory NodeDBFactory) {
+	testSpecialCaseFromJSON(t, ndb, "case-5.json")
+}
+
 func testBackend(
 	t *testing.T,
 	initBackend func(t *testing.T) (NodeDBFactory, func()),
@@ -1911,6 +1915,7 @@ func testBackend(
 		{"SpecialCase2", testSpecialCase2},
 		{"SpecialCase3", testSpecialCase3},
 		{"SpecialCase4", testSpecialCase4},
+		{"SpecialCase5", testSpecialCase5},
 		{"Errors", testErrors},
 	}
 
diff --git a/runtime/src/storage/mkvs/urkel/tree/lookup.rs b/runtime/src/storage/mkvs/urkel/tree/lookup.rs
index 68f0a3da486..4e2d53324f9 100644
--- a/runtime/src/storage/mkvs/urkel/tree/lookup.rs
+++ b/runtime/src/storage/mkvs/urkel/tree/lookup.rs
@@ -49,6 +49,11 @@ impl UrkelTree {
         let boxed_key = key.to_vec();
         let pending_root = self.cache.borrow().get_pending_root();
 
+        // If the key has been modified locally, no need to perform any lookups.
+        if let Some(PendingLogEntry { ref value, .. }) = self.pending_write_log.get(&boxed_key) {
+            return Ok(value.clone());
+        }
+
         // Remember where the path from root to target node ends (will end).
         self.cache.borrow_mut().mark_position();
 
diff --git a/runtime/src/storage/mkvs/urkel/tree/remove.rs b/runtime/src/storage/mkvs/urkel/tree/remove.rs
index 0e8b6bfa108..9d324d185e0 100644
--- a/runtime/src/storage/mkvs/urkel/tree/remove.rs
+++ b/runtime/src/storage/mkvs/urkel/tree/remove.rs
@@ -14,6 +14,11 @@ impl UrkelTree {
         let boxed_key = key.to_vec();
         let pending_root = self.cache.borrow().get_pending_root();
 
+        // If the key has already been removed locally, don't try to remove it again.
+        if let Some(PendingLogEntry { value: None, .. }) = self.pending_write_log.get(&boxed_key) {
+            return Ok(None);
+        }
+
         // Remember where the path from root to target node ends (will end).
         self.cache.borrow_mut().mark_position();