Give stream option callbacks Node's callback identities

Node passes a stream's one onwrite function to every _write and _writev
call, and a new function to every _transform call. A write or writev
option callback now receives the same closure on every call
(scr_stream_write_closure), kept or not, so a Set, WeakSet or WeakMap
finds it again. A reused transform, flush, final, construct or destroy
callback is first forgotten by weak collections, as at its free, so a
WeakSet or WeakMap that recorded the previous call's callback does not
find it on the next call.

packages/compiler/src/backend/llvm/expr-stream-callbacks.ts

@@ -684,11 +684,15 @@ export function streamCbThunkFor(
     passed.push(`ptr %sr`);
   }
   // The stream-owning completion closure (a 1-cap closure boxing the
-  // retained stream) — shared by the typed and dyn done shapes. The
-  // stream reuses it while it holds the only reference.
+  // retained stream) — shared by the typed and dyn done shapes. Every
+  // write and writev call receives the same one, as Node passes one
+  // onwrite function; other calls receive a new one, which the stream
+  // reuses while it holds the only reference.
+  const doneEntry =
+    kind === "w" || kind === "v" ? "scr_stream_write_closure" : "scr_stream_done_closure";
   const doneClosure = (fnRef: string, tag: string): string => {
-    host.declare(`declare ptr @scr_stream_done_closure(ptr, ptr)`);
-    d.push(`  %${tag}clo = call ptr @scr_stream_done_closure(ptr %s, ptr ${fnRef})`);
+    host.declare(`declare ptr @${doneEntry}(ptr, ptr)`);
+    d.push(`  %${tag}clo = call ptr @${doneEntry}(ptr %s, ptr ${fnRef})`);
     return `%${tag}clo`;
   };
   for (let i = off; i < declared.length; i++) {

packages/runtime/src/scr_stream.c

@@ -246,10 +246,14 @@ struct ScrStreamState {
   ScrClosure *destroy_cb;
   ScrStreamErrInv destroy_inv;
 
-  /* The completion closure (`callback`) the emitted option-callback
-   * thunks last passed, owned: scr_stream_done_closure hands it out again
-   * while this is its only reference. It holds the stream, so it is
-   * dropped at 'close' and traced for a stream dropped before it closes. */
+  /* The completion closures (`callback`) the emitted option-callback
+   * thunks pass, owned. write_done is the one every write and writev call
+   * receives (scr_stream_write_closure); done_cache is the last one a
+   * transform, flush, final, construct or destroy call received, handed
+   * out again while this is its only reference (scr_stream_done_closure).
+   * Both hold the stream, so they are dropped at 'close' and traced for a
+   * stream dropped before it closes. */
+  ScrClosure *write_done;
   ScrClosure *done_cache;
 
   /* finished()/pipeline() watchers: fired once at the terminal point
@@ -423,6 +427,7 @@ static void scr_stream_state_drop(ScrStreamState *st, bool gc) {
     if (st->transform_cb) scr_closure_release(st->transform_cb);
     if (st->flush_cb) scr_closure_release(st->flush_cb);
     if (st->destroy_cb) scr_closure_release(st->destroy_cb);
+    if (st->write_done) scr_closure_release(st->write_done);
     if (st->done_cache) scr_closure_release(st->done_cache);
     for (size_t i = 0; i < st->w.end_cbs_n; i++) scr_closure_release(st->w.end_cbs[i]);
     for (size_t i = 0; i < st->fin.n; i++) scr_closure_release(st->fin.cb[i]);
@@ -472,6 +477,7 @@ void scr_stream_st_trace(ScrStreamState *st, ScrTraceVisit visit, void *ctx) {
   if (st->transform_cb) visit(st->transform_cb, ctx);
   if (st->flush_cb) visit(st->flush_cb, ctx);
   if (st->destroy_cb) visit(st->destroy_cb, ctx);
+  if (st->write_done) visit(st->write_done, ctx);
   if (st->done_cache) visit(st->done_cache, ctx);
   for (size_t i = 0; i < st->w.end_cbs_n; i++) visit(st->w.end_cbs[i], ctx);
   for (size_t i = 0; i < st->fin.n; i++) visit(st->fin.cb[i], ctx);
@@ -527,24 +533,53 @@ void scr_stream_release(ScrStream *s) {
 void *scr_stream_retain_v(void *s) { return scr_stream_retain(s); }
 void scr_stream_release_v(void *s) { scr_stream_release(s); }
 
-/* The completion closure an option callback receives as `callback`: a
- * 1-cap closure calling fn over a box that holds the stream. The stream
- * keeps the last one and hands it out again while the cache holds its
- * only reference and it has no own properties; otherwise a new closure
- * replaces it, and one the user kept keeps its own stream reference.
- * Nothing is cached after 'close'. Returns +1. */
-ScrClosure *scr_stream_done_closure(ScrStream *s, void *fn) {
-  ScrStreamState *st = s->st;
-  ScrClosure *c = st->done_cache;
-  if (c != NULL && c->rc == 1 && c->fn == fn && c->props == NULL) return scr_closure_retain(c);
-  c = scr_closure_new(fn, 1);
+/* A completion closure: a 1-cap closure calling fn over a box that holds
+ * the stream. Returns +1. */
+static ScrClosure *scr_stream_done_new(ScrStream *s, void *fn) {
+  ScrClosure *c = scr_closure_new(fn, 1);
   ScrBox *bx = scr_box_new_obj(scr_stream_retain_v, scr_stream_release_v, scr_stream_trace);
   c->caps[0] = bx;
   scr_box_set_ref(bx, scr_stream_retain(s));
-  if (st->close_emitted) return c;
-  ScrClosure *old = st->done_cache;
-  st->done_cache = scr_closure_retain(c);
+  return c;
+}
+
+/* Keeps c in *slot, replacing the closure there. Nothing is kept after
+ * 'close'. Borrows c. */
+static void scr_stream_done_keep(ScrStreamState *st, ScrClosure **slot, ScrClosure *c) {
+  if (st->close_emitted) return;
+  ScrClosure *old = *slot;
+  *slot = scr_closure_retain(c);
   if (old) scr_closure_release(old);
+}
+
+/* The `callback` a write or writev option callback receives. Node passes
+ * the stream's one onwrite function to every _write and _writev call, so
+ * every call receives the same closure, kept or not. Returns +1. */
+ScrClosure *scr_stream_write_closure(ScrStream *s, void *fn) {
+  ScrStreamState *st = s->st;
+  ScrClosure *c = st->write_done;
+  if (c != NULL && c->fn == fn) return scr_closure_retain(c);
+  c = scr_stream_done_new(s, fn);
+  scr_stream_done_keep(st, &st->write_done, c);
+  return c;
+}
+
+/* The `callback` a transform, flush, final, construct or destroy option
+ * callback receives. Node passes a new function to each call. The stream
+ * keeps the last closure and hands it out again while it holds the only
+ * reference and the closure has no own properties; weak collections then
+ * forget it first, as they would at its free, so it is a new key to them.
+ * Otherwise a new closure replaces it, and one the user kept keeps its own
+ * stream reference. Returns +1. */
+ScrClosure *scr_stream_done_closure(ScrStream *s, void *fn) {
+  ScrStreamState *st = s->st;
+  ScrClosure *c = st->done_cache;
+  if (c != NULL && c->rc == 1 && c->fn == fn && c->props == NULL) {
+    scr_weak_dispose_inline(c);
+    return scr_closure_retain(c);
+  }
+  c = scr_stream_done_new(s, fn);
+  scr_stream_done_keep(st, &st->done_cache, c);
   return c;
 }
 
@@ -7989,6 +8024,7 @@ static void scr_stream_run_tick(ScrStreamTick *t) {
         if (st->transform_cb) { scr_closure_release(st->transform_cb); st->transform_cb = NULL; }
         if (st->flush_cb) { scr_closure_release(st->flush_cb); st->flush_cb = NULL; }
         if (st->destroy_cb) { scr_closure_release(st->destroy_cb); st->destroy_cb = NULL; }
+        if (st->write_done) { scr_closure_release(st->write_done); st->write_done = NULL; }
         if (st->done_cache) { scr_closure_release(st->done_cache); st->done_cache = NULL; }
       }
       break;

packages/runtime/src/scr_stream.h

@@ -681,7 +681,10 @@ void scr_stream_release(ScrStream *s);
 void *scr_stream_retain_v(void *s);
 void scr_stream_release_v(void *s);
 void scr_stream_trace(void *obj, ScrTraceVisit visit, void *ctx);
-/* The `callback` an option callback receives (+1), reused while the
+/* The `callback` a write or writev option callback receives (+1): the
+ * same closure for every call, as Node passes one onwrite function. */
+ScrClosure *scr_stream_write_closure(ScrStream *s, void *fn);
+/* The `callback` any other option callback receives (+1), reused while the
  * stream holds its only reference. */
 ScrClosure *scr_stream_done_closure(ScrStream *s, void *fn);
 void scr_stream_install(void);

tests/corpus/7517-stream-kept-callbacks.ts

@@ -3,7 +3,9 @@
 // calls its callback and keeps it (`cb(); keep(cb)`) and calls it again
 // after later writes (ERR_MULTIPLE_CALLBACK), and a Transform that calls
 // back synchronously, from setImmediate, and from a promise. Every write
-// and transform completes in Node's order.
+// and transform completes in Node's order. Every write and writev call
+// receives the same callback, as in Node (a Set, a WeakSet and a WeakMap
+// find it again); a transform call receives a new one each time.
 import { Transform, Writable } from "node:stream";
 
 type Done = (error?: Error | null) => void;
@@ -73,7 +75,65 @@ function transforms(): void {
   });
   const parts: string[] = [];
   t.on("data", (chunk: Buffer) => parts.push(chunk.toString()));
-  t.on("end", () => console.log("transform", parts.join("")));
+  t.on("end", () => {
+    console.log("transform", parts.join(""));
+    weakKeys();
+  });
   for (let i = 0; i < 9; i++) t.write(String(i));
   t.end();
 }
+
+function weakKeys(): void {
+  const seen = new WeakSet<Done>();
+  const notes = new WeakMap<Done, string>();
+  const held = new Set<Done>();
+  const w = new Writable({
+    write(chunk: Buffer, _encoding: string, callback: Done) {
+      const text = chunk.toString();
+      console.log("weak write", text, seen.has(callback), notes.get(callback) ?? "none", held.has(callback));
+      seen.add(callback);
+      notes.set(callback, text);
+      if (text === "q") held.add(callback);
+      callback();
+    },
+  });
+  for (const text of ["p", "q", "r", "s"]) w.write(text);
+  w.end(batches);
+}
+
+function batches(): void {
+  const seen = new WeakSet<Done>();
+  let first: Done | null = null;
+  const w = new Writable({
+    writev(chunks: Array<{ chunk: Buffer }>, callback: Done) {
+      const text = chunks.map((c) => c.chunk.toString()).join("");
+      console.log("writev", text, seen.has(callback), first === callback);
+      seen.add(callback);
+      if (first === null) first = callback;
+      setImmediate(() => callback());
+    },
+  });
+  w.write("a");
+  w.write("b");
+  w.write("c");
+  w.end(weakTransform);
+}
+
+function weakTransform(): void {
+  type TransformDone = (error?: Error | null, data?: Buffer) => void;
+  const seen = new WeakSet<TransformDone>();
+  const held = new Set<TransformDone>();
+  const t = new Transform({
+    transform(chunk: Buffer, _encoding: string, callback: TransformDone) {
+      const text = chunk.toString();
+      console.log("weak transform", text, seen.has(callback), held.has(callback));
+      seen.add(callback);
+      if (text === "x") held.add(callback);
+      callback(null, chunk);
+    },
+  });
+  t.resume();
+  t.on("end", () => console.log("weak transform end"));
+  for (const text of ["x", "y", "z"]) t.write(text);
+  t.end();
+}