fix: count an accepted connection in the drain before its task starts

docs/signals.md

@@ -117,7 +117,7 @@ and several subsystems compose without knowing about each other.
 | `request()` | begin a drain, from anywhere (a `/quit` route, a supervisor) |
 | `requested()` | whether a drain has begun — poll this in a long-lived handler |
 | `register_listener(fd)` / `close_listener(fd)` | listener registry; `close_listener` closes at most once |
-| `enter()` / `leave()` | bracket one unit of in-flight work |
+| `enter()` / `leave()` | bracket one unit of in-flight work — `enter()` on the accept loop, `leave()` from the task |
 | `expect_flush()` / `flush_done()` | a subsystem with shutdown work of its own |
 | `drain(deadline_ms)` | wait for both, returning what was left unfinished |
 | `set_drain_deadline(ms)` | the grace period every std server uses |
@@ -142,6 +142,27 @@ all of them, and through the same coordinator:
   one that would otherwise never get there, to a `UNAVAILABLE` status rather than
   a cut connection.
 
+### when a connection joins the count
+
+on the accept loop, in the same breath as the accept — not inside the task that
+serves it. between a `spawn` and the spawned task's first instruction a
+connection is invisible, and a drain landing in that window would see nothing
+outstanding and return, cutting off a client that had already been accepted.
+"every accepted connection finished" is the whole promise of a graceful drain,
+so the count starts at the accept.
+
+that leaves only one direction to be wrong in, and it is the safe one: work that
+is counted and then never runs holds its count up until the grace period
+expires, and `drain()` returns it as work it abandoned. bounded, and reported —
+where the other direction drops a connection silently. write your own accept
+loops the same way:
+
+```pith
+accepted := tcp.accept(fd)
+shutdown.enter()
+spawn serve(accepted.ok)   # serve() opens with `defer shutdown.leave()`
+```
+
 ## streams that never end
 
 a drain waits for work to finish, and a stream that never finishes would outlast

docs/tls.md

@@ -56,22 +56,34 @@ transfer: `Listener.close()` gives the borrow back, and whoever called
 the timing is what makes a server different in practice. an accept loop hands
 each socket to its own task, and that task reads the certificate and the private
 key out of the registry when its own handshake reaches them — which can be a
-whole handshake timeout after the accept that spawned it returned. so a server
-closes its config once it has drained, never before:
+whole handshake timeout after the accept that spawned it returned, and which for
+a connection accepted just before the shutdown may not have happened at all yet.
+so a server gives the borrow back and closes its config once it has drained,
+never before:
 
 ```pith
 cfg := tls.server_config("certs/server.crt", "certs/server.key")!
 listener := tls.listen(host, 443, cfg)!
+shutdown.register_listener(listener.handle)
 # ... accept, spawning a task per connection, until a shutdown is requested ...
-listener.close()
+shutdown.close_listener(listener.handle)   # free the port, and only the port
 drained := shutdown.drain_default()
+tls.release_listener_config(listener)      # give the borrow back
 cfg.close()
 ```
 
-closing before the drain does not corrupt anything — a handle is never reissued,
-so the handshake that loses the race fails cleanly rather than reaching for
-another config's key — but it does turn working connections into failed ones on
-the way out, which is a worse shutdown than the leak it was fixing.
+the two halves of `Listener.close()` are split on purpose, and the split is the
+point. the socket goes first, because the replacement instance of a rolling
+deploy is waiting to bind that port and the drain can take the whole grace
+period. the listener's *binding* to the config goes last, because that binding
+is how a handshake finds the certificate — dropping it while a connection is
+still in the drain count is the same mistake as closing the config early, one
+step removed.
+
+either mistake corrupts nothing — a handle is never reissued, so a handshake
+that loses the race fails cleanly rather than reaching for another config's key
+— but both turn working connections into failed ones on the way out, which is a
+worse shutdown than the leak they were fixing.
 
 the std servers do this for you: `App.listen_tls`, `http2.listen_h2_tls` and its
 streaming twin each build a config, serve on it, and close it after their drain,

