[CoreInternal] Add async flush method (#13850)

Author: ncooke3

FirebaseCore/Internal/Sources/HeartbeatLogging/HeartbeatController.swift

@@ -126,6 +126,38 @@ public final class HeartbeatController {
     }
   }
 
+  @available(iOS 13.0, macOS 10.15, macCatalyst 13.0, tvOS 13.0, watchOS 6.0, *)
+  public func flushAsync() async -> HeartbeatsPayload {
+    return await withCheckedContinuation { continuation in
+      let resetTransform = { (heartbeatsBundle: HeartbeatsBundle?) -> HeartbeatsBundle? in
+        guard let oldHeartbeatsBundle = heartbeatsBundle else {
+          return nil // Storage was empty.
+        }
+        // The new value that's stored will use the old's cache to prevent the
+        // logging of duplicates after flushing.
+        return HeartbeatsBundle(
+          capacity: self.heartbeatsStorageCapacity,
+          cache: oldHeartbeatsBundle.lastAddedHeartbeatDates
+        )
+      }
+
+      // Asynchronously gets and returns the stored heartbeats, resetting storage
+      // using the given transform.
+      storage.getAndSetAsync(using: resetTransform) { result in
+        switch result {
+        case let .success(heartbeatsBundle):
+          // If no heartbeats bundle was stored, return an empty payload.
+          continuation
+            .resume(returning: heartbeatsBundle?.makeHeartbeatsPayload() ?? HeartbeatsPayload
+              .emptyPayload)
+        case .failure:
+          // If the operation throws, assume no heartbeat(s) were retrieved or set.
+          continuation.resume(returning: HeartbeatsPayload.emptyPayload)
+        }
+      }
+    }
+  }
+
   /// Synchronously flushes the heartbeat for today.
   ///
   /// If no heartbeat was logged today, the returned payload is empty.

FirebaseCore/Internal/Sources/HeartbeatLogging/HeartbeatStorage.swift

@@ -20,6 +20,8 @@ protocol HeartbeatStorageProtocol {
   func readAndWriteAsync(using transform: @escaping (HeartbeatsBundle?) -> HeartbeatsBundle?)
   func getAndSet(using transform: (HeartbeatsBundle?) -> HeartbeatsBundle?) throws
     -> HeartbeatsBundle?
+  func getAndSetAsync(using transform: @escaping (HeartbeatsBundle?) -> HeartbeatsBundle?,
+                      completion: @escaping (Result<HeartbeatsBundle?, Error>) -> Void)
 }
 
 /// Thread-safe storage object designed for transforming heartbeat data that is persisted to disk.
@@ -134,6 +136,27 @@ final class HeartbeatStorage: HeartbeatStorageProtocol {
     return heartbeatsBundle
   }
 
+  /// Asynchronously gets the current heartbeat data from storage and resets the storage using the
+  /// given transform block.
+  /// - Parameters:
+  ///   - transform: An escaping block used to reset the currently stored heartbeat.
+  ///   - completion: An escaping block used to process the heartbeat data that
+  ///   was stored (before the `transform` was applied); otherwise, the error
+  ///   that occurred.
+  func getAndSetAsync(using transform: @escaping (HeartbeatsBundle?) -> HeartbeatsBundle?,
+                      completion: @escaping (Result<HeartbeatsBundle?, Error>) -> Void) {
+    queue.async {
+      do {
+        let oldHeartbeatsBundle = try? self.load(from: self.storage)
+        let newHeartbeatsBundle = transform(oldHeartbeatsBundle)
+        try self.save(newHeartbeatsBundle, to: self.storage)
+        completion(.success(oldHeartbeatsBundle))
+      } catch {
+        completion(.failure(error))
+      }
+    }
+  }
+
   /// Loads and decodes the stored heartbeats bundle from a given storage object.
   /// - Parameter storage: The storage container to read from.
   /// - Returns: The decoded `HeartbeatsBundle` loaded from storage; `nil` if storage is empty.

FirebaseCore/Internal/Sources/HeartbeatLogging/_ObjC_HeartbeatController.swift

@@ -45,6 +45,16 @@ public class _ObjC_HeartbeatController: NSObject {
     return _ObjC_HeartbeatsPayload(heartbeatsPayload)
   }
 
+  /// Asynchronously flushes heartbeats from storage into a heartbeats payload.
+  ///
+  /// - Note: This API is thread-safe.
+  /// - Returns: A heartbeats payload for the flushed heartbeat(s).
+  @available(iOS 13.0, macOS 10.15, macCatalyst 13.0, tvOS 13.0, watchOS 6.0, *)
+  public func flushAsync() async -> _ObjC_HeartbeatsPayload {
+    let heartbeatsPayload = await heartbeatController.flushAsync()
+    return _ObjC_HeartbeatsPayload(heartbeatsPayload)
+  }
+
   /// Synchronously flushes the heartbeat for today.
   ///
   /// If no heartbeat was logged today, the returned payload is empty.

