Add hard-link support to CopyEngine [3/4] (#14947)

Summary:
Pull Request resolved: https://github.com/facebook/rocksdb/pull/14947

Teaches the shared `CopyEngine` to hard-link files through its background pool: adds `WorkItemType::Link` and `CopyEngine::LinkFile(src, dst, dst_env)`, which returns the FileSystem's `NotSupported` when links are unavailable so callers can fall back to copying. This is what CheckpointEngine (next diff) needs for warm storage, where linking is supported but benefits from parallelism.

No existing caller links through the engine yet (BackupEngine never hard-links), so no behavior change; the path is exercised by a new `copy_engine_test.cc` (parallel copy, parallel link, single-thread copy, per-item error reporting).

Reviewed By: xingbowang

Differential Revision: D111387972

fbshipit-source-id: ec0b0e9b80749a3f51e62066ceef2df21a1811d2

BUCK

@@ -4744,6 +4744,12 @@ cpp_unittest_wrapper(name="configurable_test",
             extra_compiler_flags=[])
 
 
+cpp_unittest_wrapper(name="copy_engine_test",
+            srcs=["utilities/copy_engine/copy_engine_test.cc"],
+            deps=[":rocksdb_test_lib"],
+            extra_compiler_flags=[])
+
+
 cpp_unittest_wrapper(name="corruption_test",
             srcs=["db/corruption_test.cc"],
             deps=[":rocksdb_test_lib"],

CMakeLists.txt

@@ -1641,6 +1641,7 @@ if(WITH_TESTS)
         utilities/cassandra/cassandra_row_merge_test.cc
         utilities/cassandra/cassandra_serialize_test.cc
         utilities/checkpoint/checkpoint_test.cc
+        utilities/copy_engine/copy_engine_test.cc
         utilities/sorted_run_builder/sorted_run_builder_test.cc
         utilities/env_timed_test.cc
         utilities/fault_injection_fs_test.cc

Makefile

@@ -1829,6 +1829,9 @@ backup_engine_test: $(OBJ_DIR)/utilities/backup/backup_engine_test.o $(TEST_LIBR
 checkpoint_test: $(OBJ_DIR)/utilities/checkpoint/checkpoint_test.o $(TEST_LIBRARY) $(LIBRARY)
 	$(AM_LINK)
 
+copy_engine_test: $(OBJ_DIR)/utilities/copy_engine/copy_engine_test.o $(TEST_LIBRARY) $(LIBRARY)
+	$(AM_LINK)
+
 sorted_run_builder_test: $(OBJ_DIR)/utilities/sorted_run_builder/sorted_run_builder_test.o $(TEST_LIBRARY) $(LIBRARY)
 	$(AM_LINK)
 

src.mk

@@ -657,6 +657,7 @@ TEST_MAIN_SOURCES =                                                     \
   utilities/cassandra/cassandra_row_merge_test.cc                       \
   utilities/cassandra/cassandra_serialize_test.cc                       \
   utilities/checkpoint/checkpoint_test.cc                               \
+  utilities/copy_engine/copy_engine_test.cc                     \
   utilities/sorted_run_builder/sorted_run_builder_test.cc               \
   utilities/env_timed_test.cc                                           \
   utilities/fault_injection_fs_test.cc                                  \

utilities/copy_engine/copy_engine.cc

@@ -126,6 +126,9 @@ void CopyEngine::ThreadBody() {
           work_item.src_temperature);
       result.db_id = work_item.db_id;
       result.db_session_id = work_item.db_session_id;
+    } else if (work_item.type == Link) {
+      result.io_status =
+          LinkFile(work_item.src_path, work_item.dst_path, work_item.dst_env);
     } else {
       result.io_status = IOStatus::InvalidArgument(
           "Unknown work item type: " + std::to_string(work_item.type));
@@ -284,6 +287,15 @@ IOStatus CopyEngine::CopyOrCreateFile(
   return io_s;
 }
 
+IOStatus CopyEngine::LinkFile(const std::string& src, const std::string& dst,
+                              Env* dst_env) {
+  if (ShouldAbort()) {
+    return status_to_io_status(
+        Status::Incomplete(options_.abort_status_message));
+  }
+  return dst_env->GetFileSystem()->LinkFile(src, dst, IOOptions(), nullptr);
+}
+
 uint64_t CopyEngine::CalculateIOBufferSize(RateLimiter* rate_limiter) const {
   if (options_.io_buffer_size > 0) {
     return options_.io_buffer_size;

utilities/copy_engine/copy_engine.h

@@ -69,6 +69,7 @@ struct WorkItemResult {
 enum WorkItemType : uint64_t {
   CopyOrCreate = 1U,
   ComputeChecksum = 2U,
+  Link = 3U,
 };
 
 // Exactly one of src_path and contents must be non-empty. If src_path is
@@ -177,8 +178,9 @@ struct CopyEngineOptions {
   std::string abort_status_message = "Operation stopped";
 };
 
-// A pool of background threads that copy or checksum files pulled from a shared
-// work queue, so BackupEngine and CheckpointEngine share one implementation.
+// A pool of background threads that copy, hard-link, or checksum files pulled
+// from a shared work queue, so BackupEngine and CheckpointEngine share one
+// implementation.
 class CopyEngine {
  public:
   explicit CopyEngine(CopyEngineOptions options);
@@ -213,6 +215,11 @@ class CopyEngine {
                                       std::string* checksum_hex,
                                       Temperature src_temperature) const;
 
+  // Returns the FileSystem's NotSupported when linking is unavailable, so
+  // callers can fall back to copying.
+  IOStatus LinkFile(const std::string& src, const std::string& dst,
+                    Env* dst_env);
+
   uint64_t CalculateIOBufferSize(RateLimiter* rate_limiter) const;
 
  private:

utilities/copy_engine/copy_engine_test.cc

@@ -0,0 +1,192 @@
+//  Copyright (c) Meta Platforms, Inc. and affiliates.
+//  This source code is licensed under both the GPLv2 (found in the
+//  COPYING file in the root directory) and Apache 2.0 License
+//  (found in the LICENSE.Apache file in the root directory).
+
+#include "utilities/copy_engine/copy_engine.h"
+
+#include <future>
+#include <string>
+#include <vector>
+
+#include "file/file_util.h"
+#include "port/stack_trace.h"
+#include "rocksdb/env.h"
+#include "test_util/testharness.h"
+
+namespace ROCKSDB_NAMESPACE {
+
+class CopyEngineTest : public testing::Test {
+ protected:
+  CopyEngineTest() {
+    env_ = Env::Default();
+    dir_ = test::PerThreadDBPath(env_, "copy_engine_test");
+  }
+
+  // PerThreadDBPath is deterministic per (name, pid, thread), so every case in
+  // this fixture maps to the same directory. Destroy it around each case so
+  // files from one case (or a prior/repeated run) cannot leak into the next.
+  void SetUp() override {
+    ASSERT_OK(DestroyDir(env_, dir_));
+    ASSERT_OK(env_->CreateDirIfMissing(dir_));
+  }
+
+  void TearDown() override { ASSERT_OK(DestroyDir(env_, dir_)); }
+
+  std::string SrcPath(int i) const {
+    return dir_ + "/src_" + std::to_string(i);
+  }
+  std::string DstPath(int i) const {
+    return dir_ + "/dst_" + std::to_string(i);
+  }
+
+  void CreateSrc(int i, const std::string& content) {
+    ASSERT_OK(WriteStringToFile(env_, content, SrcPath(i),
+                                /*should_sync=*/false));
+  }
+
+  // Varied sizes exercise the dynamic pull-based pool (not a static split).
+  std::string ContentFor(int i) const {
+    return "file-" + std::to_string(i) + "-" + std::string((i + 1) * 1024, 'x');
+  }
+
+  WorkItem MakeCopyItem(int i) {
+    return WorkItem(SrcPath(i), DstPath(i), Temperature::kUnknown,
+                    Temperature::kUnknown, /*contents=*/"", env_, env_,
+                    EnvOptions(), /*sync=*/false, /*rate_limiter=*/nullptr,
+                    /*size_limit=*/0, /*stats=*/nullptr);
+  }
+
+  WorkItem MakeLinkItem(int i) {
+    WorkItem w = MakeCopyItem(i);
+    w.type = WorkItemType::Link;
+    return w;
+  }
+
+  void VerifyDstMatches(int i, const std::string& expected) {
+    std::string got;
+    ASSERT_OK(ReadFileToString(env_, DstPath(i), &got));
+    ASSERT_EQ(expected, got);
+  }
+
+  Env* env_;
+  std::string dir_;
+};
+
+TEST_F(CopyEngineTest, ParallelCopy) {
+  constexpr int kNumFiles = 16;
+  std::vector<std::string> contents;
+  for (int i = 0; i < kNumFiles; ++i) {
+    contents.push_back(ContentFor(i));
+    CreateSrc(i, contents[i]);
+  }
+
+  CopyEngineOptions options;
+  options.max_background_operations = 4;
+  CopyEngine engine(std::move(options));
+
+  std::vector<std::future<WorkItemResult>> results;
+  for (int i = 0; i < kNumFiles; ++i) {
+    WorkItem w = MakeCopyItem(i);
+    results.push_back(w.result.get_future());
+    engine.Submit(std::move(w));
+  }
+
+  for (int i = 0; i < kNumFiles; ++i) {
+    WorkItemResult r = results[i].get();
+    ASSERT_OK(r.io_status);
+    ASSERT_EQ(contents[i].size(), r.size);
+    VerifyDstMatches(i, contents[i]);
+  }
+}
+
+TEST_F(CopyEngineTest, ParallelLink) {
+  constexpr int kNumFiles = 16;
+  std::vector<std::string> contents;
+  for (int i = 0; i < kNumFiles; ++i) {
+    contents.push_back(ContentFor(i));
+    CreateSrc(i, contents[i]);
+  }
+
+  CopyEngineOptions options;
+  options.max_background_operations = 4;
+  CopyEngine engine(std::move(options));
+
+  std::vector<std::future<WorkItemResult>> results;
+  for (int i = 0; i < kNumFiles; ++i) {
+    WorkItem w = MakeLinkItem(i);
+    results.push_back(w.result.get_future());
+    engine.Submit(std::move(w));
+  }
+
+  bool link_supported = true;
+  for (int i = 0; i < kNumFiles; ++i) {
+    WorkItemResult r = results[i].get();
+    if (r.io_status.IsNotSupported()) {
+      link_supported = false;
+      continue;
+    }
+    ASSERT_OK(r.io_status);
+    VerifyDstMatches(i, contents[i]);
+  }
+  if (!link_supported) {
+    ROCKSDB_GTEST_SKIP("FileSystem does not support hard links");
+  }
+}
+
+TEST_F(CopyEngineTest, SerialCopyMatchesSingleThread) {
+  constexpr int kNumFiles = 5;
+  std::vector<std::string> contents;
+  for (int i = 0; i < kNumFiles; ++i) {
+    contents.push_back(ContentFor(i));
+    CreateSrc(i, contents[i]);
+  }
+
+  CopyEngineOptions options;
+  options.max_background_operations = 1;
+  CopyEngine engine(std::move(options));
+
+  std::vector<std::future<WorkItemResult>> results;
+  for (int i = 0; i < kNumFiles; ++i) {
+    WorkItem w = MakeCopyItem(i);
+    results.push_back(w.result.get_future());
+    engine.Submit(std::move(w));
+  }
+
+  for (int i = 0; i < kNumFiles; ++i) {
+    WorkItemResult r = results[i].get();
+    ASSERT_OK(r.io_status);
+    VerifyDstMatches(i, contents[i]);
+  }
+}
+
+TEST_F(CopyEngineTest, MissingSourceReportsErrorPerItem) {
+  CreateSrc(0, ContentFor(0));
+
+  CopyEngineOptions options;
+  options.max_background_operations = 2;
+  CopyEngine engine(std::move(options));
+
+  WorkItem ok_item = MakeCopyItem(0);
+  auto ok_future = ok_item.result.get_future();
+  engine.Submit(std::move(ok_item));
+
+  WorkItem bad_item = MakeCopyItem(1);  // no source file on disk
+  auto bad_future = bad_item.result.get_future();
+  engine.Submit(std::move(bad_item));
+
+  WorkItemResult ok_result = ok_future.get();
+  ASSERT_OK(ok_result.io_status);
+  VerifyDstMatches(0, ContentFor(0));
+
+  WorkItemResult bad_result = bad_future.get();
+  ASSERT_NOK(bad_result.io_status);
+}
+
+}  // namespace ROCKSDB_NAMESPACE
+
+int main(int argc, char** argv) {
+  ROCKSDB_NAMESPACE::port::InstallStackTraceHandler();
+  ::testing::InitGoogleTest(&argc, argv);
+  return RUN_ALL_TESTS();
+}