From 08ccae32f34f46a9d1471b1e39bc8626d566acd8 Mon Sep 17 00:00:00 2001
From: alesapin <alesapin@gmail.com>
Date: Tue, 2 Mar 2021 15:34:18 +0300
Subject: [PATCH] Fix nasty bug

---
 .../NuKeeperStorageSerializer.cpp             |  1 -
 src/Coordination/NuKeeperStorageSerializer.h  |  1 -
 src/Coordination/SnapshotableHashTable.h      |  6 ++-
 src/Coordination/tests/gtest_for_build.cpp    | 44 ++++++++++++++++++-
 4 files changed, 47 insertions(+), 5 deletions(-)

diff --git a/src/Coordination/NuKeeperStorageSerializer.cpp b/src/Coordination/NuKeeperStorageSerializer.cpp
index bf57590f448..dd4fc62d744 100644
--- a/src/Coordination/NuKeeperStorageSerializer.cpp
+++ b/src/Coordination/NuKeeperStorageSerializer.cpp
@@ -148,7 +148,6 @@ NuKeeperStorageSnapshot::NuKeeperStorageSnapshot(NuKeeperStorage * storage_, siz
     storage->enableSnapshotMode();
     snapshot_container_size = storage->container.snapshotSize();
     begin = storage->getSnapshotIteratorBegin();
-    end = storage->getSnapshotIteratorEnd();
     session_and_timeout = storage->getActiveSessions();
 }
 
diff --git a/src/Coordination/NuKeeperStorageSerializer.h b/src/Coordination/NuKeeperStorageSerializer.h
index 5e6407835ec..59acba97679 100644
--- a/src/Coordination/NuKeeperStorageSerializer.h
+++ b/src/Coordination/NuKeeperStorageSerializer.h
@@ -30,7 +30,6 @@ public:
     int64_t session_id;
     size_t snapshot_container_size;
     NuKeeperStorage::Container::const_iterator begin;
-    NuKeeperStorage::Container::const_iterator end;
     SessionAndTimeout session_and_timeout;
 };
 
diff --git a/src/Coordination/SnapshotableHashTable.h b/src/Coordination/SnapshotableHashTable.h
index 90b8bf05ced..3fca985758f 100644
--- a/src/Coordination/SnapshotableHashTable.h
+++ b/src/Coordination/SnapshotableHashTable.h
@@ -68,7 +68,8 @@ public:
                 ListElem elem{key, value, true};
                 list_itr->active_in_map = false;
                 auto new_list_itr = list.insert(list.end(), elem);
-                map[new_list_itr->key] = new_list_itr;
+                map.erase(it);
+                map.emplace(new_list_itr->key, new_list_itr);
             }
             else
             {
@@ -115,7 +116,8 @@ public:
 
             updater(elem_copy.value);
             auto itr = list.insert(list.end(), elem_copy);
-            map[itr->key] = itr;
+            map.erase(it);
+            map.emplace(itr->key, itr);
             return itr;
         }
         else
diff --git a/src/Coordination/tests/gtest_for_build.cpp b/src/Coordination/tests/gtest_for_build.cpp
index 97539017914..3bd39ba8bb0 100644
--- a/src/Coordination/tests/gtest_for_build.cpp
+++ b/src/Coordination/tests/gtest_for_build.cpp
@@ -898,7 +898,7 @@ TEST(CoordinationTest, SnapshotableHashMapTrySnapshot)
 void addNode(DB::NuKeeperStorage & storage, const std::string & path, const std::string & data, int64_t ephemeral_owner=0)
 {
     using Node = DB::NuKeeperStorage::Node;
-    storage.container.insert(path, Node{.data=data, .ephemeral_owner = ephemeral_owner});
+    storage.container.insertOrReplace(path, Node{.data=data, .ephemeral_owner = ephemeral_owner});
 }
 
 TEST(CoordinationTest, TestStorageSnapshotSimple)
@@ -1029,6 +1029,48 @@ TEST(CoordinationTest, TestStorageSnapshotManySnapshots)
     }
 }
 
+TEST(CoordinationTest, TestStorageSnapshotMode)
+{
+    ChangelogDirTest test("./snapshots");
+    DB::NuKeeperSnapshotManager manager("./snapshots", 3);
+    DB::NuKeeperStorage storage(500);
+    for (size_t i = 0; i < 50; ++i)
+    {
+        addNode(storage, "/hello_" + std::to_string(i), "world_" + std::to_string(i));
+    }
+
+    {
+        DB::NuKeeperStorageSnapshot snapshot(&storage, 50);
+        for (size_t i = 0; i < 50; ++i)
+        {
+            addNode(storage, "/hello_" + std::to_string(i), "wlrd_" + std::to_string(i));
+        }
+        for (size_t i = 0; i < 50; ++i)
+        {
+            EXPECT_EQ(storage.container.getValue("/hello_" + std::to_string(i)).data, "wlrd_" + std::to_string(i));
+        }
+        EXPECT_EQ(storage.container.size(), 51);
+        EXPECT_EQ(storage.container.snapshotSize(), 101);
+        auto buf = manager.serializeSnapshotToBuffer(snapshot);
+        manager.serializeSnapshotBufferToDisk(*buf, 50);
+    }
+    EXPECT_TRUE(fs::exists("./snapshots/snapshot_50.bin"));
+    EXPECT_EQ(storage.container.size(), 51);
+    EXPECT_EQ(storage.container.snapshotSize(), 51);
+    for (size_t i = 0; i < 50; ++i)
+    {
+        EXPECT_EQ(storage.container.getValue("/hello_" + std::to_string(i)).data, "wlrd_" + std::to_string(i));
+    }
+
+    DB::NuKeeperStorage restored_storage(500);
+    manager.restoreFromLatestSnapshot(&restored_storage);
+
+    for (size_t i = 0; i < 50; ++i)
+    {
+        EXPECT_EQ(restored_storage.container.getValue("/hello_" + std::to_string(i)).data, "world_" + std::to_string(i));
+    }
+}
+
 int main(int argc, char ** argv)
 {
     Poco::AutoPtr<Poco::ConsoleChannel> channel(new Poco::ConsoleChannel(std::cerr));