FirebaseCore/Internal/Tests/Integration/HeartbeatLoggingIntegrationTests.swift

@@ -52,6 +52,31 @@ class HeartbeatLoggingIntegrationTests: XCTestCase {
     )
   }
 
+  @available(iOS 13.0, macOS 10.15, macCatalyst 13.0, tvOS 13.0, watchOS 6.0, *)
+  func testLogAndFlushAsync() async throws {
+    // Given
+    let heartbeatController = HeartbeatController(id: #function)
+    let expectedDate = HeartbeatsPayload.dateFormatter.string(from: Date())
+    // When
+    heartbeatController.log("dummy_agent")
+    let payload = await heartbeatController.flushAsync()
+    // Then
+    try HeartbeatLoggingTestUtils.assertEqualPayloadStrings(
+      payload.headerValue(),
+      """
+      {
+        "version": 2,
+        "heartbeats": [
+          {
+            "agent": "dummy_agent",
+            "dates": ["\(expectedDate)"]
+          }
+        ]
+      }
+      """
+    )
+  }
+
   /// This test may flake if it is executed during the transition from one day to the next.
   func testDoNotLogMoreThanOnceInACalendarDay() throws {
     // Given

FirebaseCore/Internal/Tests/Unit/HeartbeatControllerTests.swift

@@ -58,6 +58,39 @@ class HeartbeatControllerTests: XCTestCase {
     assertHeartbeatControllerFlushesEmptyPayload(controller)
   }
 
+  @available(iOS 13.0, macOS 10.15, macCatalyst 13.0, tvOS 13.0, watchOS 6.0, *)
+  func testLogAndFlushAsync() async throws {
+    // Given
+    let controller = HeartbeatController(
+      storage: HeartbeatStorageFake(),
+      dateProvider: { self.date }
+    )
+
+    assertHeartbeatControllerFlushesEmptyPayload(controller)
+
+    // When
+    controller.log("dummy_agent")
+    let heartbeatPayload = await controller.flushAsync()
+
+    // Then
+    try HeartbeatLoggingTestUtils.assertEqualPayloadStrings(
+      heartbeatPayload.headerValue(),
+      """
+      {
+        "version": 2,
+        "heartbeats": [
+          {
+            "agent": "dummy_agent",
+            "dates": ["2021-11-01"]
+          }
+        ]
+      }
+      """
+    )
+
+    assertHeartbeatControllerFlushesEmptyPayload(controller)
+  }
+
   func testLogAtEndOfTimePeriodAndAcceptAtStartOfNextOne() throws {
     // Given
     var testDate = date
@@ -404,4 +437,15 @@ private class HeartbeatStorageFake: HeartbeatStorageProtocol {
     heartbeatsBundle = transform(heartbeatsBundle)
     return oldHeartbeatsBundle
   }
+
+  func getAndSetAsync(using transform: @escaping (FirebaseCoreInternal.HeartbeatsBundle?)
+    -> FirebaseCoreInternal.HeartbeatsBundle?,
+    completion: @escaping (Result<
+      FirebaseCoreInternal.HeartbeatsBundle?,
+      any Error
+    >) -> Void) {
+    let oldHeartbeatsBundle = heartbeatsBundle
+    heartbeatsBundle = transform(heartbeatsBundle)
+    completion(.success(oldHeartbeatsBundle))
+  }
 }

FirebaseCore/Internal/Tests/Unit/HeartbeatStorageTests.swift

@@ -208,6 +208,46 @@ class HeartbeatStorageTests: XCTestCase {
     wait(for: [expectation], timeout: 0.5)
   }
 
+  func testGetAndSetAsync_ReturnsOldValueAndSetsNewValue() throws {
+    // Given
+    let heartbeatStorage = HeartbeatStorage(id: #function, storage: StorageFake())
+
+    var dummyHeartbeatsBundle = HeartbeatsBundle(capacity: 1)
+    dummyHeartbeatsBundle.append(Heartbeat(agent: "dummy_agent", date: Date()))
+
+    // When
+    let expectation1 = expectation(description: #function + "_1")
+    heartbeatStorage.getAndSetAsync { heartbeatsBundle in
+      // Assert that heartbeat storage is empty.
+      XCTAssertNil(heartbeatsBundle)
+      // Write new value.
+      return dummyHeartbeatsBundle
+    } completion: { result in
+      switch result {
+      case .success: break
+      case let .failure(error): XCTFail("Error: \(error)")
+      }
+      expectation1.fulfill()
+    }
+
+    // Then
+    let expectation2 = expectation(description: #function + "_2")
+    XCTAssertNoThrow(
+      try heartbeatStorage.getAndSet { heartbeatsBundle in
+        // Assert old value is read.
+        XCTAssertEqual(
+          heartbeatsBundle?.makeHeartbeatsPayload(),
+          dummyHeartbeatsBundle.makeHeartbeatsPayload()
+        )
+        // Write some new value.
+        expectation2.fulfill()
+        return heartbeatsBundle
+      }
+    )
+
+    wait(for: [expectation1, expectation2], timeout: 0.5, enforceOrder: true)
+  }
+
   func testGetAndSet_WhenLoadFails_PassesNilToBlockAndReturnsNil() throws {
     // Given
     let expectation = expectation(description: #function)
@@ -232,6 +272,41 @@ class HeartbeatStorageTests: XCTestCase {
     wait(for: [expectation], timeout: 0.5)
   }
 
+  func testGetAndSetAsync_WhenLoadFails_PassesNilToBlockAndReturnsNil() throws {
+    // Given
+    let readExpectation = expectation(description: #function + "_1")
+    let transformExpectation = expectation(description: #function + "_2")
+    let completionExpectation = expectation(description: #function + "_3")
+
+    let storageFake = StorageFake()
+    let heartbeatStorage = HeartbeatStorage(id: #function, storage: storageFake)
+
+    // When
+    storageFake.onRead = {
+      readExpectation.fulfill()
+      return try XCTUnwrap("BAD_DATA".data(using: .utf8))
+    }
+
+    // Then
+    heartbeatStorage.getAndSetAsync { heartbeatsBundle in
+      XCTAssertNil(heartbeatsBundle)
+      transformExpectation.fulfill()
+      return heartbeatsBundle
+    } completion: { result in
+      switch result {
+      case .success: break
+      case let .failure(error): XCTFail("Error: \(error)")
+      }
+      completionExpectation.fulfill()
+    }
+
+    wait(
+      for: [readExpectation, transformExpectation, completionExpectation],
+      timeout: 0.5,
+      enforceOrder: true
+    )
+  }
+
   func testGetAndSet_WhenSaveFails_ThrowsError() throws {
     // Given
     let expectation = expectation(description: #function)
@@ -250,7 +325,42 @@ class HeartbeatStorageTests: XCTestCase {
     wait(for: [expectation], timeout: 0.5)
   }
 
-  func testOperationsAreSynrononizedSerially() throws {
+  func testGetAndSetAsync_WhenSaveFails_ThrowsError() throws {
+    // Given
+    let transformExpectation = expectation(description: #function + "_1")
+    let writeExpectation = expectation(description: #function + "_2")
+    let completionExpectation = expectation(description: #function + "_3")
+
+    let storageFake = StorageFake()
+    let heartbeatStorage = HeartbeatStorage(id: #function, storage: storageFake)
+
+    // When
+    storageFake.onWrite = { _ in
+      writeExpectation.fulfill()
+      throw StorageError.writeError
+    }
+
+    // Then
+    heartbeatStorage.getAndSetAsync { heartbeatsBundle in
+      transformExpectation.fulfill()
+      XCTAssertNil(heartbeatsBundle)
+      return heartbeatsBundle
+    } completion: { result in
+      switch result {
+      case .success: XCTFail("Error: unexpected success")
+      case .failure: break
+      }
+      completionExpectation.fulfill()
+    }
+
+    wait(
+      for: [transformExpectation, writeExpectation, completionExpectation],
+      timeout: 0.5,
+      enforceOrder: true
+    )
+  }
+
+  func testOperationsAreSyncrononizedSerially() throws {
     // Given
     let heartbeatStorage = HeartbeatStorage(id: #function, storage: StorageFake())
 
@@ -263,10 +373,24 @@ class HeartbeatStorageTests: XCTestCase {
         return heartbeatsBundle
       }
 
-      if /* randomChoice */ .random() {
+      switch Int.random(in: 1 ... 3) {
+      case 1:
         heartbeatStorage.readAndWriteAsync(using: transform)
-      } else {
+      case 2:
         XCTAssertNoThrow(try heartbeatStorage.getAndSet(using: transform))
+      case 3:
+        let getAndSet = self.expectation(description: "GetAndSetAsync_\(i)")
+        heartbeatStorage.getAndSetAsync(using: transform) { result in
+          switch result {
+          case .success: break
+          case let .failure(error):
+            XCTFail("Unexpected: Error occurred in getAndSet_\(i), \(error)")
+          }
+          getAndSet.fulfill()
+        }
+        wait(for: [getAndSet], timeout: 1.0)
+      default:
+        XCTFail("Unexpected: Random number is out of range.")
       }
 
       return expectation