From 6155400cb15b0d99b2c257d57603d61d03a817a8 Mon Sep 17 00:00:00 2001
From: Stephen Celis <stephen@stephencelis.com>
Date: Thu, 26 Oct 2023 14:34:18 -0700
Subject: [PATCH] Lock async stream inits (#11)

* Lock async stream inits

Fixes #10.

* wip
---
 Sources/ConcurrencyExtras/AsyncStream.swift         |  9 +++++++--
 Sources/ConcurrencyExtras/AsyncThrowingStream.swift |  9 +++++++--
 Sources/ConcurrencyExtras/Internal/Locking.swift    | 11 +++++++++++
 3 files changed, 25 insertions(+), 4 deletions(-)
 create mode 100644 Sources/ConcurrencyExtras/Internal/Locking.swift

diff --git a/Sources/ConcurrencyExtras/AsyncStream.swift b/Sources/ConcurrencyExtras/AsyncStream.swift
index b8d4ad0..6dfa2ea 100644
--- a/Sources/ConcurrencyExtras/AsyncStream.swift
+++ b/Sources/ConcurrencyExtras/AsyncStream.swift
@@ -1,3 +1,5 @@
+import Foundation
+
 extension AsyncStream {
   /// Produces an `AsyncStream` from an `AsyncSequence` by consuming the sequence till it
   /// terminates, ignoring any failure.
@@ -51,10 +53,13 @@ extension AsyncStream {
   ///
   /// - Parameter sequence: An async sequence.
   public init<S: AsyncSequence>(_ sequence: S) where S.Element == Element {
+    let lock = NSLock()
     var iterator: S.AsyncIterator?
     self.init {
-      if iterator == nil {
-        iterator = sequence.makeAsyncIterator()
+      lock.withLock {
+        if iterator == nil {
+          iterator = sequence.makeAsyncIterator()
+        }
       }
       return try? await iterator?.next()
     }
diff --git a/Sources/ConcurrencyExtras/AsyncThrowingStream.swift b/Sources/ConcurrencyExtras/AsyncThrowingStream.swift
index 14fd40c..829b0ae 100644
--- a/Sources/ConcurrencyExtras/AsyncThrowingStream.swift
+++ b/Sources/ConcurrencyExtras/AsyncThrowingStream.swift
@@ -1,13 +1,18 @@
+import Foundation
+
 extension AsyncThrowingStream where Failure == Error {
   /// Produces an `AsyncThrowingStream` from an `AsyncSequence` by consuming the sequence till it
   /// terminates, rethrowing any failure.
   ///
   /// - Parameter sequence: An async sequence.
   public init<S: AsyncSequence>(_ sequence: S) where S.Element == Element {
+    let lock = NSLock()
     var iterator: S.AsyncIterator?
     self.init {
-      if iterator == nil {
-        iterator = sequence.makeAsyncIterator()
+      lock.withLock {
+        if iterator == nil {
+          iterator = sequence.makeAsyncIterator()
+        }
       }
       return try await iterator?.next()
     }
diff --git a/Sources/ConcurrencyExtras/Internal/Locking.swift b/Sources/ConcurrencyExtras/Internal/Locking.swift
new file mode 100644
index 0000000..a97a9d0
--- /dev/null
+++ b/Sources/ConcurrencyExtras/Internal/Locking.swift
@@ -0,0 +1,11 @@
+import Foundation
+
+#if !(os(iOS) || os(macOS) || os(tvOS) || os(watchOS))
+  extension NSLock {
+    func withLock<R>(_ body: () throws -> R) rethrows -> R {
+      self.lock()
+      defer { self.unlock() }
+      return try body()
+    }
+  }
+#endif