Bounded buffers and ordered workers

kitchen/kitchen.py

"""Two waiters clip order tickets on a rail that holds 3, and two chefs take them and cook.

The rail is asyncio.Queue(maxsize=3): put() waits while it is full and get() while it is
empty. A waiter needs 1 simulated second to take an order and a chef 4 to cook one, so
orders arrive four times as fast as the two chefs can cook them."""
import asyncio

from fake_clock import FakeClock

ORDERS_EACH, TAKE_ORDER, COOK = 6, 1, 4


async def kitchen(rail_size, show_seconds=0):
    clock = FakeClock()
    rail = asyncio.Queue(maxsize=rail_size)     # maxsize 0: no limit
    on_rail, cooking, blocked = [], [], []      # what the timeline shows
    clipped, cooked, waited, most = [], [], {"A": 0, "B": 0}, 0

    async def waiter(name):
        nonlocal most
        for n in range(1, ORDERS_EACH + 1):
            await clock.sleep(TAKE_ORDER)
            ticket, asked = f"{name}{n}", clock.now
            blocked.append(ticket)
            await rail.put(ticket)               # waits while the rail is full
            blocked.remove(ticket)
            on_rail.append(ticket)
            waited[name] += clock.now - asked
            clipped.append(ticket)
            most = max(most, rail.qsize())

    async def chef():
        while True:
            ticket = await rail.get()            # waits while the rail is empty
            on_rail.remove(ticket)
            cooking.append(ticket)
            await clock.sleep(COOK)
            cooking.remove(ticket)
            cooked.append(ticket)
            rail.task_done()

    async def timeline():
        await clock.sleep(0.5)                   # look half-way between events
        for _ in range(show_seconds):
            print(f"{clock.now:5} s  {' '.join(on_rail) or '-':12}  {' '.join(cooking) or '-':8}  "
                  f"{' '.join(blocked) or '-'}")
            await clock.sleep(1)

    async def service():
        chefs = [asyncio.create_task(chef()) for _ in range(2)]
        watcher = asyncio.create_task(timeline())
        await asyncio.gather(waiter("A"), waiter("B"))
        await rail.join()                        # every clipped ticket has been cooked
        for task in [*chefs, watcher]:
            task.cancel()

    await clock.run(service())
    return clock.now, sorted(clipped) == sorted(cooked) and len(cooked), waited, most


async def main():
    print("time     on the rail   cooking   waiting for room")
    end, cooked, waited, most = await kitchen(3, show_seconds=8)
    print(f"a rail of 3: {cooked} tickets cooked by {end} s; at most {most} on the rail; "
          f"waiters waited {waited['A']} s and {waited['B']} s for room")
    end, cooked, waited, most = await kitchen(0)
    print(f"no limit:    {cooked} tickets cooked by {end} s; at most {most} on the rail; "
          f"waiters waited {waited['A']} s and {waited['B']} s for room")


asyncio.run(main())

time     on the rail   cooking   waiting for room
  0.5 s  -             -         -
  1.5 s  -             A1 B1     -
  2.5 s  A2 B2         A1 B1     -
  3.5 s  A2 B2 A3      A1 B1     B3
  4.5 s  A2 B2 A3      A1 B1     B3 A4
  5.5 s  A3 B3 A4      A2 B2     -
  6.5 s  A3 B3 A4      A2 B2     B4 A5
  7.5 s  A3 B3 A4      A2 B2     B4 A5
a rail of 3: 12 tickets cooked by 25 s; at most 3 on the rail; waiters waited 7 s and 11 s for room
no limit:    12 tickets cooked by 25 s; at most 8 on the rail; waiters waited 0 s and 0 s for room

kitchen/fake_clock.py

"""A fake clock for asyncio: simulated seconds that pass only while every task waits.

Tasks call `await clock.sleep(seconds)` where real code would wait for a download.
`await clock.run(main())` runs a program: each time every task is waiting, it moves the
time on to the earliest wake-up and wakes that sleeper. Nothing really waits, so the
times a program prints are the same in every run and on every computer.
"""
import asyncio
import heapq
import itertools

# Turns of the event loop that let every task run on until it waits again. The programs
# of this lesson need at most 6; 12 leave room for a longer chain of awaits.
SETTLE_TURNS = 12


class FakeClock:
    def __init__(self):
        self.now = 0
        self._sleepers = []                 # (wake-up time, order, future), earliest first
        self._order = itertools.count()

    async def sleep(self, seconds):
        wake = asyncio.get_running_loop().create_future()
        heapq.heappush(self._sleepers, (self.now + seconds, next(self._order), wake))
        await wake

    async def run(self, main):
        """Runs the coroutine main to its end, and returns its result or raises its error."""
        outcome = {}

        async def capture():
            try:
                outcome["result"] = await main
            except Exception as error:      # handed to the caller below
                outcome["error"] = error

        task = asyncio.ensure_future(capture())
        while True:
            for _ in range(SETTLE_TURNS):
                await asyncio.sleep(0)
            if task.done():
                if "error" in outcome:
                    raise outcome["error"]
                return outcome["result"]
            if not self._sleepers:
                task.cancel()
                raise RuntimeError(f"stuck at {self.now} s: every task waits, and none of them for the clock")
            self.now, _, wake = heapq.heappop(self._sleepers)
            wake.set_result(None)