fix: snapshot recover for disk table (#2174)

zhanghaohit · web-flow · commit 0769892c8b26 · 2022-07-19T17:10:19.000+08:00
* fix: snapshot recover for disk table if snapshot dir not exists

* hardlink snapshot to data dir rather than rename
diff --git a/src/base/file_util.h b/src/base/file_util.h
@@ -103,6 +103,26 @@ inline static int GetSubDir(const std::string& path,
     return 0;
 }
 
+inline static int GetSubFiles(const std::string& path, std::vector<std::string>& sub_dir) {  // NOLINT
+    if (path.empty()) {
+        return -1;
+    }
+    DIR* dir = opendir(path.c_str());
+    if (dir == NULL) {
+        return -1;
+    }
+    struct dirent* ptr;
+    while ((ptr = readdir(dir)) != NULL) {
+        if (strcmp(ptr->d_name, ".") == 0 || strcmp(ptr->d_name, "..") == 0) {
+            continue;
+        } else if (ptr->d_type == DT_REG) {
+            sub_dir.push_back(ptr->d_name);
+        }
+    }
+    closedir(dir);
+    return 0;
+}
+
 inline static int GetFileName(const std::string& path,
                               std::vector<std::string>& file_vec) {  // NOLINT
     if (path.empty()) {
@@ -287,6 +307,27 @@ __attribute__((unused)) static bool CopyFile(const std::string& src_file, const
     return has_error == false;
 }
 
+inline static int HardLinkDir(const std::string& src, const std::string& dest) {
+    if (!IsExists(src)) {
+        return -2;
+    }
+
+    if (IsExists(dest)) {
+        RemoveDirRecursive(dest);
+    }
+
+    MkdirRecur(dest);
+    std::vector<std::string> files;
+    GetSubFiles(src, files);
+    for (const auto& file : files) {
+        int ret = link((src + "/" + file).c_str(), (dest + "/" + file).c_str());
+        if (ret) {
+            return ret;
+        }
+    }
+    return 0;
+}
+
 }  // namespace base
 }  // namespace openmldb
 
diff --git a/src/replica/snapshot_replica_test.cc b/src/replica/snapshot_replica_test.cc
@@ -46,6 +46,7 @@ using ::openmldb::storage::Ticket;
 using ::openmldb::tablet::TabletImpl;
 
 DECLARE_string(db_root_path);
+DECLARE_string(ssd_root_path);
 DECLARE_string(hdd_root_path);
 DECLARE_string(endpoint);
 DECLARE_int32(make_snapshot_threshold_offset);
@@ -326,6 +327,226 @@ TEST_P(SnapshotReplicaTest, SendSnapshot) {
     sleep(2);
 }
 
+TEST_P(SnapshotReplicaTest, IncompleteSnapshot) {
+    FLAGS_make_snapshot_threshold_offset = 0;
+    uint32_t tid = 2;
+    uint32_t pid = 123;
+    auto storage_mode = GetParam();
+    uint64_t cur_time = ::baidu::common::timer::get_micros() / 1000;
+    {
+        ::openmldb::tablet::TabletImpl* tablet = new ::openmldb::tablet::TabletImpl();
+        MockClosure closure;
+        tablet->Init("");
+        brpc::Server server;
+        if (server.AddService(tablet, brpc::SERVER_OWNS_SERVICE) != 0) {
+            PDLOG(WARNING, "fail to register tablet rpc service");
+            exit(1);
+        }
+        brpc::ServerOptions options;
+        std::string leader_point = "127.0.0.1:18529";
+        if (server.Start(leader_point.c_str(), &options) != 0) {
+            PDLOG(WARNING, "fail to start server %s", leader_point.c_str());
+            exit(1);
+        }
+
+        ::openmldb::client::TabletClient client(leader_point, "");
+        client.Init();
+        std::vector<std::string> endpoints;
+        bool ret =
+            client.CreateTable("table1", tid, pid, 100000, 0, true, endpoints, ::openmldb::type::TTLType::kAbsoluteTime,
+                               16, 0, ::openmldb::type::CompressType::kNoCompress, storage_mode);
+        ASSERT_TRUE(ret);
+        ret = client.Put(tid, pid, "testkey", cur_time, ::openmldb::test::EncodeKV("testkey", "value1"));
+        ASSERT_TRUE(ret);
+
+        uint32_t count = 0;
+        while (count < 10) {
+            count++;
+            std::string key = "test";
+            client.Put(tid, pid, key, cur_time + count, ::openmldb::test::EncodeKV(key, key));
+        }
+        ::openmldb::api::GeneralRequest grq;
+        grq.set_tid(tid);
+        grq.set_pid(pid);
+        grq.set_storage_mode(storage_mode);
+        ::openmldb::api::GeneralResponse grp;
+        grp.set_code(-1);
+        tablet->MakeSnapshot(NULL, &grq, &grp, &closure);
+        sleep(5);
+    }
+
+    {
+        ::openmldb::tablet::TabletImpl* tablet = new ::openmldb::tablet::TabletImpl();
+        MockClosure closure;
+        tablet->Init("");
+        brpc::Server server;
+        if (server.AddService(tablet, brpc::SERVER_OWNS_SERVICE) != 0) {
+            PDLOG(WARNING, "fail to register tablet rpc service");
+            exit(1);
+        }
+        brpc::ServerOptions options;
+        std::string leader_point = "127.0.0.1:18529";
+        if (server.Start(leader_point.c_str(), &options) != 0) {
+            PDLOG(WARNING, "fail to start server %s", leader_point.c_str());
+            exit(1);
+        }
+
+        ::openmldb::client::TabletClient client(leader_point, "");
+        client.Init();
+
+        // load table
+        ::openmldb::api::TableMeta table_meta;
+        table_meta.set_format_version(1);
+        table_meta.set_name("table1");
+        table_meta.set_tid(tid);
+        table_meta.set_pid(pid);
+        table_meta.set_storage_mode(storage_mode);
+        client.LoadTable(table_meta, nullptr);
+        sleep(5);
+
+        ::openmldb::api::ScanRequest sr;
+        sr.set_tid(tid);
+        sr.set_pid(pid);
+        sr.set_pk("testkey");
+        sr.set_st(cur_time + 1);
+        sr.set_et(cur_time - 1);
+        sr.set_limit(10);
+        ::openmldb::api::ScanResponse srp;
+        tablet->Scan(NULL, &sr, &srp, &closure);
+        ASSERT_EQ(1, (int64_t)srp.count());
+        ASSERT_EQ(0, srp.code());
+
+        sr.set_tid(tid);
+        sr.set_pid(pid);
+        sr.set_pk("test");
+        sr.set_st(cur_time + 20);
+        sr.set_et(cur_time - 20);
+        sr.set_limit(20);
+        tablet->Scan(NULL, &sr, &srp, &closure);
+        ASSERT_EQ(10, (int64_t)srp.count());
+        ASSERT_EQ(0, srp.code());
+
+        std::string key = "test2";
+        ASSERT_TRUE(client.Put(tid, pid, key, cur_time, ::openmldb::test::EncodeKV(key, key)));
+
+        sr.set_tid(tid);
+        sr.set_pid(pid);
+        sr.set_pk("test2");
+        sr.set_st(cur_time + 20);
+        sr.set_et(cur_time - 20);
+        sr.set_limit(20);
+        tablet->Scan(NULL, &sr, &srp, &closure);
+        ASSERT_EQ(1, (int64_t)srp.count());
+        ASSERT_EQ(0, srp.code());
+    }
+
+    // remove the snapshot file, only keeping the MANIFEST
+    // i.e., corrupt the snapshot
+    std::string db_root_path;
+    if (storage_mode == common::kSSD) {
+        db_root_path = FLAGS_ssd_root_path;
+    } else if (storage_mode == common::kHDD) {
+        db_root_path = FLAGS_hdd_root_path;
+    } else {
+        db_root_path = FLAGS_db_root_path;
+    }
+    std::string snapshot_path = absl::StrCat(db_root_path, "/", tid, "_", pid, "/snapshot/");
+    std::vector<std::string> sub_dirs;
+    ::openmldb::base::GetSubDir(snapshot_path, sub_dirs);
+    for (const auto& dir : sub_dirs) {
+        auto sub_path = absl::StrCat(snapshot_path, dir);
+        DLOG(INFO) << "remove snapshot path: " << sub_path;
+        ASSERT_TRUE(::openmldb::base::RemoveDir(sub_path));
+    }
+    sleep(2);
+
+    {
+        ::openmldb::tablet::TabletImpl* tablet = new ::openmldb::tablet::TabletImpl();
+        MockClosure closure;
+        tablet->Init("");
+        brpc::Server server;
+        if (server.AddService(tablet, brpc::SERVER_OWNS_SERVICE) != 0) {
+            PDLOG(WARNING, "fail to register tablet rpc service");
+            exit(1);
+        }
+        brpc::ServerOptions options;
+        std::string leader_point = "127.0.0.1:18529";
+        if (server.Start(leader_point.c_str(), &options) != 0) {
+            PDLOG(WARNING, "fail to start server %s", leader_point.c_str());
+            exit(1);
+        }
+
+        ::openmldb::client::TabletClient client(leader_point, "");
+        client.Init();
+
+        // load table
+        ::openmldb::api::TableMeta table_meta;
+        table_meta.set_format_version(1);
+        table_meta.set_name("table1");
+        table_meta.set_tid(tid);
+        table_meta.set_pid(pid);
+        table_meta.set_storage_mode(storage_mode);
+        client.LoadTable(table_meta, nullptr);
+        sleep(5);
+
+        ::openmldb::api::ScanRequest sr;
+        sr.set_tid(tid);
+        sr.set_pid(pid);
+        sr.set_pk("testkey");
+        sr.set_st(cur_time + 1);
+        sr.set_et(cur_time - 1);
+        sr.set_limit(10);
+        ::openmldb::api::ScanResponse srp;
+        tablet->Scan(NULL, &sr, &srp, &closure);
+        ASSERT_EQ(1, (int64_t)srp.count());
+        ASSERT_EQ(0, srp.code());
+
+        sr.set_tid(tid);
+        sr.set_pid(pid);
+        sr.set_pk("test");
+        sr.set_st(cur_time + 20);
+        sr.set_et(cur_time - 20);
+        sr.set_limit(20);
+        tablet->Scan(NULL, &sr, &srp, &closure);
+        ASSERT_EQ(10, (int64_t)srp.count());
+        ASSERT_EQ(0, srp.code());
+
+        sr.set_tid(tid);
+        sr.set_pid(pid);
+        sr.set_pk("test2");
+        sr.set_st(cur_time + 20);
+        sr.set_et(cur_time - 20);
+        sr.set_limit(20);
+        tablet->Scan(NULL, &sr, &srp, &closure);
+        ASSERT_EQ(1, (int64_t)srp.count());
+        ASSERT_EQ(0, srp.code());
+
+        uint32_t count = 0;
+        while (count < 10) {
+            count++;
+            std::string key = "test3";
+            client.Put(tid, pid, key, cur_time + count, ::openmldb::test::EncodeKV(key, key));
+        }
+
+        sr.set_tid(tid);
+        sr.set_pid(pid);
+        sr.set_pk("test3");
+        sr.set_st(cur_time + 20);
+        sr.set_et(cur_time - 20);
+        sr.set_limit(20);
+        tablet->Scan(NULL, &sr, &srp, &closure);
+        ASSERT_EQ(10, (int64_t)srp.count());
+        ASSERT_EQ(0, srp.code());
+
+        ::openmldb::api::DropTableRequest dr;
+        dr.set_tid(tid);
+        dr.set_pid(pid);
+        ::openmldb::api::DropTableResponse drs;
+        tablet->DropTable(NULL, &dr, &drs, &closure);
+        sleep(2);
+    }
+}
+
 TEST_P(SnapshotReplicaTest, LeaderAndFollowerTS) {
     auto storage_mode = GetParam();
     ::openmldb::tablet::TabletImpl* tablet = new ::openmldb::tablet::TabletImpl();
diff --git a/src/tablet/tablet_impl.cc b/src/tablet/tablet_impl.cc
@@ -3121,16 +3121,17 @@ int TabletImpl::LoadDiskTableInternal(uint32_t tid, uint32_t pid, const ::openml
         std::string manifest_file = snapshot_path + "MANIFEST";
         if (Snapshot::GetLocalManifest(manifest_file, manifest) == 0) {
             std::string snapshot_dir = snapshot_path + manifest.name();
-            PDLOG(INFO, "rename dir %s to %s. tid %u pid %u", snapshot_dir.c_str(), data_path.c_str(), tid, pid);
-            if (!::openmldb::base::Rename(snapshot_dir, data_path)) {
-                PDLOG(WARNING, "rename dir failed. tid %u pid %u path %s", tid, pid, snapshot_dir.c_str());
-                break;
-            }
-            if (unlink(manifest_file.c_str()) < 0) {
-                PDLOG(WARNING, "remove manifest failed. tid %u pid %u path %s", tid, pid, manifest_file.c_str());
-                break;
+            if (::openmldb::base::IsExists(snapshot_dir)) {
+                PDLOG(INFO, "hardlink dir %s to %s (tid %u pid %u)", snapshot_dir.c_str(), data_path.c_str(), tid, pid);
+                if (::openmldb::base::HardLinkDir(snapshot_dir, data_path)) {
+                    PDLOG(WARNING, "hardlink snapshot dir %s to data dir failed (tid %u pid %u)", snapshot_dir.c_str(),
+                          tid, pid);
+                    break;
+                }
+                snapshot_offset = manifest.offset();
+            } else {
+                PDLOG(WARNING, "snapshot_dir %s with tid %u pid %u not exists", snapshot_dir.c_str(), tid, pid);
             }
-            snapshot_offset = manifest.offset();
         }
         std::string msg;
         if (CreateTableInternal(&table_meta, msg) < 0) {
@@ -3178,7 +3179,6 @@ int TabletImpl::LoadDiskTableInternal(uint32_t tid, uint32_t pid, const ::openml
             task_pool_.DelayTask(FLAGS_binlog_delete_interval,
                                  boost::bind(&TabletImpl::SchedDelBinlog, this, tid, pid));
             PDLOG(INFO, "load table success. tid %u pid %u", tid, pid);
-            MakeSnapshotInternal(tid, pid, 0, std::shared_ptr<::openmldb::api::TaskInfo>());
             std::string old_data_path = table_path + "/old_data";
             if (::openmldb::base::IsExists(old_data_path)) {
                 if (!::openmldb::base::RemoveDir(old_data_path)) {
@@ -3191,6 +3191,8 @@ int TabletImpl::LoadDiskTableInternal(uint32_t tid, uint32_t pid, const ::openml
                 task_ptr->set_status(::openmldb::api::TaskStatus::kDone);
                 return 0;
             }
+            PDLOG(INFO, "Recover table with tid %u and pid %u from binlog offset %u to %u", tid, pid, snapshot_offset,
+                  latest_offset);
         } else {
             DeleteTableInternal(tid, pid, std::shared_ptr<::openmldb::api::TaskInfo>());
         }
