Bound the drain's send, and stage what times out

GraphcodeKit/Sources/GraphStore.swift

@@ -176,6 +176,7 @@ public actor GraphStore {
   /// A deferred delivery: a `--follow-up` message, or a Mailroom watcher's wake for one
   /// post (`watchedPostID`), waiting for its target to go idle.
   struct PendingFollowUp: Equatable {
+    let id: UUID
     let nodeID: UUID
     let text: String
     /// The post a watcher's wake is about — what lets `mail watch --off` drop the wakes
@@ -189,6 +190,8 @@ public actor GraphStore {
   }
 
   private var pendingFollowUps: [PendingFollowUp] = []
+  private var pendingDeliveryAttempts: Set<UUID> = []
+  private var completedTimedOutDeliveries: [UUID: Bool] = [:]
   /// `drainPendingFollowUps` runs across several awaits, and the presence poll that
   /// calls it runs every fifteen seconds: without this, a second drain started while
   /// the first was suspended delivered the same items again and, on finishing, wrote
@@ -233,6 +236,8 @@ public actor GraphStore {
   /// one poll, not the process.
   private let presenceReadDeadline: Duration
 
+  private let deliveryDeadline: Duration
+
   /// How long a drain may hold the queue before another is allowed to take over. Far
   /// longer than any healthy pass (milliseconds) and longer than a pass that meets
   /// several wedged presence reads, each of which is bounded by `presenceReadDeadline`;
@@ -282,6 +287,7 @@ public actor GraphStore {
   /// `ZmxSessionLauncher`; tests leave it `nil` or capture the calls.
   public init(
     graph: LoopGraph = LoopGraph(project: ProjectRef(path: "", name: "Untitled")),
+    deliveryDeadline: Duration = .seconds(45),
     onGraphChanged: (@Sendable (LoopGraph) -> Void)? = nil,
     onEnsureSession: (@Sendable (LoopNode, String?) -> Void)? = nil,
     onTerminateSession: (@Sendable (LoopNode, String?) -> Void)? = nil,
@@ -345,6 +351,7 @@ public actor GraphStore {
     self.goalCache = goalCache ?? GoalEvaluationCache()
     self.recurrence = recurrence
     self.presenceReadDeadline = presenceReadDeadline
+    self.deliveryDeadline = deliveryDeadline
     self.drainLeaseDuration = drainLeaseDuration
   }
 
@@ -2733,7 +2740,8 @@ public actor GraphStore {
     // dropping it.
     if followUp, deliversLater(to: target) {
       pendingFollowUps.append(
-        PendingFollowUp(nodeID: nodeID, text: message, watchedPostID: watchedPostID))
+        PendingFollowUp(
+          id: UUID(), nodeID: nodeID, text: message, watchedPostID: watchedPostID))
       return
     }
 
@@ -2873,7 +2881,15 @@ public actor GraphStore {
     // the queue rather than in a variable about to go out of scope, whatever the exit.
     defer {
       if drainOwner == owner {
-        pendingFollowUps = remaining + Array(batch[index...]) + pendingFollowUps
+        var retained: [PendingFollowUp] = []
+        for pending in remaining + Array(batch[index...]) {
+          if let confirmed = completedTimedOutDeliveries.removeValue(forKey: pending.id) {
+            if !confirmed { _ = staged(pending) }
+          } else {
+            retained.append(pending)
+          }
+        }
+        pendingFollowUps = retained + pendingFollowUps
       }
     }
     // One reading per target per pass, taken the first time this pass reaches that
@@ -2892,6 +2908,12 @@ public actor GraphStore {
       index += 1
       drainBatch = Array(batch[index...])
       drainInFlight = pending
+      if pendingDeliveryAttempts.contains(pending.id) {
+        drainInFlight = nil
+        remaining.append(pending)
+        drainDeferred = remaining
+        continue
+      }
       guard let node = graph.nodes[id: pending.nodeID], !node.isResolved else {
         // Its work is over, but the message must not go with the queue entry: the log is
         // what the loop's next wake reads, and for a resolved loop that is all there is.
@@ -2946,9 +2968,38 @@ public actor GraphStore {
         drainInFlight = nil
         continue
       }
-      let delivered = await deliverToSession(node, pending.text)
+      pendingDeliveryAttempts.insert(pending.id)
+      let attempt = DeliveryAttempt()
+      let delivered = await withDeadline(deliveryDeadline) {
+        let result = await self.deliverToSession(node, pending.text)
+        await attempt.complete(result)
+        return result
+      }
       guard drainOwner == owner else { return }
+      guard let delivered else {
+        remaining.append(pending)
+        drainDeferred = remaining
+        Task { [self] in
+          let result = await attempt.wait()
+          finishTimedOutDelivery(pending.id, confirmed: result)
+        }
+        DaemonLog.shared.record(
+          "delivery-stall",
+          DaemonRequestContext.fields + [
+            ("node", node.id.uuidString),
+            ("deadline_ms", DaemonLog.milliseconds(deliveryDeadline.timeInterval)),
+          ])
+        drainInFlight = nil
+        continue
+      }
+      pendingDeliveryAttempts.remove(pending.id)
       if !delivered {
+        DaemonLog.shared.record(
+          "delivery-stall",
+          DaemonRequestContext.fields + [
+            ("node", node.id.uuidString),
+            ("deadline_ms", DaemonLog.milliseconds(deliveryDeadline.timeInterval)),
+          ])
         remaining.append(staged(pending))
         drainDeferred = remaining
         if !pending.recorded {
@@ -2960,6 +3011,16 @@ public actor GraphStore {
     }
   }
 
+  private func finishTimedOutDelivery(_ id: UUID, confirmed: Bool) {
+    guard pendingDeliveryAttempts.remove(id) != nil else { return }
+    if let index = pendingFollowUps.firstIndex(where: { $0.id == id }) {
+      let pending = pendingFollowUps.remove(at: index)
+      if !confirmed { _ = staged(pending) }
+    } else {
+      completedTimedOutDeliveries[id] = confirmed
+    }
+  }
+
   private func announceError(_ message: String) {
     if let frame = Self.encode(.errorOccurred(message)) {
       for id in connections.keys {
@@ -3782,3 +3843,21 @@ public actor GraphStore {
     }
   }
 }
+
+private actor DeliveryAttempt {
+  private var result: Bool?
+  private var waiter: CheckedContinuation<Bool, Never>?
+
+  func complete(_ result: Bool) {
+    guard self.result == nil else { return }
+    self.result = result
+    guard let waiter else { return }
+    self.waiter = nil
+    waiter.resume(returning: result)
+  }
+
+  func wait() async -> Bool {
+    if let result { return result }
+    return await withCheckedContinuation { waiter = $0 }
+  }
+}

graphcode/Tests/DrainWedgeVerificationTests.swift

@@ -0,0 +1,232 @@
+import ComposableArchitecture
+import Foundation
+import GraphcodeKit
+import Testing
+
+/// The stable-release check's wedge tests, inverted against #316.
+///
+/// On `main` at 0.1.64-beta5 each of these described a freeze; here they describe the
+/// fix. The setup is byte-for-byte what was measured on beta5 — one loop whose presence
+/// read never returns, one innocent bystander — so a regression restores the original
+/// failure rather than a new one.
+@Suite
+struct DrainWedgeVerificationTests {
+  private struct Fixture {
+    let graph: LoopGraph
+    var hung: UUID { graph.nodes[0].id }
+    var bystander: UUID { graph.nodes[1].id }
+  }
+
+  private func fixture() -> Fixture {
+    Fixture(
+      graph: LoopGraph(
+        project: ProjectRef(path: "/tmp/drainwedge", name: "drainwedge"),
+        nodes: [
+          LoopNode(
+            title: "Hung", loopType: .goalBased, goal: GoalSpec(summary: "hangs"),
+            presence: PresenceReading(presence: .busy, confidence: .reported),
+            state: .running),
+          LoopNode(
+            title: "Bystander", loopType: .goalBased, goal: GoalSpec(summary: "innocent"),
+            presence: PresenceReading(presence: .busy, confidence: .reported),
+            state: .running),
+        ]))
+  }
+
+  /// Beta5: never delivered, for as long as a client stayed attached.
+  /// Here: delivered, once the bounded read times out.
+  @Test
+  func aHungPresenceReadNoLongerFreezesTheBystandersFollowUp() async {
+    let fixture = fixture()
+    let delivered = LockIsolated<[String]>([])
+    let release = LockIsolated(false)
+    let store = GraphStore(
+      graph: fixture.graph,
+      onDeliverMessage: { _, message, _ in
+        delivered.withValue { $0.append(message) }
+        return true
+      },
+      onReadPresence: { node, _ in
+        guard node.id == fixture.hung else {
+          return PresenceReading(presence: .idle, confidence: .reported)
+        }
+        while !release.value { try? await Task.sleep(for: .milliseconds(20)) }
+        return PresenceReading(presence: .busy, confidence: .reported)
+      },
+      presenceReadDeadline: .milliseconds(50))
+
+    await store.handle(
+      .messageNode(fixture.hung, text: "for the hung one", from: nil, followUp: true))
+    await store.handle(
+      .messageNode(fixture.bystander, text: "for the bystander", from: nil, followUp: true))
+
+    #expect(delivered.value == ["[graphcode] for the bystander"])
+    release.setValue(true)
+  }
+
+  /// The hung loop's own message is owed, not dropped: a timed-out read is `.unknown`,
+  /// which is not a state, so the item stays queued for a pass that gets an answer.
+  @Test
+  func theHungLoopsOwnMessageIsStillOwedAndLandsWhenItAnswers() async {
+    let fixture = fixture()
+    let delivered = LockIsolated<[String]>([])
+    let release = LockIsolated(false)
+    let store = GraphStore(
+      graph: fixture.graph,
+      onDeliverMessage: { _, message, _ in
+        delivered.withValue { $0.append(message) }
+        return true
+      },
+      onReadPresence: { node, _ in
+        guard node.id == fixture.hung else {
+          return PresenceReading(presence: .busy, confidence: .reported)
+        }
+        while !release.value { try? await Task.sleep(for: .milliseconds(20)) }
+        return PresenceReading(presence: .idle, confidence: .reported)
+      },
+      presenceReadDeadline: .milliseconds(50))
+
+    await store.handle(
+      .messageNode(fixture.hung, text: "for the hung one", from: nil, followUp: true))
+    #expect(delivered.value.isEmpty)
+
+    release.setValue(true)
+    await store.handle(.refreshUsage)
+
+    #expect(delivered.value == ["[graphcode] for the hung one"])
+  }
+
+  /// Point 3 of the stable check: on beta5 `refreshPresence` walked the graph serially
+  /// through the same unbounded await, so nothing behind the hung node was ever read.
+  @Test
+  func thePollTickReachesLoopsBehindAHungOne() async {
+    let fixture = fixture()
+    let asked = LockIsolated<[UUID]>([])
+    let release = LockIsolated(false)
+    let store = GraphStore(
+      graph: fixture.graph,
+      onReadPresence: { node, _ in
+        asked.withValue { $0.append(node.id) }
+        guard node.id == fixture.hung else {
+          return PresenceReading(presence: .idle, confidence: .reported)
+        }
+        while !release.value { try? await Task.sleep(for: .milliseconds(20)) }
+        return PresenceReading(presence: .busy, confidence: .reported)
+      },
+      presenceReadDeadline: .milliseconds(50))
+
+    await store.handle(.refreshUsage)
+
+    #expect(asked.value.contains(fixture.bystander))
+    release.setValue(true)
+  }
+}
+
+/// The presence read is bounded now. The **delivery** in the same loop is not.
+///
+/// `deliverToSession` → `onDeliverMessage` → `CLISessionBackend.deliverMessage` →
+/// `ZmxSessionLauncher.send` → `PTYProcessSession.waitCollectingOutput` is the same
+/// unbounded chain the presence read had, and `GraphcodeKit/Sources/Sessions/` is
+/// untouched by #316. So `drainPendingFollowUps` still holds `isDrainingFollowUps`
+/// across an `await` that can never return — which is the premise the "no lease needed"
+/// argument rests on.
+///
+/// A send is if anything the likelier of the two to wedge: a message is chunked into
+/// several sequential `zmx send` invocations, so one message is several chances to hang
+/// on a wedged `ControlMaster` rather than one.
+@Suite
+struct DeliveryWedgeTests {
+  private struct Fixture {
+    let graph: LoopGraph
+    var hung: UUID { graph.nodes[0].id }
+    var bystander: UUID { graph.nodes[1].id }
+  }
+
+  private func fixture() -> Fixture {
+    Fixture(
+      graph: LoopGraph(
+        project: ProjectRef(path: "/tmp/deliverywedge", name: "deliverywedge"),
+        nodes: [
+          LoopNode(
+            title: "Hung", loopType: .goalBased, goal: GoalSpec(summary: "hangs"),
+            presence: PresenceReading(presence: .busy, confidence: .reported),
+            state: .running),
+          LoopNode(
+            title: "Bystander", loopType: .goalBased, goal: GoalSpec(summary: "innocent"),
+            presence: PresenceReading(presence: .busy, confidence: .reported),
+            state: .running),
+        ]))
+  }
+
+  private func settle(until flag: LockIsolated<Bool>) async {
+    for _ in 0..<2000 where !flag.value { try? await Task.sleep(for: .milliseconds(5)) }
+  }
+
+  /// The control: every send returns, so the bystander is served.
+  @Test
+  func aBystandersFollowUpIsDeliveredWhenEverySendReturns() async {
+    let fixture = fixture()
+    let delivered = LockIsolated<[String]>([])
+    let store = GraphStore(
+      graph: fixture.graph,
+      onDeliverMessage: { _, message, _ in
+        delivered.withValue { $0.append(message) }
+        return true
+      },
+      onReadPresence: { _, _ in PresenceReading(presence: .idle, confidence: .reported) })
+
+    await store.handle(
+      .messageNode(fixture.hung, text: "for the hung one", from: nil, followUp: true))
+    await store.handle(
+      .messageNode(fixture.bystander, text: "for the bystander", from: nil, followUp: true))
+
+    #expect(delivered.value.contains("[graphcode] for the bystander"))
+  }
+
+  /// The same wedge as beta5, reached through the send instead of the read: one loop's
+  /// `zmx send` never returns and every other loop's staged mail stops moving, silently,
+  /// for as long as the process lives.
+  @Test
+  func aHungDeliveryStillFreezesEveryOtherLoopsFollowUps() async {
+    let fixture = fixture()
+    let delivered = LockIsolated<[String]>([])
+    let entered = LockIsolated(false)
+    let release = LockIsolated(false)
+    let store = GraphStore(
+      graph: fixture.graph,
+      onDeliverMessage: { node, message, _ in
+        guard node.id == fixture.hung else {
+          delivered.withValue { $0.append(message) }
+          return true
+        }
+        entered.setValue(true)
+        while !release.value { try? await Task.sleep(for: .milliseconds(5)) }
+        return true
+      },
+      // Both loops read idle, so the drain tries to deliver to both.
+      onReadPresence: { _, _ in PresenceReading(presence: .idle, confidence: .reported) },
+      deliveryDeadline: .milliseconds(50))
+
+    // Queued while both are mid-turn (cached `busy`), then delivered by the drain.
+    let wedging = Task {
+      await store.handle(
+        .messageNode(fixture.hung, text: "for the hung one", from: nil, followUp: true))
+    }
+    await settle(until: entered)
+    #expect(entered.value)
+
+    // The store keeps answering, and the bystander's mail moves after the timeout.
+    await store.handle(
+      .messageNode(fixture.bystander, text: "for the bystander", from: nil, followUp: true))
+    try? await Task.sleep(for: .milliseconds(100))
+    await store.handle(.refreshUsage)
+
+    #expect(delivered.value == ["[graphcode] for the bystander"])
+    #expect(!wedging.isCancelled)
+
+    release.setValue(true)
+    _ = await wedging.value
+    try? await Task.sleep(for: .milliseconds(50))
+    #expect(delivered.value == ["[graphcode] for the bystander", "[graphcode] for the hung one"])
+  }
+}