RCBC-512: Invoke fork hooks to protect SDK internal state

Author: avsej

.gitignore

@@ -11,7 +11,7 @@
 /.yardoc
 /_yardoc/
 /cmake-build-*/
-/build-*/
+/build*
 /coverage/
 /doc/
 /ext/cache/

ext/.gitignore

@@ -1,7 +1,7 @@
 /.idea/sonarlint/
 /.idea/editor.xml
 /.idea/workspace.xml
-/build/
+/build*
 /cmake-build-*/
 /cmake-build-report.tar.gz
 /revisions.rb

ext/couchbase

@@ -41,6 +41,7 @@ __declspec(dllexport)
 void
 Init_libcouchbase(void)
 {
+  couchbase::ruby::install_terminate_handler();
   couchbase::ruby::init_logger();
 
   VALUE mCouchbase = rb_define_module("Couchbase");

ext/couchbase.cxx

@@ -154,6 +154,16 @@ def sys(*cmd)
             File.join(Dir.tmpdir, "cb-#{build_type}-#{RUBY_VERSION}-#{RUBY_PATCHLEVEL}-#{RUBY_PLATFORM}-#{SDK_VERSION}")
 FileUtils.rm_rf(build_dir, verbose: true) unless ENV['CB_PRESERVE_BUILD_DIR']
 FileUtils.mkdir_p(build_dir, verbose: true)
+if ENV["CB_CREATE_BUILD_DIR_LINK"]
+  links = [
+    File.expand_path(File.join(project_path, "..", "build")),
+    File.expand_path(File.join(project_path, "build"))
+  ]
+  links.each do |link|
+    next if link == build_dir
+    FileUtils.ln_sf(build_dir, link, verbose: true)
+  end
+end
 Dir.chdir(build_dir) do
   puts "-- build #{build_type} extension #{SDK_VERSION} for ruby #{RUBY_VERSION}-#{RUBY_PATCHLEVEL}-#{RUBY_PLATFORM}"
   sys(cmake, *cmake_flags, "-B#{build_dir}", "-S#{project_path}")

ext/extconf.rb

@@ -17,6 +17,9 @@
 
 #include <couchbase/cluster.hxx>
 
+#include <couchbase/fork_event.hxx>
+#include <couchbase/ip_protocol.hxx>
+
 #include <core/cluster.hxx>
 #include <core/logger/logger.hxx>
 #include <core/utils/connection_string.hxx>
@@ -25,11 +28,12 @@
 #include <spdlog/fmt/bundled/core.h>
 
 #include <future>
+#include <list>
 #include <memory>
+#include <mutex>
 
 #include <ruby.h>
 
-#include "couchbase/ip_protocol.hxx"
 #include "rcb_backend.hxx"
 #include "rcb_exceptions.hxx"
 #include "rcb_logger.hxx"
@@ -44,10 +48,79 @@ struct cb_backend_data {
   std::unique_ptr<cluster> instance{ nullptr };
 };
 
+class instance_registry
+{
+public:
+  void add(cluster* instance)
+  {
+    std::scoped_lock lock(instances_mutex_);
+    known_instances_.push_back(instance);
+  }
+
+  void remove(cluster* instance)
+  {
+    std::scoped_lock lock(instances_mutex_);
+    known_instances_.remove(instance);
+  }
+
+  void notify_fork(couchbase::fork_event event)
+  {
+    if (event != couchbase::fork_event::prepare) {
+      init_logger();
+    }
+
+    {
+      std::scoped_lock lock(instances_mutex_);
+      for (auto* instance : known_instances_) {
+        instance->notify_fork(event);
+      }
+    }
+
+    if (event == couchbase::fork_event::prepare) {
+      flush_logger();
+      couchbase::core::logger::shutdown();
+    }
+  }
+
+private:
+  std::mutex instances_mutex_;
+  std::list<cluster*> known_instances_;
+};
+
+instance_registry instances;
+
+VALUE
+cb_Backend_notify_fork(VALUE self, VALUE event)
+{
+  static const auto id_prepare{ rb_intern("prepare") };
+  static const auto id_parent{ rb_intern("parent") };
+  static const auto id_child{ rb_intern("child") };
+
+  try {
+    cb_check_type(event, T_SYMBOL);
+
+    if (rb_sym2id(event) == id_prepare) {
+      instances.notify_fork(couchbase::fork_event::prepare);
+    } else if (rb_sym2id(event) == id_parent) {
+      instances.notify_fork(couchbase::fork_event::parent);
+    } else if (rb_sym2id(event) == id_child) {
+      instances.notify_fork(couchbase::fork_event::child);
+    } else {
+      throw ruby_exception(rb_eTypeError,
+                           rb_sprintf("unexpected fork event type %" PRIsVALUE "", event));
+    }
+  } catch (const ruby_exception& e) {
+    rb_exc_raise(e.exception_object());
+  }
+
+  return Qnil;
+}
+
 void
 cb_backend_close(cb_backend_data* backend)
 {
   if (auto instance = std::move(backend->instance); instance) {
+    instances.remove(instance.get());
     auto promise = std::make_shared<std::promise<void>>();
     auto f = promise->get_future();
     instance->close([promise = std::move(promise)]() mutable {
@@ -446,6 +519,7 @@ cb_Backend_open(VALUE self, VALUE connstr, VALUE credentials, VALUE options)
         error, fmt::format("failed to connect to the Couchbase Server \"{}\"", connection_string));
     }
     backend->instance = std::make_unique<couchbase::cluster>(std::move(cluster));
+    instances.add(backend->instance.get());
   } catch (const std::system_error& se) {
     rb_exc_raise(cb_map_error_code(
       se.code(), fmt::format("failed to perform {}: {}", __func__, se.what()), false));
@@ -509,6 +583,8 @@ init_backend(VALUE mCouchbase)
   rb_define_method(cBackend, "open", cb_Backend_open, 3);
   rb_define_method(cBackend, "open_bucket", cb_Backend_open_bucket, 2);
   rb_define_method(cBackend, "close", cb_Backend_close, 0);
+
+  rb_define_singleton_method(cBackend, "notify_fork", cb_Backend_notify_fork, 1);
   return cBackend;
 }
 

ext/rcb_backend.cxx

@@ -268,13 +268,18 @@ cb_Backend_install_logger_shim(VALUE self, VALUE logger, VALUE log_level)
 } // namespace
 
 void
-init_logger()
+install_terminate_handler()
 {
   if (auto env_val =
         spdlog::details::os::getenv("COUCHBASE_BACKEND_DONT_INSTALL_TERMINATE_HANDLER");
       env_val.empty()) {
     core::platform::install_backtrace_terminate_handler();
   }
+}
+
+void
+init_logger()
+{
   if (auto env_val = spdlog::details::os::getenv("COUCHBASE_BACKEND_DONT_USE_BUILTIN_LOGGER");
       env_val.empty()) {
     auto default_log_level = core::logger::level::info;

ext/rcb_logger.cxx

@@ -22,6 +22,9 @@
 
 namespace couchbase::ruby
 {
+void
+install_terminate_handler();
+
 void
 init_logger();
 

ext/rcb_logger.hxx

@@ -245,13 +245,11 @@ cb_Backend_document_scan_create(VALUE self,
     std::promise<tl::expected<couchbase::core::topology::configuration, std::error_code>> promise;
     auto f = promise.get_future();
     cluster.with_bucket_configuration(
-      bucket_name,
-      [promise = std::move(promise)](
-        std::error_code ec, const couchbase::core::topology::configuration& config) mutable {
+      bucket_name, [promise = std::move(promise)](std::error_code ec, const auto& config) mutable {
         if (ec) {
           return promise.set_value(tl::unexpected(ec));
         }
-        promise.set_value(config);
+        promise.set_value(*config);
       });
     auto config = cb_wait_for_future(f);
     if (!config.has_value()) {

ext/rcb_range_scan.cxx

@@ -16,6 +16,7 @@
 
 require "couchbase/version"
 require "couchbase/libcouchbase"
+require "couchbase/fork_hooks"
 require "couchbase/logger"
 require "couchbase/cluster"
 require "couchbase/deprecations"

lib/couchbase.rb

@@ -0,0 +1,32 @@
+# frozen_string_literal: true
+
+#  Copyright 2020-2025 Couchbase, Inc.
+#
+#  Licensed under the Apache License, Version 2.0 (the "License");
+#  you may not use this file except in compliance with the License.
+#  You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+#  Unless required by applicable law or agreed to in writing, software
+#  distributed under the License is distributed on an "AS IS" BASIS,
+#  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+#  See the License for the specific language governing permissions and
+#  limitations under the License.
+
+module Couchbase
+  module ForkHooks
+    def _fork
+      Couchbase::Backend.notify_fork(:prepare)
+      pid = super
+      if pid
+        Couchbase::Backend.notify_fork(:parent)
+      else
+        Couchbase::Backend.notify_fork(:child)
+      end
+      pid
+    end
+  end
+end
+
+Process.singleton_class.prepend(Couchbase::ForkHooks)