1 files changed, 30 insertions, 18 deletions
diff --git a/src/core/lib/iomgr/fork_posix.cc b/src/core/lib/iomgr/fork_posix.cc
index f8645ab157..b37384b8db 100644
--- a/src/core/lib/iomgr/fork_posix.cc
+++ b/src/core/lib/iomgr/fork_posix.cc
@@ -28,7 +28,7 @@
 #include <grpc/support/log.h>
 
 #include "src/core/lib/gpr/env.h"
-#include "src/core/lib/gpr/fork.h"
+#include "src/core/lib/gprpp/fork.h"
 #include "src/core/lib/gprpp/thd.h"
 #include "src/core/lib/iomgr/ev_posix.h"
 #include "src/core/lib/iomgr/executor.h"
@@ -41,47 +41,59 @@
  *       AROUND VERY SPECIFIC USE CASES.
  */
 
+namespace {
+bool skipped_handler = true;
+bool registered_handlers = false;
+}  // namespace
+
 void grpc_prefork() {
-  if (!grpc_fork_support_enabled()) {
+  grpc_core::ExecCtx exec_ctx;
+  skipped_handler = true;
+  if (!grpc_is_initialized()) {
+    return;
+  }
+  if (!grpc_core::Fork::Enabled()) {
     gpr_log(GPR_ERROR,
             "Fork support not enabled; try running with the "
             "environment variable GRPC_ENABLE_FORK_SUPPORT=1");
     return;
   }
-  if (grpc_is_initialized()) {
-    grpc_core::ExecCtx exec_ctx;
-    grpc_timer_manager_set_threading(false);
-    grpc_executor_set_threading(false);
-    grpc_core::ExecCtx::Get()->Flush();
-    if (!grpc_core::Thread::AwaitAll(
-            gpr_time_add(gpr_now(GPR_CLOCK_REALTIME),
-                         gpr_time_from_seconds(3, GPR_TIMESPAN)))) {
-      gpr_log(GPR_ERROR, "gRPC thread still active! Cannot fork!");
-    }
+  if (!grpc_core::Fork::BlockExecCtx()) {
+    gpr_log(GPR_INFO,
+            "Other threads are currently calling into gRPC, skipping fork() "
+            "handlers");
+    return;
   }
+  grpc_timer_manager_set_threading(false);
+  grpc_executor_set_threading(false);
+  grpc_core::ExecCtx::Get()->Flush();
+  grpc_core::Fork::AwaitThreads();
+  skipped_handler = false;
 }
 
 void grpc_postfork_parent() {
-  if (grpc_is_initialized()) {
-    grpc_timer_manager_set_threading(true);
+  if (!skipped_handler) {
+    grpc_core::Fork::AllowExecCtx();
     grpc_core::ExecCtx exec_ctx;
+    grpc_timer_manager_set_threading(true);
     grpc_executor_set_threading(true);
   }
 }
 
 void grpc_postfork_child() {
-  if (grpc_is_initialized()) {
-    grpc_timer_manager_set_threading(true);
+  if (!skipped_handler) {
+    grpc_core::Fork::AllowExecCtx();
     grpc_core::ExecCtx exec_ctx;
+    grpc_timer_manager_set_threading(true);
     grpc_executor_set_threading(true);
-    grpc_core::ExecCtx::Get()->Flush();
   }
 }
 
 void grpc_fork_handlers_auto_register() {
-  if (grpc_fork_support_enabled()) {
+  if (grpc_core::Fork::Enabled() & !registered_handlers) {
 #ifdef GRPC_POSIX_FORK_ALLOW_PTHREAD_ATFORK
     pthread_atfork(grpc_prefork, grpc_postfork_parent, grpc_postfork_child);
+    registered_handlers = true;
 #endif  // GRPC_POSIX_FORK_ALLOW_PTHREAD_ATFORK
   }
 }