[embedded] Make embedded swift_once thread-safe

stdlib/public/core/EmbeddedRuntime.swift

@@ -196,23 +196,31 @@ fileprivate func refcountPointer(for object: UnsafeMutablePointer<HeapObject>) -
   return UnsafeMutablePointer<Int>(UnsafeRawPointer(object).advanced(by: MemoryLayout<Int>.size)._rawValue)
 }
 
-fileprivate func loadRelaxed(_ refcount: UnsafeMutablePointer<Int>) -> Int {
-  Int(Builtin.atomicload_monotonic_Word(refcount._rawValue))
+fileprivate func loadRelaxed(_ atomic: UnsafeMutablePointer<Int>) -> Int {
+  Int(Builtin.atomicload_monotonic_Word(atomic._rawValue))
 }
 
-fileprivate func loadAcquire(_ refcount: UnsafeMutablePointer<Int>) -> Int {
-  Int(Builtin.atomicload_acquire_Word(refcount._rawValue))
+fileprivate func loadAcquire(_ atomic: UnsafeMutablePointer<Int>) -> Int {
+  Int(Builtin.atomicload_acquire_Word(atomic._rawValue))
 }
 
-fileprivate func subFetchAcquireRelease(_ refcount: UnsafeMutablePointer<Int>, n: Int) -> Int {
-  let oldValue = Int(Builtin.atomicrmw_sub_acqrel_Word(refcount._rawValue, n._builtinWordValue))
+fileprivate func subFetchAcquireRelease(_ atomic: UnsafeMutablePointer<Int>, n: Int) -> Int {
+  let oldValue = Int(Builtin.atomicrmw_sub_acqrel_Word(atomic._rawValue, n._builtinWordValue))
   return oldValue - n
 }
 
-fileprivate func addRelaxed(_ refcount: UnsafeMutablePointer<Int>, n: Int) {
-  _ = Builtin.atomicrmw_add_monotonic_Word(refcount._rawValue, n._builtinWordValue)
+fileprivate func addRelaxed(_ atomic: UnsafeMutablePointer<Int>, n: Int) {
+  _ = Builtin.atomicrmw_add_monotonic_Word(atomic._rawValue, n._builtinWordValue)
 }
 
+fileprivate func compareExchangeRelaxed(_ atomic: UnsafeMutablePointer<Int>, expectedOldValue: Int, desiredNewValue: Int) -> Bool {
+  let (_, won) = Builtin.cmpxchg_monotonic_monotonic_Word(atomic._rawValue, expectedOldValue._builtinWordValue, desiredNewValue._builtinWordValue)
+  return Bool(won)
+}
+
+fileprivate func storeRelease(_ atomic: UnsafeMutablePointer<Int>, newValue: Int) {
+  Builtin.atomicstore_release_Word(atomic._rawValue, newValue._builtinWordValue)
+}
 
 
 /// Exclusivity checking
@@ -233,11 +241,18 @@ public func swift_endAccess(buffer: UnsafeMutableRawPointer) {
 
 @_silgen_name("swift_once")
 public func swift_once(predicate: UnsafeMutablePointer<Int>, fn: (@convention(c) (UnsafeMutableRawPointer)->()), context: UnsafeMutableRawPointer) {
-  // TODO/FIXME: The following only works in single-threaded environments.
-  if predicate.pointee == 0 {
-    predicate.pointee = 1
+  if loadAcquire(predicate) < 0 { return }
+
+  let won = compareExchangeRelaxed(predicate, expectedOldValue: 0, desiredNewValue: 1)
+  if won {
     fn(context)
-    predicate.pointee = -1
+    storeRelease(predicate, newValue: -1)
+    return
+  }
+
+  // TODO: This should really use an OS provided lock
+  while loadAcquire(predicate) >= 0 {
+    // spin
   }
 }
 

test/embedded/once-multithreaded.swift

@@ -0,0 +1,66 @@
+// RUN: %empty-directory(%t)
+// RUN: %{python} %utils/split_file.py -o %t %s
+
+// RUN: %target-swift-frontend %t/Main.swift %S/Inputs/print.swift -import-bridging-header %t/BridgingHeader.h -Rmodule-loading -enable-experimental-feature Embedded -c -o %t/main.o
+// RUN: %target-clang %t/main.o -o %t/a.out -dead_strip
+// RUN: %target-run %t/a.out | %FileCheck %s
+
+// REQUIRES: executable_test
+// REQUIRES: optimized_stdlib
+// REQUIRES: VENDOR=apple
+// REQUIRES: OS=macosx
+
+// BEGIN BridgingHeader.h
+
+typedef void *pthread_t;
+int pthread_create(pthread_t *thread, const void *attr, void *(*start_routine)(void *), void *arg);
+int pthread_join(pthread_t thread, void **value_ptr);
+unsigned int sleep(unsigned int);
+int usleep(unsigned int);
+
+// BEGIN Main.swift
+
+public struct MyStructWithAnExpensiveInitializer {
+  static var singleton = MyStructWithAnExpensiveInitializer()
+  static var n = 0
+
+  init() {
+    print("MyStructWithAnExpensiveInitializer.init")
+    usleep(100_000)
+    MyStructWithAnExpensiveInitializer.n += 1
+    print("MyStructWithAnExpensiveInitializer.init done")
+  }
+
+  func foo() {
+    precondition(MyStructWithAnExpensiveInitializer.n == 1)
+  }
+}
+
+@main
+struct Main {
+  static func main() {
+    print("Start")
+
+    var t1 = pthread_t(bitPattern: 0)
+    pthread_create(&t1, nil, { _ in
+      MyStructWithAnExpensiveInitializer.singleton.foo()
+      return nil
+    }, nil)
+
+    var t2 = pthread_t(bitPattern: 0)
+    pthread_create(&t2, nil, { _ in
+      MyStructWithAnExpensiveInitializer.singleton.foo()
+      return nil
+    }, nil)
+
+    pthread_join(t1, nil)
+    pthread_join(t2, nil)
+
+    print("All done")
+    
+    // CHECK: Start
+    // CHECK-NEXT: MyStructWithAnExpensiveInitializer.init
+    // CHECK-NEXT: MyStructWithAnExpensiveInitializer.init done
+    // CHECK-NEXT: All done
+  }
+}