From ffc93a95220a3a9f27a2b926c0e3fcd6e33573a8 Mon Sep 17 00:00:00 2001
From: Guilherme Souza <grsouza@pm.me>
Date: Tue, 20 Feb 2024 05:20:23 -0300
Subject: [PATCH] feat: add AuthStateChangeListenerRegistration type

---
 Sources/Auth/AuthClient.swift                 | 22 +---------
 Sources/Auth/AuthStateChangeListener.swift    | 41 +++++++++++++++++
 Sources/Auth/Internal/EventEmitter.swift      |  2 +-
 .../AuthStateChangeListenerHandleTests.swift  | 44 +++++++++++++++++++
 4 files changed, 88 insertions(+), 21 deletions(-)
 create mode 100644 Sources/Auth/AuthStateChangeListener.swift
 create mode 100644 Tests/AuthTests/AuthStateChangeListenerHandleTests.swift

diff --git a/Sources/Auth/AuthClient.swift b/Sources/Auth/AuthClient.swift
index 9d3b579cb..2a2eaf735 100644
--- a/Sources/Auth/AuthClient.swift
+++ b/Sources/Auth/AuthClient.swift
@@ -5,24 +5,6 @@ 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 (
@@ -216,7 +198,7 @@ public actor AuthClient {
   @discardableResult
   public func onAuthStateChange(
     _ listener: @escaping AuthStateChangeListener
-  ) async -> AuthStateChangeListenerHandle {
+  ) async -> AuthStateChangeListenerRegistration {
     let handle = eventEmitter.attachListener(listener)
     await emitInitialSession(forHandle: handle)
     return handle
@@ -240,7 +222,7 @@ public actor AuthClient {
       }
 
       continuation.onTermination = { _ in
-        handle.cancel()
+        handle.remove()
       }
     }
 
diff --git a/Sources/Auth/AuthStateChangeListener.swift b/Sources/Auth/AuthStateChangeListener.swift
new file mode 100644
index 000000000..8b053653e
--- /dev/null
+++ b/Sources/Auth/AuthStateChangeListener.swift
@@ -0,0 +1,41 @@
+//
+//  AuthStateChangeListener.swift
+//
+//
+//  Created by Guilherme Souza on 17/02/24.
+//
+
+import ConcurrencyExtras
+import Foundation
+
+/// A listener that can be removed by calling ``AuthStateChangeListenerRegistration/remove()``.
+///
+/// - Note: Listener is automatically removed on deinit.
+public protocol AuthStateChangeListenerRegistration: Sendable, AnyObject {
+  /// Removes the listener. After the initial call, subsequent calls have no effect.
+  func remove()
+}
+
+final class AuthStateChangeListenerHandle: AuthStateChangeListenerRegistration {
+  let _onRemove = LockIsolated((@Sendable () -> Void)?.none)
+
+  public func remove() {
+    _onRemove.withValue {
+      if $0 == nil {
+        return
+      }
+
+      $0?()
+      $0 = nil
+    }
+  }
+
+  deinit {
+    remove()
+  }
+}
+
+public typealias AuthStateChangeListener = @Sendable (
+  _ event: AuthChangeEvent,
+  _ session: Session?
+) -> Void
diff --git a/Sources/Auth/Internal/EventEmitter.swift b/Sources/Auth/Internal/EventEmitter.swift
index 3dcbc53d5..6b1ba0174 100644
--- a/Sources/Auth/Internal/EventEmitter.swift
+++ b/Sources/Auth/Internal/EventEmitter.swift
@@ -35,7 +35,7 @@ final class DefaultEventEmitter: EventEmitter {
     let handle = AuthStateChangeListenerHandle()
     let key = ObjectIdentifier(handle)
 
-    handle.onCancel = { [weak self] in
+    handle._onRemove.setValue { [weak self] in
       self?.listeners.withValue {
         $0[key] = nil
       }
diff --git a/Tests/AuthTests/AuthStateChangeListenerHandleTests.swift b/Tests/AuthTests/AuthStateChangeListenerHandleTests.swift
new file mode 100644
index 000000000..a5d9048c2
--- /dev/null
+++ b/Tests/AuthTests/AuthStateChangeListenerHandleTests.swift
@@ -0,0 +1,44 @@
+//
+//  AuthStateChangeListenerHandleTests.swift
+//
+//
+//  Created by Guilherme Souza on 17/02/24.
+//
+
+@testable import Auth
+import ConcurrencyExtras
+import Foundation
+import XCTest
+
+final class AuthStateChangeListenerHandleTests: XCTestCase {
+  func testRemove() {
+    let handle = AuthStateChangeListenerHandle()
+
+    let onRemoveCallCount = LockIsolated(0)
+    handle._onRemove.setValue {
+      onRemoveCallCount.withValue {
+        $0 += 1
+      }
+    }
+
+    handle.remove()
+    handle.remove()
+
+    XCTAssertEqual(onRemoveCallCount.value, 1)
+  }
+
+  func testDeinit() {
+    var handle: AuthStateChangeListenerHandle? = AuthStateChangeListenerHandle()
+
+    let onRemoveCallCount = LockIsolated(0)
+    handle?._onRemove.setValue {
+      onRemoveCallCount.withValue {
+        $0 += 1
+      }
+    }
+
+    handle = nil
+
+    XCTAssertEqual(onRemoveCallCount.value, 1)
+  }
+}