Add a FATAL message upon fork()-ing once tensorstore has started any threads.

PiperOrigin-RevId: 716823316
Change-Id: Ifcf976769fdbbd9fc2af8f4d79718055c33b0a3d

Author: laramiel

tensorstore/internal/thread/BUILD

@@ -45,9 +45,13 @@ tensorstore_cc_test(
 
 tensorstore_cc_library(
     name = "thread",
-    srcs = ["thread.cc"],
+    srcs = [
+        "thread.cc",
+        "thread_on_fork.cc",
+    ],
     hdrs = ["thread.h"],
     deps = [
+        "@com_google_absl//absl/base",
         "@com_google_absl//absl/functional:any_invocable",
         "@com_google_absl//absl/log:absl_check",
     ],
@@ -62,6 +66,15 @@ tensorstore_cc_test(
     ],
 )
 
+tensorstore_cc_test(
+    name = "thread_death_test",
+    srcs = ["thread_death_test.cc"],
+    deps = [
+        ":thread",
+        "@com_google_googletest//:gtest_main",
+    ],
+)
+
 THREAD_POOL_DEFINES = []
 
 THREAD_POOL_DEPS = []

tensorstore/internal/thread/thread.cc

@@ -12,26 +12,27 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
-#if defined(__linux__) || defined(__APPLE__)
+#if defined(__has_include)
+#if __has_include(<pthread.h>)
+#define TENSORSTORE_INTERNAL_USE_PTHREAD
 #include <pthread.h>
 #endif
+#endif
 
 #include <thread>  // NOLINT
 #include <type_traits>
 
 namespace tensorstore {
 namespace internal {
 
-/// See
-/// https://stackoverflow.com/questions/2369738/how-to-set-the-name-of-a-thread-in-linux-pthreads
-
 void TrySetCurrentThreadName(const char* name) {
   if (name == nullptr) return;
-#if defined(__linux__)
-  pthread_setname_np(pthread_self(), name);
-#endif
+#if defined(TENSORSTORE_INTERNAL_USE_PTHREAD)
 #if defined(__APPLE__)
   pthread_setname_np(name);
+#else
+  pthread_setname_np(pthread_self(), name);
+#endif
 #endif
   // TODO: Add windows via SetThreadDescription()
 }

tensorstore/internal/thread/thread.h

@@ -29,6 +29,9 @@ namespace internal {
 /// Helper functions to set the thread name.
 void TrySetCurrentThreadName(const char* name);
 
+/// Helper function to log a fatal error if fork() is called.
+void SetupLogFatalOnFork();
+
 // Tensorstore-specific Thread class to be used instead of std::thread.
 // This exposes a limited subset of the std::thread api.
 class Thread {
@@ -89,12 +92,13 @@ class Thread {
   // factory method.
   template <class Function, class... Args>
   Thread(private_t, Options options, Function&& f, Args&&... args)
-      : thread_(
+      : thread_((
+            SetupLogFatalOnFork(),
             [name = options.name, fn = std::bind(std::forward<Function>(f),
                                                  std::forward<Args>(args)...)] {
               TrySetCurrentThreadName(name);
               std::move(fn)();
-            }) {}
+            })) {}
 
   std::thread thread_;
 };

tensorstore/internal/thread/thread_death_test.cc

@@ -0,0 +1,48 @@
+// Copyright 2025 The TensorStore Authors
+//
+// 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.
+
+#if !defined(_WIN32) && !defined(__APPLE__) && defined(__has_include)
+#if __has_include(<pthread.h>)
+#define TENSORSTORE_INTERNAL_ENABLE_FORK_TEST 1
+// In order to run the death test, both pthreads and fork are required,
+// and the test must be run in gunit's thread-safe mode.
+#endif
+#endif
+
+#if defined(TENSORSTORE_INTERNAL_ENABLE_FORK_TEST)
+
+#include <pthread.h>
+#include <unistd.h>
+
+#include <gtest/gtest.h>
+#include "tensorstore/internal/thread/thread.h"
+
+namespace {
+
+TEST(ThreadDeathTest, Fork) {
+  GTEST_FLAG_SET(death_test_style, "threadsafe");
+  // Span a thread and join it; this should register pthread_at_fork()
+  // handler which will crash if fork() is called.
+  int x = 0;
+  tensorstore::internal::Thread my_thread({}, [&x]() { x = 1; });
+  my_thread.Join();
+  EXPECT_EQ(1, x);
+
+  // fork()-ing after starting a thread is not supported; forcibly crash.
+  EXPECT_DEATH(fork(), "");
+}
+
+}  // namespace
+
+#endif  // TENSORSTORE_INTERNAL_ENABLE_FORK_TEST

tensorstore/internal/thread/thread_on_fork.cc

@@ -0,0 +1,52 @@
+// Copyright 2025 The TensorStore Authors
+//
+// 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.
+
+#if !defined(_WIN32)
+#if defined(__has_include)
+#if __has_include(<pthread.h>)
+#define TENSORSTORE_INTERNAL_USE_PTHREAD
+#include <pthread.h>
+#endif
+#endif
+#endif
+
+#include <cstdio>
+#include <cstdlib>
+
+#include "absl/base/call_once.h"
+
+namespace tensorstore {
+namespace internal {
+namespace {
+
+[[noreturn]] void LogFatalOnFork() {
+  std::fprintf(stderr,
+               "aborting: fork() is not allowed since tensorstore uses "
+               "internal threading\n");
+  std::fflush(stderr);
+  std::abort();
+}
+
+}  // namespace
+
+void SetupLogFatalOnFork() {
+#if defined(TENSORSTORE_INTERNAL_USE_PTHREAD)
+  static absl::once_flag g_setup_pthread;
+  absl::call_once(g_setup_pthread, &pthread_atfork, LogFatalOnFork, nullptr,
+                  nullptr);
+#endif
+}
+
+}  // namespace internal
+}  // namespace tensorstore

tensorstore/util/status.cc

@@ -90,6 +90,7 @@ absl::Status MaybeAnnotateStatusImpl(absl::Status source,
                               SourceLocation loc) {
   std::fprintf(stderr, "%s:%d: %s: %s\n", loc.file_name(), loc.line(), message,
                status.ToString().c_str());
+  std::fflush(stderr);
   std::terminate();
 }