Skip to content

← Concurrency with Threads and asyncio step 34 of 56

Hard Primitives

Release a Connection However the Stream Ends

stream_rows holds a pooled connection while it yields rows. Rewrite the consumer, _run, so the connection goes back to the pool before you build the report. That must hold whether the loop finishes, breaks, raises or is cancelled.

The starter gives you Pool and stream_rows. Keep them as they are:

  • pool.acquire(name) is an @asynccontextmanager. On entry it adds one to pool.live and logs "acquire:conn". Its finally takes one off and logs "release:conn", without awaiting.
  • stream_rows(pool, count) enters pool.acquire("conn"), then yields count rows, 400, 401 and so on. With a count of zero or less it yields nothing, but it still acquires and releases once you loop over it.

The starter's _run loops without aclosing. After a break, the connection is still out when it builds the report.

What to write

Rewrite _run(count, stop_after, mode). solve already runs it with asyncio.run.

  • Loop over stream_rows(pool, count) inside async with aclosing(...), and append each row to seen.
  • After you append a row, check whether stop_after >= 0 and len(seen) >= stop_after. If so, act on mode as the table says. A stop_after of 0 triggers after the first row. A negative stop_after never triggers.
mode When the check triggers outcome
"complete" Keep looping. "completed"
"break" Break out of the loop. "completed"
"error" Raise RuntimeError, and catch it outside the consumer. "error"
"cancel" Signal the driver, then wait. The driver cancels the consumer and awaits it. "cancelled"

If the check never triggers, every mode ends with "completed".

For "cancel", run the consumer as its own task with asyncio.create_task, so the driver can cancel it. Signal with an asyncio.Event. The consumer sets it, then waits on a long asyncio.sleep, and the driver cancels the consumer. The driver must also wake up when the consumer finishes without signalling, for example when the stream has fewer rows than stop_after. If it waits only for the signal, it waits forever. Cancel the consumer only if it is still running.

Build the report only after the consumer and its cleanup have finished:

{"seen": seen, "log": list(log), "outcome": outcome, "live": pool.live}

In every case, live must be 0 and log must be ["acquire:conn", "release:conn"]. Keep the list(log) copy. It records the log at that moment. A shared list would also pick up a release that happens later, while asyncio.run shuts down, and make a late release look on time.

For count=5, stop_after=2 and mode="cancel", seen is [400, 401] and outcome is "cancelled".

Your code must pass mypy --strict. aclosing needs an object with an aclose() method. The return type of stream_rows, AsyncGenerator[int, None], includes one. AsyncIterator would not.