std/net/http2/server.pith

@@ -1173,17 +1173,37 @@ pub fn serve_h2c_connection(fd: Int, handler: fn(http.HttpRequestBytes) -> http.
     conn.close()
 
 # serve_h2c_connection with the accept loop's concurrency permit released when
-# the connection ends, however it ends.
-#
-# the same bracket doubles as the drain count: a graceful shutdown does not
-# finish until every connection that reached here has run out, so a rolling
-# deploy cannot sever a stream that is still exchanging frames.
+# the connection ends, however it ends, and the drain count it was handed
+# released with it. both were taken on the accept loop, in spawn_h2c_conn; this
+# `defer` is the other half of each, and runs on every exit the task has.
 fn serve_h2c_slot(fd: Int, handler: fn(http.HttpRequestBytes) -> http.HttpResponse):
-    shutdown.enter()
     defer shutdown.leave()
     defer connection_slots.release()
     serve_h2c_connection(fd, handler)
 
+# hand an accepted h2c socket to its own task, having taken the two things a
+# connection holds for its whole life: the accept loop's concurrency permit and
+# a place in the drain count.
+#
+# ## why the drain count is taken here and not inside the task
+#
+# between `spawn` and the spawned task's first instruction the connection is
+# invisible. a shutdown that lands in that window drains to zero and returns,
+# and the connection — accepted, acknowledged to the peer, never served — is cut
+# off silently. "every accepted connection finished" is the whole promise of a
+# graceful drain, so the count has to start at the accept, on this task, before
+# the spawn.
+#
+# the failure mode this trades into is bounded and loud. a task that somehow
+# never runs leaves its count up, and `drain()` gives up at the grace period and
+# reports it as work it abandoned — the same number a rolling deploy already
+# logs. counting too little loses a connection with no trace; counting too much
+# costs at most the grace period and says so.
+fn spawn_h2c_conn(fd: Int, handler: fn(http.HttpRequestBytes) -> http.HttpResponse):
+    connection_slots.acquire()
+    shutdown.enter()
+    spawn serve_h2c_slot(fd, handler)
+
 # the connection lifecycle, factored out so the offline tests can drive it over a
 # ServerConn without a socket. reads the preface, sends the server SETTINGS,
 # spawns the writer, and runs the reader loop on the calling task; when the reader
@@ -1308,8 +1328,7 @@ pub fn listen_h2c(host: String, port: Int, handler: fn(http.HttpRequestBytes) ->
             tcp.back_off_after_accept_failure(failures, accepted.err)!
             continue
         failures = 0
-        connection_slots.acquire()
-        spawn serve_h2c_slot(accepted.ok, handler)
+        spawn_h2c_conn(accepted.ok, handler)
     # free the port before draining, not after: the replacement instance of a
     # rolling deploy is waiting to bind it, and the drain can take the whole
     # grace period. the defer above is then a no-op — close_listener closes at
@@ -1354,11 +1373,19 @@ fn handshake_and_serve_h2_tls(listener_handle: Int, fd: Int, handler: fn(http.Ht
 # connection stops the server accepting after MAX_CONNECTIONS attempts, which is
 # the same denial of service the accept loop was fixed for, only slower.
 fn serve_h2_tls_slot(listener_handle: Int, fd: Int, handler: fn(http.HttpRequestBytes) -> http.HttpResponse):
-    shutdown.enter()
     defer shutdown.leave()
     defer connection_slots.release()
     handshake_and_serve_h2_tls(listener_handle, fd, handler) catch false
 
+# hand an accepted tls socket to its own task, permit and drain count taken here
+# on the accept loop — see spawn_h2c_conn for why the count cannot wait for the
+# task to start. a connection still mid-handshake when a shutdown lands is a
+# client waiting for an answer, so it counts from the accept like any other.
+fn spawn_h2_tls_conn(listener_handle: Int, fd: Int, handler: fn(http.HttpRequestBytes) -> http.HttpResponse):
+    connection_slots.acquire()
+    shutdown.enter()
+    spawn serve_h2_tls_slot(listener_handle, fd, handler)
+
 # bind host:port for http/2 over tls and serve forever. it builds a tls server
 # config from the pem `cert`/`key` files, offers alpn "h2" (and "http/1.1", which
 # a peer may pick but this server does not yet serve — see serve_h2_tls_connection),
@@ -1394,15 +1421,18 @@ pub fn listen_h2_tls(host: String, port: Int, cert: String, key: String, handler
             tcp.back_off_after_accept_failure(failures, accepted.err)!
             continue
         failures = 0
-        connection_slots.acquire()
-        spawn serve_h2_tls_slot(listener.handle, accepted.ok, handler)
-    # free the port before draining (see listen_h2c above).
-    release_tls_listener(listener)
+        spawn_h2_tls_conn(listener.handle, accepted.ok, handler)
+    # free the port before draining (see listen_h2c above), but only the port.
+    # the listener's binding to the tls config outlives the accept loop: a
+    # connection accepted a moment before the shutdown has not necessarily
+    # reached its handshake yet, and the handshake finds the certificate and the
+    # key through that binding.
+    shutdown.close_listener(listener.handle)
     drained := shutdown.drain_default()
-    # this call built the config, so this call closes it — and only here, once
-    # the drain has finished. the listener merely borrowed it, and each
-    # connection task reads the certificate and key out of the registry when its
-    # own handshake gets that far. see docs/tls.md.
+    # the drain is over, so every connection that borrowed the config has
+    # finished with it: drop the listener's binding, then close the config this
+    # call built. see docs/tls.md.
+    release_tls_listener(listener)
     config.close()
     return drained
 
@@ -1426,13 +1456,20 @@ pub fn serve_h2c_streaming_connection(fd: Int):
     run_connection(conn, streaming_unused_handler, true) catch false
     conn.close()
 
-# serve_h2c_streaming_connection with the accept loop's permit released at the end.
+# serve_h2c_streaming_connection with the accept loop's permit and drain count
+# released at the end. mirrors serve_h2c_slot.
 fn serve_h2c_streaming_slot(fd: Int):
-    shutdown.enter()
     defer shutdown.leave()
     defer connection_slots.release()
     serve_h2c_streaming_connection(fd)
 
+# hand an accepted h2c socket to its own streaming task, permit and drain count
+# taken here on the accept loop. mirrors spawn_h2c_conn.
+fn spawn_h2c_streaming_conn(fd: Int):
+    connection_slots.acquire()
+    shutdown.enter()
+    spawn serve_h2c_streaming_slot(fd)
+
 # bind host:port for cleartext h2c and serve forever, handing each request
 # stream to `handler` while it is still arriving. mirrors listen_h2c.
 pub fn listen_h2c_streaming(host: String, port: Int, handler: fn(ServerStream) -> Int) -> Int!:
@@ -1449,8 +1486,7 @@ pub fn listen_h2c_streaming(host: String, port: Int, handler: fn(ServerStream) -
             tcp.back_off_after_accept_failure(failures, accepted.err)!
             continue
         failures = 0
-        connection_slots.acquire()
-        spawn serve_h2c_streaming_slot(accepted.ok)
+        spawn_h2c_streaming_conn(accepted.ok)
     # free the port before draining, not after: the replacement instance of a
     # rolling deploy is waiting to bind it, and the drain can take the whole
     # grace period. the defer above is then a no-op — close_listener closes at
@@ -1477,11 +1513,17 @@ fn handshake_and_serve_h2_tls_streaming(listener_handle: Int, fd: Int) -> Bool!:
 # handshake_and_serve_h2_tls_streaming with the accept loop's permit released at
 # the end, on the failed-handshake path included. mirrors serve_h2_tls_slot.
 fn serve_h2_tls_streaming_slot(listener_handle: Int, fd: Int):
-    shutdown.enter()
     defer shutdown.leave()
     defer connection_slots.release()
     handshake_and_serve_h2_tls_streaming(listener_handle, fd) catch false
 
+# hand an accepted tls socket to its own streaming task, permit and drain count
+# taken here on the accept loop. mirrors spawn_h2_tls_conn.
+fn spawn_h2_tls_streaming_conn(listener_handle: Int, fd: Int):
+    connection_slots.acquire()
+    shutdown.enter()
+    spawn serve_h2_tls_streaming_slot(listener_handle, fd)
+
 # bind host:port for http/2 over tls and serve forever on the streaming path.
 # mirrors listen_h2_tls, handshake off the accept loop included.
 pub fn listen_h2_tls_streaming(host: String, port: Int, cert: String, key: String, handler: fn(ServerStream) -> Int) -> Int!:
@@ -1509,12 +1551,12 @@ pub fn listen_h2_tls_streaming(host: String, port: Int, cert: String, key: Strin
             tcp.back_off_after_accept_failure(failures, accepted.err)!
             continue
         failures = 0
-        connection_slots.acquire()
-        spawn serve_h2_tls_streaming_slot(listener.handle, accepted.ok)
-    # free the port before draining (see listen_h2c above).
-    release_tls_listener(listener)
+        spawn_h2_tls_streaming_conn(listener.handle, accepted.ok)
+    # free the port before draining, and drop the listener's binding to the tls
+    # config only once the drain is over — see listen_h2_tls.
+    shutdown.close_listener(listener.handle)
     drained := shutdown.drain_default()
-    # this call built the config, so this call closes it — see listen_h2_tls.
+    release_tls_listener(listener)
     config.close()
     return drained
 
@@ -1947,3 +1989,102 @@ test "handle_streaming_stream parses the head and drives the streaming handler":
     assert(has_field(tail.fields, "x-done", "1"))
     # the stream was forgotten when the handler returned.
     assert_eq(server.active_count(), 0)
+
+# --- accept-time drain registration ---
+#
+# these drive the accept loop's hand-off (spawn_h2c_conn and its siblings)
+# directly, over real sockets, without an accept loop: the loop's only other job
+# is `tcp.accept`, and calling the hand-off by hand is what lets a test look at
+# the drain count in the window between the spawn and the task's first
+# instruction — the window the count exists to cover.
+
+# open `count` connected socket pairs against a listener on `port` and return the
+# server-side fds. the client ends are pushed onto `clients`, so the test can
+# close them and let the connection tasks finish; the listener is closed here,
+# since every connection it will ever carry has already been accepted.
+fn accepted_pairs(port: Int, count: Int, clients: List[Int]) -> List[Int]:
+    listener := tcp_listen("127.0.0.1", port) catch -1
+    assert(listener > 0)
+    mut served: List[Int] := []
+    mut i := 0
+    while i < count:
+        connected := tcp.connect("127.0.0.1", port)
+        assert(not connected.is_err)
+        clients.push(connected.ok)
+        accepted := tcp.accept(listener)
+        assert(not accepted.is_err)
+        served.push(accepted.ok)
+        i = i + 1
+    tcp_close(listener)
+    return served
+
+test "an accepted h2c connection is in the drain count before its task runs":
+    shutdown.reset()
+    mut clients: List[Int] := []
+    served := accepted_pairs(18871, 8, clients)
+    for fd in served:
+        spawn_h2c_conn(fd, ok_handler)
+    # the count is taken on this task, so it is complete the instant the last
+    # hand-off returns — it cannot depend on how many of the eight tasks the
+    # runtime has got round to starting. with the count taken inside the task
+    # instead, most of these are still invisible here, and a shutdown landing
+    # now would drain to zero over eight unserved connections.
+    assert_eq(shutdown.inflight(), 8)
+    # none of them can finish (their peers have sent nothing and are still
+    # open), so a drain must give up at its deadline and report all eight rather
+    # than hang or claim a clean shutdown.
+    assert_eq(shutdown.drain(20), 8)
+    for c in clients:
+        tcp_close(c)
+    # each task leaves exactly once, on the path where the connection dies
+    # before its first frame: no double count, nothing left behind.
+    assert_eq(shutdown.drain(5000), 0)
+    assert_eq(shutdown.inflight(), 0)
+    shutdown.reset()
+
+test "an accepted h2c streaming connection is in the drain count before its task runs":
+    shutdown.reset()
+    stream_handler_fn = echo_stream_handler
+    mut clients: List[Int] := []
+    served := accepted_pairs(18872, 4, clients)
+    for fd in served:
+        spawn_h2c_streaming_conn(fd)
+    assert_eq(shutdown.inflight(), 4)
+    for c in clients:
+        tcp_close(c)
+    assert_eq(shutdown.drain(5000), 0)
+    shutdown.reset()
+
+test "an accepted tls connection is in the drain count before its handshake starts":
+    shutdown.reset()
+    config := tls.server_config("tests/live/fixtures/localhost.crt", "tests/live/fixtures/localhost.key")!.with_alpn(["h2"])
+    bound := tls.listen("127.0.0.1", 18873, config)
+    assert(not bound.is_err)
+    listener := bound.ok
+    mut clients: List[Int] := []
+    mut served: List[Int] := []
+    mut i := 0
+    while i < 4:
+        # a plain tcp peer that never sends a client hello, so every handshake
+        # parks on its first read and no task can reach its `leave`.
+        connected := tcp.connect("127.0.0.1", 18873)
+        assert(not connected.is_err)
+        clients.push(connected.ok)
+        accepted := tls.accept_socket(listener)
+        assert(not accepted.is_err)
+        served.push(accepted.ok)
+        i = i + 1
+    for fd in served:
+        spawn_h2_tls_conn(listener.handle, fd, ok_handler)
+    # a connection still mid-handshake is a client waiting for an answer, and it
+    # counts from the accept — not from whenever the runtime starts its task.
+    assert_eq(shutdown.inflight(), 4)
+    for c in clients:
+        tcp_close(c)
+    # the handshakes fail once their peers are gone; that path leaves the count
+    # exactly once too, so the drain finishes clean.
+    assert_eq(shutdown.drain(5000), 0)
+    assert_eq(shutdown.inflight(), 0)
+    release_tls_listener(listener)
+    config.close()
+    shutdown.reset()

std/shutdown.pith

@@ -23,8 +23,9 @@
 ##   `shutdown(2)` on each, which both stops new connections and wakes the
 ##   waiting accept loop immediately — no polling for the flag. the loop still
 ##   closes its own listener, so nothing closes an fd another task is using.
-## - **in-flight work.** each connection task brackets itself with
-##   `enter()`/`leave()`, so `drain()` knows when the last one is done.
+## - **in-flight work.** a server calls `enter()` on its accept loop, the moment
+##   a connection is accepted, and the task it hands that connection to calls
+##   `leave()` when it is done — so `drain()` knows when the last one is out.
 ## - **flushes.** a subsystem with its own shutdown work — the telemetry
 ##   exporter is the one in std — calls `expect_flush()` up front and
 ##   `flush_done()` when it has finished. `drain()` waits for those too, which
@@ -157,8 +158,16 @@ pub fn listeners() -> Int:
     mu.unlock()
     return n
 
-## mark one unit of in-flight work started. a server calls this as a connection
-## task begins; `drain()` will not report a clean shutdown until it leaves.
+## mark one unit of in-flight work started; `drain()` will not report a clean
+## shutdown until it leaves.
+##
+## call it *before* handing the work to a task, not as the task's first act.
+## between a `spawn` and the spawned task's first instruction the work is
+## invisible: it has not entered, a drain landing in that window sees nothing
+## outstanding and returns, and the connection it was about to serve is cut with
+## no response and no trace. counting from the accept can only ever count too
+## much — work that never starts — and that costs at most the grace period and
+## is reported by `drain()`, where counting too little loses a client silently.
 pub fn enter():
     mu.lock()
     inflight_count = inflight_count + 1