Make the iof's XON/XOFF actually stop the producer

config/prte_setup_pmix.m4

@@ -309,6 +309,24 @@ AC_DEFUN([PRTE_CHECK_PMIX],[
                        [$prte_pmix_iof_deliver_local],
                        [Whether PMIx honors PMIX_IOF_LOCAL_OUTPUT on an IOF delivery])
 
+    dnl The iof has always had the shape of XON/XOFF back-pressure on stdin -
+    dnl a daemon whose stdin sink passes PRTE_IOF_MAX_INPUT_BUFFERS tells the
+    dnl HNP so - but the HNP could not act on it.  Stdin originates in a PMIx
+    dnl server's read of its own stdin or in a tool's PMIx_IOF_push, and until
+    dnl PMIx grew PMIx_server_IOF_flow_control there was no way to stop either
+    dnl at the source: a refusal anywhere would have dropped bytes rather than
+    dnl slowed the producer.  Without this the HNP consumes the message and
+    dnl the daemon's sink simply queues, bounded only by iof_base_output_limit.
+    AC_MSG_CHECKING([for PMIx IOF stdin flow control])
+    PRTE_CHECK_PMIX_CAP([IOF_FLOW_CONTROL],
+                        [AC_MSG_RESULT([yes])
+                         prte_pmix_iof_flow_control=1],
+                        [AC_MSG_RESULT([no])
+                         prte_pmix_iof_flow_control=0])
+    AC_DEFINE_UNQUOTED([PRTE_PMIX_IOF_FLOW_CONTROL],
+                       [$prte_pmix_iof_flow_control],
+                       [Whether PMIx supports PMIx_server_IOF_flow_control])
+
     AC_MSG_CHECKING([for LTO compatibility])
     PRTE_CHECK_PMIX_CAP([LTO],
                         [PRTE_PMIX_LTO_CAPABILITY=1

contrib/dockerswarm/AGENTS.md

@@ -418,6 +418,18 @@ unpack error. Do not weaken the reader's pace: at `cat` speed, or even at
 the pace the two cases above use, the backlog may never cross 50 and the
 case passes without having tested anything.
 
+A fourth pass runs the same reader with `--prtemca iof_base_verbose 1` and
+**counts** the two halves of the flow control in the HNP's log. That is
+the case that covers `PMIx_server_IOF_flow_control` actually being
+reached, and it is a count rather than a grep on purpose: PMIx has no
+status meaning "resume", so an implementation that asserts XOFF fifty
+times and releases once looks identical to a correct one under a presence
+test — and the job it produces is *hung*, not slow. The case fails if the
+counts differ, if flow control never engaged at all (which would mean the
+reader was too fast and the case proved nothing), or if a byte went
+missing across the suspensions. See the framework guide's *Flow control*
+section for where each half lives.
+
 **Grow** (`elastic grow node2:2,node3:2`): phase-1 `PMIX_SUCCESS`, then phase-2
 `PMIX_DVM_IS_READY`, and `prted` now running on node2 and node3.
 

contrib/dockerswarm/run-tests.sh

@@ -5007,6 +5007,43 @@ gcc -o /root/staged_marker /root/staged_marker.c' >/dev/null 2>&1
             || ok "no unpack error at the HNP while stdin was backed up"
         ON 2 'rm -f /tmp/iof_xoff_out.txt' >/dev/null 2>&1
 
+        # The same reader again, this time with the iof talking, to assert
+        # that flow control actually engaged AND that every XOFF was paired
+        # with an XON.
+        #
+        # The pairing is the part that matters. PMIx has no status meaning
+        # "resume", so a suspension the HNP asserts and forgets to release
+        # stalls the producer for the life of the job - a hang, not a
+        # slowdown. On the HNP side that release lives in
+        # release_flow_control(), called from every path a backed-up sink
+        # can leave that state by, including the ones where the sink is torn
+        # down rather than drained.
+        #
+        # Counting rather than merely grepping is deliberate: an
+        # implementation that asserts XOFF once and releases once looks
+        # identical to a correct one under a presence test, and so does one
+        # that asserts fifty times and releases once.
+        RUN 'cp /tmp/prte.out /tmp/prte.out.mark2 2>/dev/null || : ' >/dev/null 2>&1
+        out=$(RUN 'cd /tmp && timeout 300 prun --prtemca iof_base_verbose 1 -n 1 \
+                     '"$SC"' /tmp/iof_pair_out.txt 2048 3000 < iof_stdin_in.txt \
+                     2>iof_pair_err.txt; echo "rc=$?"')
+        rc=$(echo "$out" | sed -n 's/^rc=//p')
+        got=$(echo "$out" | sed -n 's/^SLOWCAT-BYTES //p')
+        pairlog=$(RUN 'diff /tmp/prte.out.mark2 /tmp/prte.out 2>/dev/null | grep "^>" \
+                       || cat /tmp/prte.out 2>/dev/null')
+        nxoff=$(echo "$pairlog" | grep -c 'buffer backed up - holding')
+        nxon=$(echo "$pairlog" | grep -c 'releasing stdin flow control')
+        [ "$nxoff" -gt 0 ] \
+            && ok "stdin flow control engaged ($nxoff XOFF)" \
+            || bad "stdin flow control never engaged - the reader was not slow enough, or the backlog never crossed PRTE_IOF_MAX_INPUT_BUFFERS"
+        [ "$nxoff" = "$nxon" ] \
+            && ok "every XOFF was paired with an XON ($nxon)" \
+            || bad "unpaired stdin flow control: $nxoff XOFF vs $nxon XON - a producer is left suspended"
+        [ "$rc" = 0 ] && [ "$got" = "$insz" ] \
+            && ok "delivery stayed exact across $nxoff suspensions" \
+            || bad "flow-controlled stdin lost data (rc=$rc, sent=$insz received=$got)"
+        RUN 'rm -f /tmp/iof_pair_out.txt' >/dev/null 2>&1
+
         # And the same slow reader on the HNP node, where push_stdin writes
         # into the proc sink directly instead of going out over the RML: the
         # HNP and the daemon carry separate copies of the write handler, so a

src/mca/iof/AGENTS.md

@@ -419,56 +419,98 @@ tree or an old commit is **retired**.
 
 ---
 
-## Flow control — read this before you trust it
-
-stdin can outrun a slow reader, and the framework has the *shape* of
-XON/XOFF back-pressure keyed on `PRTE_IOF_MAX_INPUT_BUFFERS` (50 queued
-chunks). **It does not actually throttle anything today.** Know that
-before you reason about a stdin backlog, and do not "fix" a symptom by
-leaning on it.
-
-What is really wired:
-
-- On a daemon, when `prte_iof_base_write_output` reports the stdin sink
-  backlog has crossed 50 (or a write errors out),
-  `prte_iof_prted_send_xonxoff(PRTE_IOF_XOFF)` sends the HNP a buffer
-  holding nothing but the tag, latched by
-  `prte_mca_iof_prted_component.xoff`; when the backlog drains below 50 it
-  sends `PRTE_IOF_XON`.
-- The HNP **recognizes** those messages (`prte_iof_hnp_recv` screens the
-  `PRTE_IOF_XON | PRTE_IOF_XOFF` mask ahead of everything else) and traces
-  them, but takes no action.
-- On the HNP, `push_stdin` to a *local* proc returns
-  `PRTE_ERR_OUT_OF_RESOURCE` when its own sink passes the same threshold.
-
-Why nothing happens, and why that is currently correct: stdin originates
-in the PMIx server's read of its own stdin, or in a tool's
-`PMIx_IOF_push`. Neither end honors a refusal. PRRTE's own glue
-(`pmix_server_stdin_push` in
-[`src/prted/pmix/pmix_server_gen.c`](../../prted/pmix/pmix_server_gen.c))
-discards whatever `prte_iof.push_stdin` returns and reports
-`PMIX_SUCCESS`; PMIx's completion handler for the server-side read
-(`opcbfn` in its `src/common/pmix_iof.c`) explicitly discards the status
-and re-arms the read regardless. So a refusal anywhere in this framework
-would not slow the producer — it would only **drop bytes**. Queueing is
-the only lossless behavior available, and queueing is what happens.
-
-The consequence worth knowing: a daemon's stdin sink is bounded only by
+## Flow control
+
+stdin can outrun a slow reader, so the framework applies XON/XOFF
+back-pressure keyed on `PRTE_IOF_MAX_INPUT_BUFFERS` (50 queued chunks).
+**It reaches all the way back to the producer, and it did not always** —
+until PMIx grew `PMIx_server_IOF_flow_control` the whole mechanism was a
+signal the HNP logged and discarded. If you are reading an older tree, or
+an older copy of this guide, that is what it is describing.
+
+The producer is never ours. stdin originates in a PMIx server's read of
+its own stdin, or in a tool's `PMIx_IOF_push`, so the only way to slow it
+is to ask PMIx to stop it at the source. That is what
+`PMIx_server_IOF_flow_control` does: it leaves the read un-armed, so the
+bytes stay in the producer's own input stream and the OS applies the
+back-pressure. **Nothing is buffered on behalf of a suspended stream and
+nothing is dropped** — an XOFF is not permission to lose data.
+
+The two halves, which are separate mechanisms that happen to share a
+vocabulary:
+
+- **A daemon** whose stdin sink crosses 50 (or whose write errors out)
+  calls `prte_iof_prted_send_xonxoff(PRTE_IOF_XOFF)`, latched by
+  `prte_mca_iof_prted_component.xoff`, and sends `PRTE_IOF_XON` when the
+  backlog drains. `prte_iof_hnp_recv` screens the
+  `PRTE_IOF_XON | PRTE_IOF_XOFF` mask ahead of everything else and turns
+  it into a wildcard `PMIx_server_IOF_flow_control` call. Wildcard because
+  the message says only *that* this daemon is behind, never which producer
+  filled it — so every process feeding us stdin is suspended, which is the
+  conservative reading.
+- **The HNP's own local procs** never involve the RML at all.
+  `push_stdin` returns `PRTE_ERR_OUT_OF_RESOURCE` when a local sink passes
+  the same threshold, latched by `prte_mca_iof_hnp_component.xoff`, and
+  the glue in
+  [`src/prted/pmix/pmix_server_gen.c`](../../prted/pmix/pmix_server_gen.c)
+  turns that into `PMIX_ERR_IOF_XOFF` on the `push_stdin` completion —
+  which PMIx reads as "I have the data, suspend the stream". The matching
+  release is `release_flow_control()` in
+  [`hnp/iof_hnp.c`](hnp/iof_hnp.c).
+
+**Every XOFF must be paired with an XON, and that is on us.** PMIx has no
+status meaning "resume": a suspension persists until somebody calls the
+API with `xoff` false. So a release has to run on *every* path a
+backed-up sink can leave that state by — not just the one where it
+drains, but the ones where it is torn down. `release_flow_control()` is
+called from the `check:` and `finish:` arms of `stdin_write_handler`
+**and** from `hnp_close`/`hnp_complete`, which release a stdin sink
+directly and never reach the write handler at all — which is exactly the
+case that matters, since a sink is backed up precisely when its proc
+stopped reading, and that is the proc a teardown is likely to be
+retiring. It is a no-op when no XOFF is outstanding, so it is safe to
+call from anywhere, which is what lets it be called from everywhere it
+must be.
+
+**The blast radius is the job, not the DVM** — know why, because it is
+not obvious and it is what makes the failure tolerable. PMIx keeps no
+sticky suspension state: `pmix_iof_flow_control` walks the peers that
+exist *at the moment of the call* and pushes the request to them, and
+the only durable state is the `xoff` flag on the read event inside the
+producer itself. The producer is `prun`, which is per-job, so a
+suspension it is still carrying dies when it does. A `prun` that
+connects afterwards has never been told anything and starts reading
+normally. So the worst an unreleased XOFF can do is hang **that** job;
+the next one is unaffected, and the first sink to drain clears the stale
+latch. Do not read that as license to skip a release — a hung job is
+still a bug, and the reasoning above is a property of PMIx's current
+design rather than a guarantee it owes us. The dockerswarm suite counts
+the two and fails on a mismatch rather than merely checking that flow
+control happened.
+
+The oscillation is expected. Several procs take stdin at different rates,
+so one draining proc can turn the producers back on while another is
+still behind; the one still behind asserts XOFF again on its next write.
+That is the intended failure mode — the alternative to oscillating is
+stalling, and nothing is dropped either way. The prted module's `CHECK`
+block carries the same caveat in its own comment.
+
 `prte_iof_base_output_limit` (MCA param `iof_base_output_limit`, default
-`INT_MAX`, i.e. unbounded). Pipe a very large file into a process that
-reads it very slowly and the daemon's backlog grows without limit. Set
-`iof_base_output_limit` to a finite value and the write handler declares
-IOF hopelessly behind and fires `PRTE_JOB_STATE_FORCED_EXIT` instead.
-
-Making XOFF *mean* something requires a way to stop the read at the
-source, which is a PMIx-side change (honor the `push_stdin` completion
-status and hold the stdin read event until the host says go). Until that
-exists, leave the latch alone: it is a signal the HNP logs, not a control
-loop. What is **not** optional is that the HNP keep screening the tag
-first — a flow-control message carries no proc and no payload, so falling
-through to the output unpack reports the daemon's XOFF to the user as a
-corrupted message, which is exactly what used to happen every time a
-process read its stdin slowly.
+`INT_MAX`) is still the harder ceiling underneath all of this: if a
+sink's backlog exceeds it, the write handler concludes something is
+permanently wedged and fires `PRTE_JOB_STATE_FORCED_EXIT`. With flow
+control working, reaching it means the producer ignored the suspension,
+not merely that the reader is slow.
+
+**Against a PMIx that predates the capability** (`PRTE_PMIX_IOF_FLOW_CONTROL`
+is 0, from `PRTE_CHECK_PMIX_CAP([IOF_FLOW_CONTROL])` in
+`config/prte_setup_pmix.m4`) every one of these paths compiles away and
+the old behavior returns: the message is consumed quietly, the sink
+queues, and `iof_base_output_limit` is the only bound. What is **not**
+conditional is that the HNP screen the tag first — a flow-control message
+carries no proc and no payload, so falling through to the output unpack
+reports the daemon's XOFF to the user as a corrupted message, which is
+what used to happen every time a process read its stdin slowly.
 
 ---
 

src/mca/iof/hnp/AGENTS.md

@@ -103,9 +103,11 @@ before the proc unpack.** The order in the handler is deliberate:
    `PMIx_Data_unpack(…PMIX_PROC)` fails and reports the daemon's XOFF to
    the user as a corrupted message, which is what happened every time a
    process read its stdin slowly. We recognize it by mask (the control
-   bits are disjoint from every stream bit), trace it, and return. We do
-   not act on it — see the framework guide's *Flow control* section for
-   why acting on it would drop bytes rather than slow the producer.
+   bits are disjoint from every stream bit), trace it, and hand it to
+   `PMIx_server_IOF_flow_control` — wildcard, because the message says
+   only *that* this daemon is behind and never which producer filled it.
+   See the framework guide's *Flow control* section, and note the pairing
+   obligation it describes: an XOFF we never release hangs the producer.
 2. **Relayed stdin.** A leading `PRTE_IOF_STDIN` is stdin a daemon is
    relaying on behalf of a tool attached to *it* rather than to us, and
    the proc that follows is the intended **recipient**, not a source. That

src/mca/iof/hnp/iof_hnp.c

@@ -75,6 +75,7 @@ static void hnp_complete(const prte_job_t *jdata);
 static int finalize(void);
 
 static int push_stdin(const pmix_proc_t *dst_name, uint8_t *data, size_t sz);
+static void release_flow_control(void);
 
 /* The API's in this module are solely used to support LOCAL
  * procs - i.e., procs that are co-located to the HNP. Remote
@@ -273,10 +274,14 @@ static int push_stdin(const pmix_proc_t *dst_name, uint8_t *data, size_t sz)
                 if (PRTE_IOF_MAX_INPUT_BUFFERS < prte_iof_base_write_output(&proct->name,
                                                                             PRTE_IOF_STDIN, data, sz,
                                                                             proct->stdinev->wev)) {
-                    /* getting too backed up - stop the read event for now if it is still active */
-
+                    /* getting too backed up - the data is queued, so this is
+                     * "taken, now slow down" rather than a refusal. Latch it
+                     * so stdin_write_handler knows to send the matching XON
+                     * when this sink drains; without that pairing the
+                     * producer would stay suspended forever */
                     PMIX_OUTPUT_VERBOSE((1, prte_iof_base_framework.framework_output,
                                          "buffer backed up - holding"));
+                    prte_mca_iof_hnp_component.xoff = true;
                     return PRTE_ERR_OUT_OF_RESOURCE;
                 }
             }
