rocksdb-native 2.2.0 → 2.3.0

package/vendor/librocksdb/vendor/rocksdb/db/write_thread.cc

@@ -404,6 +404,8 @@ void WriteThread::JoinBatchGroup(Writer* w) {
 
   bool linked_as_leader = LinkOne(w, &newest_writer_);
 
+  w->CheckWriteEnqueuedCallback();
+
   if (linked_as_leader) {
     SetState(w, STATE_GROUP_LEADER);
   }
@@ -428,6 +430,7 @@ void WriteThread::JoinBatchGroup(Writer* w) {
     TEST_SYNC_POINT_CALLBACK("WriteThread::JoinBatchGroup:BeganWaiting", w);
     AwaitState(w,
                STATE_GROUP_LEADER | STATE_MEMTABLE_WRITER_LEADER |
+                   STATE_PARALLEL_MEMTABLE_CALLER |
                    STATE_PARALLEL_MEMTABLE_WRITER | STATE_COMPLETED,
                &jbg_ctx);
     TEST_SYNC_POINT_CALLBACK("WriteThread::JoinBatchGroup:DoneWaiting", w);
@@ -656,12 +659,57 @@ void WriteThread::ExitAsMemTableWriter(Writer* /*self*/,
   SetState(leader, STATE_COMPLETED);
 }
 
+void WriteThread::SetMemWritersEachStride(Writer* w) {
+  WriteGroup* write_group = w->write_group;
+  Writer* last_writer = write_group->last_writer;
+
+  // The stride is the same for each writer in write_group, so w will
+  // call the writers with the same number in write_group mod total size
+  size_t stride = static_cast<size_t>(std::sqrt(write_group->size));
+  size_t count = 0;
+  while (w) {
+    if (count++ % stride == 0) {
+      SetState(w, STATE_PARALLEL_MEMTABLE_WRITER);
+    }
+    w = (w == last_writer) ? nullptr : w->link_newer;
+  }
+}
+
 void WriteThread::LaunchParallelMemTableWriters(WriteGroup* write_group) {
   assert(write_group != nullptr);
-  write_group->running.store(write_group->size);
-  for (auto w : *write_group) {
-    SetState(w, STATE_PARALLEL_MEMTABLE_WRITER);
+  size_t group_size = write_group->size;
+  write_group->running.store(group_size);
+
+  // The minimum number to allow the group use parallel caller mode.
+  // The number must no lower than 3;
+  const size_t MinParallelSize = 20;
+
+  // The group_size is too small, and there is no need to have
+  // the parallel partial callers.
+  if (group_size < MinParallelSize) {
+    for (auto w : *write_group) {
+      SetState(w, STATE_PARALLEL_MEMTABLE_WRITER);
+    }
+    return;
   }
+
+  // The stride is equal to std::sqrt(group_size) which can minimize
+  // the total number of leader SetSate.
+  // Set the leader itself STATE_PARALLEL_MEMTABLE_WRITER, and set
+  // (stride-1) writers to be STATE_PARALLEL_MEMTABLE_CALLER.
+  size_t stride = static_cast<size_t>(std::sqrt(group_size));
+  auto w = write_group->leader;
+  SetState(w, STATE_PARALLEL_MEMTABLE_WRITER);
+
+  for (size_t i = 1; i < stride; i++) {
+    w = w->link_newer;
+    SetState(w, STATE_PARALLEL_MEMTABLE_CALLER);
+  }
+
+  // After setting all STATE_PARALLEL_MEMTABLE_CALLER, the leader also
+  // does the job as STATE_PARALLEL_MEMTABLE_CALLER.
+  w = w->link_newer;
+  SetMemWritersEachStride(w);
 }
 
 static WriteThread::AdaptationContext cpmtw_ctx(
@@ -788,8 +836,8 @@ void WriteThread::ExitAsBatchGroupLeader(WriteGroup& write_group,
     }
 
     AwaitState(leader,
-               STATE_MEMTABLE_WRITER_LEADER | STATE_PARALLEL_MEMTABLE_WRITER |
-                   STATE_COMPLETED,
+               STATE_MEMTABLE_WRITER_LEADER | STATE_PARALLEL_MEMTABLE_CALLER |
+                   STATE_PARALLEL_MEMTABLE_WRITER | STATE_COMPLETED,
                &eabgl_ctx);
   } else {
     Writer* head = newest_writer_.load(std::memory_order_acquire);

package/vendor/librocksdb/vendor/rocksdb/db/write_thread.h

@@ -22,7 +22,9 @@
 #include "rocksdb/options.h"
 #include "rocksdb/status.h"
 #include "rocksdb/types.h"
+#include "rocksdb/user_write_callback.h"
 #include "rocksdb/write_batch.h"
+#include "util/aligned_storage.h"
 #include "util/autovector.h"
 
 namespace ROCKSDB_NAMESPACE {
@@ -71,6 +73,12 @@ class WriteThread {
     // A state indicating that the thread may be waiting using StateMutex()
     // and StateCondVar()
     STATE_LOCKED_WAITING = 32,
+
+    // The state used to inform a waiting writer that it has become a
+    // caller to call some other waiting writers to write to memtable
+    // by calling SetMemWritersEachStride. After doing
+    // this, it will also write to memtable.
+    STATE_PARALLEL_MEMTABLE_CALLER = 64,
   };
 
   struct Writer;
@@ -127,6 +135,7 @@ class WriteThread {
     uint64_t log_used;  // log number that this batch was inserted into
     uint64_t log_ref;   // log number that memtable insert should reference
     WriteCallback* callback;
+    UserWriteCallback* user_write_cb;
     bool made_waitable;          // records lazy construction of mutex and cv
     std::atomic<uint8_t> state;  // write under StateMutex() or pre-link
     WriteGroup* write_group;
@@ -134,8 +143,8 @@ class WriteThread {
     Status status;
     Status callback_status;  // status returned by callback->Callback()
 
-    std::aligned_storage<sizeof(std::mutex)>::type state_mutex_bytes;
-    std::aligned_storage<sizeof(std::condition_variable)>::type state_cv_bytes;
+    aligned_storage<std::mutex>::type state_mutex_bytes;
+    aligned_storage<std::condition_variable>::type state_cv_bytes;
     Writer* link_older;  // read/write only before linking, or as leader
     Writer* link_newer;  // lazy, read/write only before linking, or as leader
 
@@ -153,6 +162,7 @@ class WriteThread {
           log_used(0),
           log_ref(0),
           callback(nullptr),
+          user_write_cb(nullptr),
           made_waitable(false),
           state(STATE_INIT),
           write_group(nullptr),
@@ -161,8 +171,8 @@ class WriteThread {
           link_newer(nullptr) {}
 
     Writer(const WriteOptions& write_options, WriteBatch* _batch,
-           WriteCallback* _callback, uint64_t _log_ref, bool _disable_memtable,
-           size_t _batch_cnt = 0,
+           WriteCallback* _callback, UserWriteCallback* _user_write_cb,
+           uint64_t _log_ref, bool _disable_memtable, size_t _batch_cnt = 0,
            PreReleaseCallback* _pre_release_callback = nullptr,
            PostMemTableCallback* _post_memtable_callback = nullptr)
         : batch(_batch),
@@ -180,6 +190,7 @@ class WriteThread {
           log_used(0),
           log_ref(_log_ref),
           callback(_callback),
+          user_write_cb(_user_write_cb),
           made_waitable(false),
           state(STATE_INIT),
           write_group(nullptr),
@@ -203,6 +214,18 @@ class WriteThread {
       return callback_status.ok();
     }
 
+    void CheckWriteEnqueuedCallback() {
+      if (user_write_cb != nullptr) {
+        user_write_cb->OnWriteEnqueued();
+      }
+    }
+
+    void CheckPostWalWriteCallback() {
+      if (user_write_cb != nullptr) {
+        user_write_cb->OnWalWriteFinish();
+      }
+    }
+
     void CreateMutex() {
       if (!made_waitable) {
         // Note that made_waitable is tracked separately from state
@@ -323,10 +346,19 @@ class WriteThread {
   // Causes JoinBatchGroup to return STATE_PARALLEL_MEMTABLE_WRITER for all of
   // the non-leader members of this write batch group.  Sets Writer::sequence
   // before waking them up.
+  // If the size of write_group n is not small, the leader will call n^0.5
+  // members to be PARALLEL_MEMTABLE_CALLER in the write_group to help to set
+  // other's status parallel. This ensures that the cost to call SetState
+  // sequentially does not exceed 2(n^0.5).
   //
   // WriteGroup* write_group: Extra state used to coordinate the parallel add
   void LaunchParallelMemTableWriters(WriteGroup* write_group);
 
+  // One of the every stride=N number writer in the WriteGroup are set to the
+  // MemTableWriters, where N is equal to square of the total number of this
+  // write_group, and all of these MemTableWriters will write to memtable.
+  void SetMemWritersEachStride(Writer* w);
+
   // Reports the completion of w's batch to the parallel group leader, and
   // waits for the rest of the parallel batch to complete.  Returns true
   // if this thread is the last to complete, and hence should advance

package/vendor/librocksdb/vendor/rocksdb/env/composite_env_wrapper.h

@@ -18,6 +18,13 @@
 
 namespace ROCKSDB_NAMESPACE {
 
+// This class supports abstracting different types of an `Env`'s functionality
+// into separate interfaces. It is constructed with a `FileSystem` and a
+// `SystemClock` and delegates:
+// * File system operations to member `file_system_`.
+// * Time related misc operations to member `clock_`.
+// A subclass needs to inherit `CompositeEnv` and provide implementations for
+// the thread management related APIs.
 class CompositeEnv : public Env {
  public:
   // Initialize a CompositeEnvWrapper that delegates all thread/time related
@@ -250,6 +257,20 @@ class CompositeEnv : public Env {
   }
 };
 
+// A `CompositeEnvWrapper` is constructed with a target `Env` object, an
+// optional `FileSystem` object and an optional `SystemClock` object.
+// `Env::GetFileSystem()` is a fallback file system if no such object is
+// explicitly provided. Similarly, `Env::GetSystemClock()` is a fallback system
+// clock.
+// Besides delegating corresponding functionality to `file_system_` and `clock_`
+// which is inherited from `CompositeEnv`, it also implements the thread
+// management APIs by delegating them to the target `Env` object.
+//
+// Effectively, this class helps to support using customized file system
+// implementations such as a remote file system instead of the default file
+// system provided by the operating system.
+//
+// Also see public API `NewCompositeEnv` in rocksdb/include/env.h
 class CompositeEnvWrapper : public CompositeEnv {
  public:
   // Initialize a CompositeEnvWrapper that delegates all thread/time related

package/vendor/librocksdb/vendor/rocksdb/env/env.cc

@@ -355,6 +355,12 @@ class LegacyDirectoryWrapper : public FSDirectory {
   std::unique_ptr<Directory> target_;
 };
 
+// A helper class to make legacy `Env` implementations be backward compatible
+// now that all `Env` implementations are expected to have a `FileSystem` type
+// member `file_system_` and a `SystemClock` type member `clock_`.
+// This class wraps a legacy `Env` object and expose its file system related
+// APIs as a `FileSystem` interface. Also check `LegacySystemClock` that does
+// the same thing for the clock related APIs.
 class LegacyFileSystemWrapper : public FileSystem {
  public:
   // Initialize an EnvWrapper that delegates all calls to *t
@@ -825,6 +831,15 @@ WritableFile::~WritableFile() = default;
 
 MemoryMappedFileBuffer::~MemoryMappedFileBuffer() = default;
 
+// This const variable can be used in public headers without introducing the
+// possibility of ODR violations due to varying macro definitions.
+const InfoLogLevel Logger::kDefaultLogLevel =
+#ifdef NDEBUG
+    INFO_LEVEL;
+#else
+    DEBUG_LEVEL;
+#endif  // NDEBUG
+
 Logger::~Logger() = default;
 
 Status Logger::Close() {

package/vendor/librocksdb/vendor/rocksdb/env/fs_on_demand.cc

@@ -0,0 +1,331 @@
+// Copyright (c) 2024-present, Facebook, Inc.  All rights reserved.
+//  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 "env/fs_on_demand.h"
+
+#include <algorithm>
+#include <set>
+
+#include "file/filename.h"
+#include "port/port.h"
+#include "rocksdb/types.h"
+
+namespace ROCKSDB_NAMESPACE {
+// Check if the input path is under orig (typically the local directory), and if
+// so, change it to the equivalent path under replace (typically the remote
+// directory). For example, if orig is "/data/follower", replace is
+// "/data/leader", and the given path is "/data/follower/000010.sst", on return
+// the path would be changed to
+// "/data/leader/000010.sst".
+// Return value is true if the path was modified, false otherwise
+bool OnDemandFileSystem::CheckPathAndAdjust(const std::string& orig,
+                                            const std::string& replace,
+                                            std::string& path) {
+  size_t pos = path.find(orig);
+  if (pos > 0) {
+    return false;
+  }
+  path.replace(pos, orig.length(), replace);
+  return true;
+}
+
+bool OnDemandFileSystem::LookupFileType(const std::string& name,
+                                        FileType* type) {
+  std::size_t found = name.find_last_of('/');
+  std::string file_name = name.substr(found);
+  uint64_t number = 0;
+  return ParseFileName(file_name, &number, type);
+}
+
+// RocksDB opens non-SST files for reading in sequential file mode. This
+// includes CURRENT, OPTIONS, MANIFEST etc. For these files, we open them
+// in place in the source directory. For files that are appendable or
+// can be renamed, which is MANIFEST and CURRENT files, we wrap the
+// underlying FSSequentialFile with another class that checks when EOF
+// has been reached and re-opens the file to see the latest data. On some
+// distributed file systems, this is necessary.
+IOStatus OnDemandFileSystem::NewSequentialFile(
+    const std::string& fname, const FileOptions& file_opts,
+    std::unique_ptr<FSSequentialFile>* result, IODebugContext* dbg) {
+  FileType type;
+  static std::unordered_set<FileType> valid_types(
+      {kWalFile, kDescriptorFile, kCurrentFile, kIdentityFile, kOptionsFile});
+  if (!LookupFileType(fname, &type) ||
+      (valid_types.find(type) == valid_types.end())) {
+    return IOStatus::NotSupported();
+  }
+
+  IOStatus s;
+  std::string rname = fname;
+  if (CheckPathAndAdjust(local_path_, remote_path_, rname)) {
+    // First clear any local directory cache as it may be out of date
+    target()->DiscardCacheForDirectory(rname);
+
+    std::unique_ptr<FSSequentialFile> inner_file;
+    s = target()->NewSequentialFile(rname, file_opts, &inner_file, dbg);
+    if (s.ok() && type == kDescriptorFile) {
+      result->reset(new OnDemandSequentialFile(std::move(inner_file), this,
+                                               file_opts, rname));
+    } else {
+      *result = std::move(inner_file);
+    }
+  } else {
+    s = target()->NewSequentialFile(fname, file_opts, result, dbg);
+  }
+  return s;
+}
+
+// This is only supported for SST files. If the file is present locally,
+// i.e in the destination dir, we just open it and return. If its in the
+// remote, i.e source dir, we link it locally and open the link.
+// TODO: Add support for blob files belonging to the new BlobDB
+IOStatus OnDemandFileSystem::NewRandomAccessFile(
+    const std::string& fname, const FileOptions& file_opts,
+    std::unique_ptr<FSRandomAccessFile>* result, IODebugContext* dbg) {
+  FileType type;
+  if (!LookupFileType(fname, &type) || type != kTableFile) {
+    return IOStatus::NotSupported();
+  }
+
+  IOStatus s = target()->FileExists(fname, file_opts.io_options, nullptr);
+  if (s.IsNotFound() || s.IsPathNotFound()) {
+    std::string rname = fname;
+    if (CheckPathAndAdjust(local_path_, remote_path_, rname)) {
+      // First clear any local directory cache as it may be out of date
+      target()->DiscardCacheForDirectory(rname);
+
+      s = target()->LinkFile(rname, fname, IOOptions(), nullptr);
+      if (!s.ok()) {
+        return s;
+      }
+    }
+  }
+
+  return s.ok() ? target()->NewRandomAccessFile(fname, file_opts, result, dbg)
+                : s;
+}
+
+// We don't expect to create any writable file other than info LOG files.
+IOStatus OnDemandFileSystem::NewWritableFile(
+    const std::string& fname, const FileOptions& file_opts,
+    std::unique_ptr<FSWritableFile>* result, IODebugContext* dbg) {
+  FileType type;
+  if (!LookupFileType(fname, &type) || type != kInfoLogFile) {
+    return IOStatus::NotSupported();
+  }
+
+  std::string rname = fname;
+  if (CheckPathAndAdjust(local_path_, remote_path_, rname)) {
+    // First clear any local directory cache as it may be out of date
+    target()->DiscardCacheForDirectory(rname);
+
+    IOStatus s = target()->FileExists(rname, file_opts.io_options, dbg);
+    if (s.ok()) {
+      return IOStatus::InvalidArgument(
+          "Writing to a file present in the remote directory not supoprted");
+    }
+  }
+
+  return target()->NewWritableFile(fname, file_opts, result, dbg);
+}
+
+// Currently not supported, as there's no need for RocksDB to create a
+// directory object for a DB in read-only mode.
+IOStatus OnDemandFileSystem::NewDirectory(
+    const std::string& /*name*/, const IOOptions& /*io_opts*/,
+    std::unique_ptr<FSDirectory>* /*result*/, IODebugContext* /*dbg*/) {
+  return IOStatus::NotSupported();
+}
+
+// Check if the given file exists, either locally or remote. If the file is an
+// SST file, then link it locally. We assume if the file existence is being
+// checked, its for verification purposes, for example while replaying the
+// MANIFEST. The file will be opened for reading  some time in the future.
+IOStatus OnDemandFileSystem::FileExists(const std::string& fname,
+                                        const IOOptions& options,
+                                        IODebugContext* dbg) {
+  IOStatus s = target()->FileExists(fname, options, dbg);
+  if (!s.IsNotFound() && !s.IsPathNotFound()) {
+    return s;
+  }
+
+  std::string rname = fname;
+  if (CheckPathAndAdjust(local_path_, remote_path_, rname)) {
+    // First clear any local directory cache as it may be out of date
+    target()->DiscardCacheForDirectory(rname);
+
+    FileType type;
+    if (LookupFileType(fname, &type) && type == kTableFile) {
+      s = target()->LinkFile(rname, fname, options, dbg);
+    } else {
+      s = target()->FileExists(rname, options, dbg);
+    }
+  }
+  return s;
+}
+
+// Doa listing of both the local and remote directories and merge the two.
+IOStatus OnDemandFileSystem::GetChildren(const std::string& dir,
+                                         const IOOptions& options,
+                                         std::vector<std::string>* result,
+                                         IODebugContext* dbg) {
+  std::string rdir = dir;
+  IOStatus s = target()->GetChildren(dir, options, result, dbg);
+  if (!s.ok() || !CheckPathAndAdjust(local_path_, remote_path_, rdir)) {
+    return s;
+  }
+
+  std::vector<std::string> rchildren;
+  // First clear any local directory cache as it may be out of date
+  target()->DiscardCacheForDirectory(rdir);
+  s = target()->GetChildren(rdir, options, &rchildren, dbg);
+  if (s.ok()) {
+    std::for_each(rchildren.begin(), rchildren.end(), [&](std::string& name) {
+      // Adjust name
+      (void)CheckPathAndAdjust(remote_path_, local_path_, name);
+    });
+    std::sort(result->begin(), result->end());
+    std::sort(rchildren.begin(), rchildren.end());
+
+    std::vector<std::string> output;
+    output.reserve(result->size() + rchildren.size());
+    std::set_union(result->begin(), result->end(), rchildren.begin(),
+                   rchildren.end(), std::back_inserter(output));
+    *result = std::move(output);
+  }
+  return s;
+}
+
+// Doa listing of both the local and remote directories and merge the two.
+IOStatus OnDemandFileSystem::GetChildrenFileAttributes(
+    const std::string& dir, const IOOptions& options,
+    std::vector<FileAttributes>* result, IODebugContext* dbg) {
+  std::string rdir = dir;
+  IOStatus s = target()->GetChildrenFileAttributes(dir, options, result, dbg);
+  if (!s.ok() || !CheckPathAndAdjust(local_path_, remote_path_, rdir)) {
+    return s;
+  }
+
+  std::vector<FileAttributes> rchildren;
+  // First clear any local directory cache as it may be out of date
+  target()->DiscardCacheForDirectory(rdir);
+  s = target()->GetChildrenFileAttributes(rdir, options, &rchildren, dbg);
+  if (s.ok()) {
+    struct FileAttributeSorter {
+      bool operator()(const FileAttributes& lhs, const FileAttributes& rhs) {
+        return lhs.name < rhs.name;
+      }
+    } file_attr_sorter;
+
+    std::for_each(
+        rchildren.begin(), rchildren.end(), [&](FileAttributes& file) {
+          // Adjust name
+          (void)CheckPathAndAdjust(remote_path_, local_path_, file.name);
+        });
+    std::sort(result->begin(), result->end(), file_attr_sorter);
+    std::sort(rchildren.begin(), rchildren.end(), file_attr_sorter);
+
+    std::vector<FileAttributes> output;
+    output.reserve(result->size() + rchildren.size());
+    std::set_union(rchildren.begin(), rchildren.end(), result->begin(),
+                   result->end(), std::back_inserter(output), file_attr_sorter);
+    *result = std::move(output);
+  }
+  return s;
+}
+
+IOStatus OnDemandFileSystem::GetFileSize(const std::string& fname,
+                                         const IOOptions& options,
+                                         uint64_t* file_size,
+                                         IODebugContext* dbg) {
+  uint64_t local_size = 0;
+  IOStatus s = target()->GetFileSize(fname, options, &local_size, dbg);
+  if (!s.ok() && !s.IsNotFound() && !s.IsPathNotFound()) {
+    return s;
+  }
+
+  if (s.IsNotFound() || s.IsPathNotFound()) {
+    std::string rname = fname;
+    if (CheckPathAndAdjust(local_path_, remote_path_, rname)) {
+      // First clear any local directory cache as it may be out of date
+      target()->DiscardCacheForDirectory(rname);
+
+      FileType type;
+      if (LookupFileType(fname, &type) && type == kTableFile) {
+        s = target()->LinkFile(rname, fname, options, dbg);
+        if (s.ok()) {
+          s = target()->GetFileSize(fname, options, &local_size, dbg);
+        }
+      } else {
+        s = target()->GetFileSize(rname, options, &local_size, dbg);
+      }
+    }
+  }
+  *file_size = local_size;
+  return s;
+}
+
+// An implementation of Read that tracks whether we've reached EOF. If so,
+// re-open the file to try to read past the previous EOF offset. After
+// re-opening, positing it back to the last read offset.
+IOStatus OnDemandSequentialFile::Read(size_t n, const IOOptions& options,
+                                      Slice* result, char* scratch,
+                                      IODebugContext* dbg) {
+  IOStatus s;
+  if (eof_) {
+    // Reopen the file. With some distributed file systems, this is required
+    // in order to get the new size
+    file_.reset();
+    s = fs_->NewSequentialFile(path_, file_opts_, &file_, dbg);
+    if (!s.ok()) {
+      return IOStatus::IOError("While opening file after relinking, got error ",
+                               s.ToString());
+    }
+    s = file_->Skip(offset_);
+    if (!s.ok()) {
+      return IOStatus::IOError(
+          "While seeking to offset" + std::to_string(offset_) + "got error",
+          s.ToString());
+    }
+    eof_ = false;
+  }
+
+  s = file_->Read(n, options, result, scratch, dbg);
+  if (s.ok()) {
+    offset_ += result->size();
+    if (result->size() < n) {
+      // We reached EOF. Mark it so we know to relink next time
+      eof_ = true;
+    }
+  }
+  return s;
+}
+
+IOStatus OnDemandSequentialFile::Skip(uint64_t n) {
+  IOStatus s = file_->Skip(n);
+  if (s.ok()) {
+    offset_ += n;
+  }
+  return s;
+}
+
+bool OnDemandSequentialFile::use_direct_io() const {
+  return file_->use_direct_io();
+}
+
+size_t OnDemandSequentialFile::GetRequiredBufferAlignment() const {
+  return file_->GetRequiredBufferAlignment();
+}
+
+Temperature OnDemandSequentialFile::GetTemperature() const {
+  return file_->GetTemperature();
+}
+
+std::shared_ptr<FileSystem> NewOnDemandFileSystem(
+    const std::shared_ptr<FileSystem>& fs, std::string src_path,
+    std::string dest_path) {
+  return std::make_shared<OnDemandFileSystem>(fs, src_path, dest_path);
+}
+}  // namespace ROCKSDB_NAMESPACE