1 files changed, 45 insertions, 34 deletions
diff --git a/src/google/protobuf/stubs/once.cc b/src/google/protobuf/stubs/once.cc
index 5b7af9c..1e24b85 100644
--- a/src/google/protobuf/stubs/once.cc
+++ b/src/google/protobuf/stubs/once.cc
@@ -35,54 +35,65 @@
 // This header is intended to be included only by internal .cc files and
 // generated .pb.cc files.  Users should not use this directly.
 
+#include <google/protobuf/stubs/once.h>
+
+#ifndef GOOGLE_PROTOBUF_NO_THREAD_SAFETY
+
 #ifdef _WIN32
 #include <windows.h>
+#else
+#include <sched.h>
 #endif
 
-#include <google/protobuf/stubs/once.h>
+#include <google/protobuf/stubs/atomicops.h>
 
 namespace google {
 namespace protobuf {
 
-#ifdef _WIN32
-
-struct ProtobufOnceInternal {
-  ProtobufOnceInternal() {
-    InitializeCriticalSection(&critical_section);
-  }
-  ~ProtobufOnceInternal() {
-    DeleteCriticalSection(&critical_section);
-  }
-  CRITICAL_SECTION critical_section;
-};
-
-ProtobufOnceType::~ProtobufOnceType()
-{
-  delete internal_;
-  internal_ = NULL;
-}
+namespace {
 
-ProtobufOnceType::ProtobufOnceType() {
-  // internal_ may be non-NULL if Init() was already called.
-  if (internal_ == NULL) internal_ = new ProtobufOnceInternal;
+void SchedYield() {
+#ifdef _WIN32
+  Sleep(0);
+#else  // POSIX
+  sched_yield();
+#endif
 }
 
-void ProtobufOnceType::Init(void (*init_func)()) {
-  // internal_ may be NULL if we're still in dynamic initialization and the
-  // constructor has not been called yet.  As mentioned in once.h, we assume
-  // that the program is still single-threaded at this time, and therefore it
-  // should be safe to initialize internal_ like so.
-  if (internal_ == NULL) internal_ = new ProtobufOnceInternal;
+}  // namespace
 
-  EnterCriticalSection(&internal_->critical_section);
-  if (!initialized_) {
-    init_func();
-    initialized_ = true;
+void GoogleOnceInitImpl(ProtobufOnceType* once, Closure* closure) {
+  internal::AtomicWord state = internal::Acquire_Load(once);
+  // Fast path. The provided closure was already executed.
+  if (state == ONCE_STATE_DONE) {
+    return;
+  }
+  // The closure execution did not complete yet. The once object can be in one
+  // of the two following states:
+  //   - UNINITIALIZED: We are the first thread calling this function.
+  //   - EXECUTING_CLOSURE: Another thread is already executing the closure.
+  //
+  // First, try to change the state from UNINITIALIZED to EXECUTING_CLOSURE
+  // atomically.
+  state = internal::Acquire_CompareAndSwap(
+      once, ONCE_STATE_UNINITIALIZED, ONCE_STATE_EXECUTING_CLOSURE);
+  if (state == ONCE_STATE_UNINITIALIZED) {
+    // We are the first thread to call this function, so we have to call the
+    // closure.
+    closure->Run();
+    internal::Release_Store(once, ONCE_STATE_DONE);
+  } else {
+    // Another thread has already started executing the closure. We need to
+    // wait until it completes the initialization.
+    while (state == ONCE_STATE_EXECUTING_CLOSURE) {
+      // Note that futex() could be used here on Linux as an improvement.
+      SchedYield();
+      state = internal::Acquire_Load(once);
+    }
   }
-  LeaveCriticalSection(&internal_->critical_section);
 }
 
-#endif
-
 }  // namespace protobuf
 }  // namespace google
+
+#endif  // GOOGLE_PROTOBUF_NO_THREAD_SAFETY