@@ -359,6 +364,12 @@ static int hnp_close(const pmix_proc_t *peer, prte_iof_tag_t source_tag)
         if (PMIX_CHECK_PROCID(&proct->name, peer)) {
             if (PRTE_IOF_STDIN & source_tag) {
                 if (NULL != proct->stdinev) {
+                    /* a sink released here never reaches the write handler,
+                     * so this is the only chance to release an XOFF it was
+                     * responsible for. A sink is backed up precisely when
+                     * its proc stopped reading, which is precisely the proc
+                     * a close is likely to be tearing down */
+                    release_flow_control();
                     PMIX_RELEASE(proct->stdinev);
                 }
                 proct->stdinev = NULL;
@@ -409,6 +420,10 @@ static void hnp_complete(const prte_job_t *jdata)
              * killed does not, and that is the case this exists for.
              */
             if (NULL != proct->stdinev) {
+                /* as in close(): this sink will never reach the write
+                 * handler, so an XOFF it was responsible for has to be
+                 * released here or it outlives the job that caused it */
+                release_flow_control();
                 PMIX_RELEASE(proct->stdinev);
                 proct->stdinev = NULL;
             }
@@ -436,6 +451,37 @@ static int finalize(void)
 /* this function is called by the event library and thus
  * can access information global to the state machine
  */
+/* Release an XOFF we asserted because one of our own local procs' stdin
+ * sinks was backed up. A no-op unless we actually asserted one, so it is
+ * safe to call from every path a sink can leave the backed-up state by -
+ * including the paths where the sink is torn down rather than drained,
+ * which would otherwise leave the producers suspended forever.
+ *
+ * Against a PMIx that predates PMIx_server_IOF_flow_control there is
+ * nothing to release: push_stdin's PRTE_ERR_OUT_OF_RESOURCE is discarded
+ * by the glue rather than turned into an XOFF, so nothing was suspended. */
+static void release_flow_control(void)
+{
+#if PRTE_PMIX_IOF_FLOW_CONTROL
+    pmix_status_t prc;
+
+    if (!prte_mca_iof_hnp_component.xoff) {
+        return;
+    }
+    /* clear the latch first: the API completes inline, and a re-entrant
+     * assert would otherwise be lost */
+    prte_mca_iof_hnp_component.xoff = false;
+    PMIX_OUTPUT_VERBOSE((1, prte_iof_base_framework.framework_output,
+                         "%s iof:hnp releasing stdin flow control",
+                         PRTE_NAME_PRINT(PRTE_PROC_MY_NAME)));
+    prc = PMIx_server_IOF_flow_control(NULL, PMIX_FWD_STDIN_CHANNEL, false,
+                                       NULL, 0, NULL, NULL);
+    if (PMIX_SUCCESS != prc && PMIX_OPERATION_SUCCEEDED != prc) {
+        PMIX_ERROR_LOG(prc);
+    }
+#endif
+}
+
 static void stdin_write_handler(int fd, short event, void *cbdata)
 {
     prte_iof_sink_t *sink = (prte_iof_sink_t *) cbdata;
@@ -525,13 +571,22 @@ static void stdin_write_handler(int fd, short event, void *cbdata)
     PRTE_IOF_SINK_ACTIVATE(wev);
 
 check:
+    if (pmix_list_get_size(&wev->outputs) < PRTE_IOF_MAX_INPUT_BUFFERS) {
+        /* this proc has absorbed enough to justify restarting the producers
+         * we suspended */
+        release_flow_control();
+    }
     if (sink->closed && 0 == pmix_list_get_size(&wev->outputs)) {
         /* the sink has already been closed and everything was written, time to release it */
         PMIX_RELEASE(sink);
     }
     return;
 
 finish:
+    /* this sink is going away, so it will never drain - if it is the one
+     * that suspended the producers, they have to be let go here or they
+     * stay suspended for the life of the job */
+    release_flow_control();
     PMIX_RELEASE(wev);
     sink->wev = NULL;
     return;