1 files changed, 69 insertions, 11 deletions
diff --git a/src/9p_io.zig b/src/9p_io.zig
index d6eef9d4..2ca11a87 100644
--- a/src/9p_io.zig
+++ b/src/9p_io.zig
@@ -174,6 +174,16 @@ pub const Listener = struct {
     core: *pardes.Pardes,
     runner: Runner = undefined,
     stopping: std.atomic.Value(bool) = .init(false),
+    /// A connection wrote the Restore and has yet to answer it.
+    restore_writer: std.atomic.Value(bool) = .init(false),
+    /// Connections accepted so far, and the count each slot's connection
+    /// was accepted at: a Restore cuts those accepted before it (`reset`),
+    /// and a new client in a freed slot has a later stamp. Written on the
+    /// connection's task as it opens, read with the turn.
+    accepts: std.atomic.Value(u64) = .init(0),
+    accepted: [max_conns]std.atomic.Value(u64) = @splat(.init(0)),
+    /// Connections accepted before this count are being cut by a Restore.
+    cut_before: u64 = 0,
     tcp_address: ?std.Io.net.IpAddress = null,
     quic: if (quic_enabled) ?quic.Listener else void = if (quic_enabled) null else {},
     quic_address: ?std.Io.net.IpAddress = null,
@@ -207,6 +217,16 @@ pub const Listener = struct {
     /// answers. A request that would change a pane while the editor is out
     /// in a syscall mid-step is parked in the engine instead, and retried
     /// when the turn is next given up between steps (`wakeParked`).
+    fn onOpened(ctx: ?*anyopaque, conn: *Runner.Conn) void {
+        const l = of(ctx);
+        l.accepted[conn.index].store(l.accepts.fetchAdd(1, .acq_rel) + 1, .release);
+    }
+
+    /// A live connection accepted before the Restore under way.
+    fn cutting(l: *Listener, conn: *Runner.Conn) bool {
+        return conn.live() and l.accepted[conn.index].load(.acquire) <= l.cut_before and l.cut_before != 0;
+    }
+
     fn onServe(ctx: ?*anyopaque, conn: *Runner.Conn, req: pardes.ctlfs.Req) void {
         const l = of(ctx);
         // Once `stop` has begun the editor is tearing down and may never
@@ -217,6 +237,12 @@ pub const Listener = struct {
         }
         const quiet = pardes.turn.take();
         defer pardes.turn.give();
+        // A connection a Restore is cutting (`reset`) is the old editor's
+        // client: its fids name the old editor's panes and opens.
+        if (l.cutting(conn)) {
+            const refused = pardes.ctlfs.Reply.fail(req.tag, pardes.ctlfs.E.IO);
+            return conn.reply(&refused, "");
+        }
         if (!quiet and pardes.ctlfs.needsQuiet(req)) {
             pardes.turn.parked = true;
             const later: pardes.ctlfs.Reply = .{ .tag = req.tag, .status = .again };
@@ -267,13 +293,20 @@ pub const Listener = struct {
         // and `deinit` lets the answer out before it cuts the connections.
         if (core.quit) return conn.reply(&reply, "");
         const restoring = core.restore_req != null;
+        // Set with the turn still held, so `reset` sees it and waits.
+        if (restoring) l.restore_writer.store(true, .release);
+        defer if (restoring) l.restore_writer.store(false, .release);
         if (core.effects_len != 0) pardes.turn.awaitSettled(epoch);
+        // The Restore's own write is answered once it is done, and before
+        // `reset` hangs every connection up, this one too, so the writer
+        // hears that it happened rather than a cut; a failed one leaves the
+        // connection and says why on the message row and in the log.
+        if (restoring) {
+            if (l.core == core) pardes.turn.awaitRestored(restores);
+            return conn.reply(&reply, "");
+        }
         // Once a wait returns `core` may be gone: a Restore meanwhile put a
-        // replacement in (`reset`) and is hanging this connection up, which
-        // is a Restore's answer. Only one that failed leaves this connection
-        // here to answer.
-        if (l.core != core) return;
-        if (restoring) pardes.turn.awaitRestored(restores);
+        // replacement in and is hanging this connection up.
         if (l.core != core) return;
         conn.reply(&reply, "");
     }
@@ -484,9 +517,9 @@ pub const Listener = struct {
     /// Hangs every connection up, so that `replacement` starts with no
     /// client holding anything. From here requests are answered from the
     /// replacement -- a client that connects meanwhile is a client of the
-    /// replacement -- and every task waiting on the old core is let go: it
-    /// finds the core changed and answers nothing, its connection being
-    /// hung up. The connections pay their releases on their own tasks, so
+    /// replacement -- and every task waiting on the old core is let go: the
+    /// Restore's own writer answers, the rest find the core changed and
+    /// answer nothing, their connections being hung up. The connections pay their releases on their own tasks, so
     /// the editor rests while they do.
     pub fn reset(l: *Listener, replacement: *pardes.Pardes) void {
         for (0..quic_slots) |i| l.drop(@intCast(i));
@@ -496,13 +529,38 @@ pub const Listener = struct {
         replacement.fs.tcp_address = l.tcp_address;
         replacement.fs.quic_address = l.quic_address;
         if (comptime supported) {
-            l.runner.closeAll();
+            // The Restore's writer answers as it wakes (`onServe`), and its
+            // answer gets 200 ms to leave before the cut; every other request
+            // from the old editor's clients meanwhile is refused, since the
+            // fids it names are the old editor's. A client that connects now
+            // is the replacement's, and stays.
+            l.cut_before = l.accepts.load(.acquire);
             pardes.turn.settle();
             pardes.turn.restoreSettled();
             pardes.turn.rest();
+            const flush_by = nowMs() +| 200;
+            while (nowMs() < flush_by) {
+                const pending = l.restore_writer.load(.acquire) or for (&l.runner.conns) |*conn| {
+                    if (!l.cutting(conn)) continue;
+                    conn.lock();
+                    const n = conn.engine.output().len;
+                    conn.unlock();
+                    if (n != 0) break true;
+                } else false;
+                if (!pending) break;
+                Client.nap(1);
+            }
+            for (&l.runner.conns) |*conn| if (l.cutting(conn)) conn.close();
             const deadline = nowMs() +| 2000;
-            while (l.runner.count() != 0 and nowMs() < deadline) Client.nap(1);
+            while (nowMs() < deadline) {
+                const left = for (&l.runner.conns) |*conn| {
+                    if (l.cutting(conn)) break true;
+                } else false;
+                if (!left) break;
+                Client.nap(1);
+            }
             pardes.turn.wake();
+            l.cut_before = 0;
         }
         l.collectOs();
         for (&l.conns) |*conn| conn.accepted_ms = 0;
@@ -758,7 +816,7 @@ pub fn listen(io: std.Io, gpa: std.mem.Allocator, core: *pardes.Pardes, named: [
     l.runner.init(.{
         .io = io,
         .root = pardes.ctlfs.root,
-        .handler = .{ .ctx = l, .serve = Listener.onServe },
+        .handler = .{ .ctx = l, .serve = Listener.onServe, .opened = Listener.onOpened },
         .greet_timeout_ms = Listener.greet_deadline_ms,
     });
     const entry_name = if (named.len != 0) named else fallback;