Improve EventEmitter implementation

Sources/Auth/AuthClient.swift

@@ -5,10 +5,29 @@ import Foundation
   import FoundationNetworking
 #endif
 
+public final class AuthStateChangeListenerHandle {
+  var onCancel: (@Sendable () -> Void)?
+
+  public func cancel() {
+    onCancel?()
+    onCancel = nil
+  }
+
+  deinit {
+    cancel()
+  }
+}
+
+public typealias AuthStateChangeListener = @Sendable (
+  _ event: AuthChangeEvent,
+  _ session: Session?
+) -> Void
+
 public actor AuthClient {
   /// FetchHandler is a type alias for asynchronous network request handling.
-  public typealias FetchHandler =
-    @Sendable (_ request: URLRequest) async throws -> (Data, URLResponse)
+  public typealias FetchHandler = @Sendable (
+    _ request: URLRequest
+  ) async throws -> (Data, URLResponse)
 
   /// Configuration struct represents the client configuration.
   public struct Configuration: Sendable {
@@ -150,7 +169,7 @@ public actor AuthClient {
       sessionManager: .live,
       codeVerifierStorage: .live,
       api: api,
-      eventEmitter: .live,
+      eventEmitter: EventEmitter(),
       sessionStorage: .live,
       logger: configuration.logger
     )
@@ -187,19 +206,38 @@ public actor AuthClient {
     )
   }
 
+  /// Listen for auth state changes.
+  ///
+  /// An `.initialSession` is always emitted when this method is called.
+  @discardableResult
+  public func onAuthStateChange(
+    _ listener: @escaping AuthStateChangeListener
+  ) -> AuthStateChangeListenerHandle {
+    let handle = eventEmitter.attachListener(listener)
+    Task {
+      await emitInitialSession(forHandle: handle)
+    }
+    return handle
+  }
+
   /// Listen for auth state changes.
   ///
   /// An `.initialSession` is always emitted when this method is called.
   public var authStateChanges: AsyncStream<(
     event: AuthChangeEvent,
     session: Session?
   )> {
-    let (id, stream) = eventEmitter.attachListener()
-    logger?.debug("auth state change listener with id '\(id.uuidString)' attached.")
+    let (stream, continuation) = AsyncStream<(
+      event: AuthChangeEvent,
+      session: Session?
+    )>.makeStream()
+
+    let handle = onAuthStateChange { event, session in
+      continuation.yield((event, session))
+    }
 
-    Task { [id] in
-      await emitInitialSession(forStreamWithID: id)
-      logger?.debug("initial session for listener with id '\(id.uuidString)' emitted.")
+    continuation.onTermination = { _ in
+      handle.cancel()
     }
 
     return stream
@@ -884,9 +922,9 @@ public actor AuthClient {
     return session
   }
 
-  private func emitInitialSession(forStreamWithID id: UUID) async {
+  private func emitInitialSession(forHandle handle: AuthStateChangeListenerHandle) async {
     let session = try? await session
-    eventEmitter.emit(.initialSession, session, id)
+    eventEmitter.emit(.initialSession, session: session, handle: handle)
   }
 
   private func prepareForPKCE() -> (codeChallenge: String?, codeChallengeMethod: String?) {

Sources/Auth/Internal/EventEmitter.swift

@@ -1,62 +1,50 @@
 import ConcurrencyExtras
 import Foundation
 
-struct EventEmitter: Sendable {
-  var attachListener: @Sendable () -> (
-    id: UUID,
-    stream: AsyncStream<(event: AuthChangeEvent, session: Session?)>
-  )
-  var emit: @Sendable (_ event: AuthChangeEvent, _ session: Session?, _ id: UUID?) -> Void
-}
+class EventEmitter: @unchecked Sendable {
+  let listeners = LockIsolated<[ObjectIdentifier: AuthStateChangeListener]>([:])
+
+  func attachListener(_ listener: @escaping AuthStateChangeListener)
+    -> AuthStateChangeListenerHandle
+  {
+    let handle = AuthStateChangeListenerHandle()
+    let key = ObjectIdentifier(handle)
+
+    handle.onCancel = { [weak self] in
+      self?.listeners.withValue {
+        $0[key] = nil
+      }
+    }
 
-extension EventEmitter {
-  func emit(_ event: AuthChangeEvent, session: Session?) {
-    emit(event, session, nil)
+    listeners.withValue {
+      $0[key] = listener
+    }
+
+    return handle
   }
-}
 
-extension EventEmitter {
-  static var live: Self = {
-    let continuations = LockIsolated(
-      [UUID: AsyncStream<(event: AuthChangeEvent, session: Session?)>.Continuation]()
+  func emit(
+    _ event: AuthChangeEvent,
+    session: Session?,
+    handle: AuthStateChangeListenerHandle? = nil
+  ) {
+    NotificationCenter.default.post(
+      name: AuthClient.didChangeAuthStateNotification,
+      object: nil,
+      userInfo: [
+        AuthClient.authChangeEventInfoKey: event,
+        AuthClient.authChangeSessionInfoKey: session as Any,
+      ]
     )
 
-    return Self(
-      attachListener: {
-        let id = UUID()
-
-        let (stream, continuation) = AsyncStream<(event: AuthChangeEvent, session: Session?)>
-          .makeStream()
-
-        continuation.onTermination = { [id] _ in
-          continuations.withValue {
-            $0[id] = nil
-          }
-        }
-
-        continuations.withValue {
-          $0[id] = continuation
-        }
-
-        return (id, stream)
-      },
-      emit: { event, session, id in
-        NotificationCenter.default.post(
-          name: AuthClient.didChangeAuthStateNotification,
-          object: nil,
-          userInfo: [
-            AuthClient.authChangeEventInfoKey: event,
-            AuthClient.authChangeSessionInfoKey: session as Any,
-          ]
-        )
-        if let id {
-          continuations.value[id]?.yield((event, session))
-        } else {
-          for continuation in continuations.value.values {
-            continuation.yield((event, session))
-          }
-        }
+    let listeners = listeners.value
+
+    if let handle {
+      listeners[ObjectIdentifier(handle)]?(event, session)
+    } else {
+      for listener in listeners.values {
+        listener(event, session)
       }
-    )
-  }()
+    }
+  }
 }