Improve GoogleOnceInit() in Protocol Buffers.

It is based on V8's new CallOnce(): http://codereview.chromium.org/9447052/.

This patch includes the following changes:
  - POD (no static initializer generated) and faster implementation on Windows.
  - GoogleOnceInit() can now take an additional parameter which is forwarded to
    the function provided by the user.

This patch is part of the static initializers removal initiative.

2 files changed, 98 insertions, 63 deletions

diff --git a/src/google/protobuf/stubs/once.cc b/src/google/protobuf/stubs/once.cc
index 5b7af9ce..1e24b855 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
diff --git a/src/google/protobuf/stubs/once.h b/src/google/protobuf/stubs/once.h
index 0dee4076..c8cbe246 100644
--- a/src/google/protobuf/stubs/once.h
+++ b/src/google/protobuf/stubs/once.h
@@ -37,16 +37,22 @@
 //
 // This is basically a portable version of pthread_once().
 //
-// This header declares three things:
+// This header declares:
 // * A type called ProtobufOnceType.
 // * A macro GOOGLE_PROTOBUF_DECLARE_ONCE() which declares a variable of type
 //   ProtobufOnceType.  This is the only legal way to declare such a variable.
-//   The macro may only be used at the global scope (you cannot create local
-//   or class member variables of this type).
-// * A function GogoleOnceInit(ProtobufOnceType* once, void (*init_func)()).
+//   The macro may only be used at the global scope (you cannot create local or
+//   class member variables of this type).
+// * A function GoogleOnceInit(ProtobufOnceType* once, void (*init_func)()).
 //   This function, when invoked multiple times given the same ProtobufOnceType
 //   object, will invoke init_func on the first call only, and will make sure
 //   none of the calls return before that first call to init_func has finished.
+// * The user can provide a parameter which GoogleOnceInit() forwards to the
+//   user-provided function when it is called. Usage example:
+//     int a = 10;
+//     GoogleOnceInit(&my_once, &MyFunctionExpectingIntArgument, &a);
+// * This implementation guarantees that ProtobufOnceType is a POD (i.e. no
+//   static initializer generated).
 //
 // This implements a way to perform lazy initialization.  It's more efficient
 // than using mutexes as no lock is needed if initialization has already
@@ -72,50 +78,68 @@
 #ifndef GOOGLE_PROTOBUF_STUBS_ONCE_H__
 #define GOOGLE_PROTOBUF_STUBS_ONCE_H__
 
+#include <google/protobuf/stubs/atomicops.h>
 #include <google/protobuf/stubs/common.h>
 
-#ifndef _WIN32
-#include <pthread.h>
-#endif
-
 namespace google {
 namespace protobuf {
 
-#ifdef _WIN32
-
-struct ProtobufOnceInternal;
+#ifdef GOOGLE_PROTOBUF_NO_THREAD_SAFETY
 
-struct LIBPROTOBUF_EXPORT ProtobufOnceType {
-  ProtobufOnceType();
-  ~ProtobufOnceType();
-  void Init(void (*init_func)());
-
-  volatile bool initialized_;
-  ProtobufOnceInternal* internal_;
-};
+typedef bool ProtobufOnceType;
 
-#define GOOGLE_PROTOBUF_DECLARE_ONCE(NAME)                    \
-  ::google::protobuf::ProtobufOnceType NAME
+#define GOOGLE_PROTOBUF_ONCE_INIT false
 
 inline void GoogleOnceInit(ProtobufOnceType* once, void (*init_func)()) {
-  // Note:  Double-checked locking is safe on x86.
-  if (!once->initialized_) {
-    once->Init(init_func);
+  if (!*once) {
+    *once = true;
+    init_func();
+  }
+}
+
+template <typename Arg>
+inline void GoogleOnceInit(ProtobufOnceType* once, void (*init_func)(Arg),
+    Arg arg) {
+  if (!*once) {
+    *once = true;
+    init_func(arg);
   }
 }
 
 #else
 
-typedef pthread_once_t ProtobufOnceType;
+enum {
+  ONCE_STATE_UNINITIALIZED = 0,
+  ONCE_STATE_EXECUTING_CLOSURE = 1,
+  ONCE_STATE_DONE = 2
+};
 
-#define GOOGLE_PROTOBUF_DECLARE_ONCE(NAME)                    \
-  pthread_once_t NAME = PTHREAD_ONCE_INIT
+typedef internal::AtomicWord ProtobufOnceType;
+
+#define GOOGLE_PROTOBUF_ONCE_INIT ::google::protobuf::ONCE_STATE_UNINITIALIZED
+
+void GoogleOnceInitImpl(ProtobufOnceType* once, Closure* closure);
 
 inline void GoogleOnceInit(ProtobufOnceType* once, void (*init_func)()) {
-  pthread_once(once, init_func);
+  if (internal::Acquire_Load(once) != ONCE_STATE_DONE) {
+    internal::FunctionClosure0 func(init_func, false);
+    GoogleOnceInitImpl(once, &func);
+  }
 }
 
-#endif
+template <typename Arg>
+inline void GoogleOnceInit(ProtobufOnceType* once, void (*init_func)(Arg*),
+    Arg* arg) {
+  if (internal::Acquire_Load(once) != ONCE_STATE_DONE) {
+    internal::FunctionClosure1<Arg*> func(init_func, false, arg);
+    GoogleOnceInitImpl(once, &func);
+  }
+}
+
+#endif  // GOOGLE_PROTOBUF_NO_THREAD_SAFETY
+
+#define GOOGLE_PROTOBUF_DECLARE_ONCE(NAME) \
+  ::google::protobuf::ProtobufOnceType NAME = GOOGLE_PROTOBUF_ONCE_INIT
 
 }  // namespace protobuf
 }  // namespace google