Implement continuation queue in AsyncReadWriteLock (apollographql/apollo-ios-dev#851)

Author: AnthonyMDev

Sources/Apollo/Caching/RecordSet.swift

@@ -55,17 +55,19 @@ public struct RecordSet: Sendable, Hashable {
   }
 
   @discardableResult public mutating func merge(record: Record) -> Set<CacheKey> {
-    if var oldRecord = storage.removeValue(forKey: record.key) {
+    if let oldRecord = storage[record.key] {
       var changedKeys: Set<CacheKey> = Set()
-
+      var updatedRecord = oldRecord
+      
       for (key, value) in record.fields {
         if let oldValue = oldRecord.fields[key], AnyHashable(oldValue) == AnyHashable(value) {
           continue
         }
-        oldRecord[key] = value
+        updatedRecord[key] = value
         changedKeys.insert([record.key, key].joined(separator: "."))
       }
-      storage[record.key] = oldRecord
+
+      storage[record.key] = updatedRecord
       return changedKeys
     } else {
       storage[record.key] = record

Sources/Apollo/Internal Utilities/AsyncReadWriteLock.swift

@@ -1,6 +1,12 @@
 import Foundation
 
 actor AsyncReadWriteLock {
+  private enum State {
+    case idle
+    case reading
+    case writing
+  }
+
   private final class ReadTask: Sendable {
     let task: Task<Void, any Swift.Error>
 
@@ -13,19 +19,32 @@ actor AsyncReadWriteLock {
 
   private var currentReadTasks: [ObjectIdentifier: ReadTask] = [:]
   private var currentWriteTask: Task<Void, any Swift.Error>?
+  private var queue: [(handle: CheckedContinuation<Void, Never>, isWriter: Bool)] = []
+
+  private var state: State {
+    if currentWriteTask != nil { return .writing }
+    if !currentReadTasks.isEmpty { return .reading }
+    return .idle
+  }
 
   /// Waits for all current reads/writes to be completed, then calls the provided closure while preventing
   /// any other reads/writes from beginning.
   ///
   /// This function should be `rethrows` but the compiler doesn't understand that when passing the `block` into a Task.
   /// If the `body` provided does not throw, this function will not throw.
   func write(_ body: @Sendable @escaping () async throws -> Void) async throws {
-    while currentWriteTask != nil || !currentReadTasks.isEmpty {
-      await Task.yield()
-      continue
+    switch state {
+    case .writing, .reading:
+      await addToQueueAndWait(isWriter: true)
+
+    case .idle:
+      break
     }
 
-    defer { currentWriteTask = nil }
+    defer {
+      currentWriteTask = nil
+      writeTaskDidFinish()
+    }
     let writeTask = Task {
       try await body()
     }
@@ -40,17 +59,68 @@ actor AsyncReadWriteLock {
   /// This function should be `rethrows` but the compiler doesn't understand that when passing the `block` into a Task.
   /// If the `body` provided does not throw, this function will not throw.
   func read(_ body: @Sendable @escaping () async throws -> Void) async throws {
-    while currentWriteTask != nil {
-      await Task.yield()
-      continue
+    switch state {
+    case .writing:
+      await addToQueueAndWait(isWriter: false)
+
+    case .reading:
+      // If we are currently reading, there will only be a queue if there is at least one waiting writer. If a writer
+      // is waiting, we should queue the new reader for after the write.
+      if !queue.isEmpty {
+        await addToQueueAndWait(isWriter: false)
+      }
+    case .idle:
+      break
     }
 
     let readTask = ReadTask(body)
     let taskID = ObjectIdentifier(readTask)
-    defer { currentReadTasks[taskID] = nil }
+    defer {
+      currentReadTasks[taskID] = nil
+      readTaskDidFinish()
+    }
     currentReadTasks[taskID] = readTask
 
     try await readTask.task.value
   }
 
+  private func addToQueueAndWait(isWriter: Bool) async {
+    await withCheckedContinuation { continuation in
+      queue.append((handle: continuation, isWriter: isWriter))
+    }
+  }
+
+  private func readTaskDidFinish() {
+    if state == .idle {
+      wakeNext()
+    }
+  }
+
+  private func writeTaskDidFinish() {
+    wakeNext()
+  }
+
+  private func wakeNext() {
+    guard !queue.isEmpty else {
+      return
+    }
+
+    let next = queue[0]
+    next.handle.resume(returning: ())
+
+    if next.isWriter {
+      queue.remove(at: 0)
+      return
+
+    } else {
+      var lastReader = 0
+      for i in 1..<queue.count {
+        guard !queue[i].isWriter else { break }
+        queue[i].handle.resume(returning: ())
+        lastReader = i
+      }
+      queue.removeSubrange(0...lastReader)
+    }
+  }
+
 }