jmri_server.programming_queue

  1"""``ProgrammingQueue`` — serialized per-loco CV writes with progress + cancel.
  2
  3DCC programming is slow: direct-mode writes on the service track take
  4several seconds *per CV*. A full loco can have dozens of CVs the user
  5wants to write. Firing them all off in parallel would overwhelm the
  6programming track (which serializes at the CS anyway) and race the
  7LocoCvStore atomic writes.
  8
  9This module owns a per-loco queue. One batch runs at a time. A cancel
 10flag causes the queue to stop after the *in-flight* write completes —
 11matching the user's spec ("cancel after the completion of the in-flight
 12operation") because DCC hardware can't safely abort a write already on
 13the wire.
 14
 15Client flow:
 16
 17    POST /programmer/loco/<id>/write-batch  → 202 + batch id
 18    GET  /programmer/loco/<id>/batch/status → progress
 19    POST /programmer/loco/<id>/batch/cancel → stop after current write
 20    POST /programmer/loco/<id>/batch/clear  → clear terminal-state batch
 21"""
 22from __future__ import annotations
 23
 24import asyncio
 25import logging
 26import time
 27import uuid
 28from collections.abc import Awaitable, Callable
 29from dataclasses import dataclass, field
 30from enum import StrEnum
 31from typing import Any
 32
 33_log = logging.getLogger(__name__)
 34
 35
 36class BatchState(StrEnum):
 37    """High-level lifecycle for a queued batch."""
 38    PENDING     = "pending"
 39    IN_PROGRESS = "in_progress"
 40    DONE        = "done"
 41    CANCELLED   = "cancelled"
 42    ERROR       = "error"
 43
 44
 45class WriteResult(StrEnum):
 46    """Per-write outcome inside a batch."""
 47    PENDING   = "pending"
 48    WRITING   = "writing"
 49    DONE      = "done"
 50    ERROR     = "error"
 51    CANCELLED = "cancelled"
 52
 53
 54@dataclass
 55class BatchWrite:
 56    """One CV operation inside a batch.
 57
 58    * ``kind`` — ``"write"`` (default) or ``"read"``. Selects which
 59      callback the queue dispatches to. Reads leave ``value`` as 0 on
 60      the request side and store the read byte in ``value`` once the
 61      operation completes so downstream consumers see a uniform
 62      ``{cv, value, item, result}`` snapshot.
 63    * ``cv`` — target CV number
 64    * ``value`` — byte value: for writes this is the requested value;
 65      for reads it starts at 0 and holds the read result on completion.
 66    * ``item`` — display label for the progress UI (variable ``item``
 67      name from the decoder catalog)
 68    * ``result`` — updated as the queue processes this op
 69    * ``error`` — filled when ``result == ERROR``
 70    """
 71    cv: int
 72    value: int = 0
 73    item: str = ""
 74    kind: str = "write"
 75    result: WriteResult = WriteResult.PENDING
 76    error: str = ""
 77
 78
 79# Write callback: returns the value the CS confirmed committed.
 80WriteCallback = Callable[[int, int, str, dict[str, Any]], Awaitable[int]]
 81
 82# Read callback: returns the byte read from the decoder.
 83ReadCallback = Callable[[int, str, dict[str, Any]], Awaitable[int]]
 84
 85
 86@dataclass
 87class Batch:
 88    """One in-flight batch of writes."""
 89
 90    batch_id: str
 91    loco_id: str
 92    writes: list[BatchWrite]
 93    mode: str = "service"
 94    address: int | None = None
 95    long_address: bool = False
 96    file_name: str | None = None
 97    state: BatchState = BatchState.PENDING
 98    error: str = ""
 99    started_at: float = field(default_factory=time.monotonic)
100    finished_at: float | None = None
101    _cancel: asyncio.Event = field(default_factory=asyncio.Event)
102    _task: asyncio.Task[None] | None = None
103
104    @property
105    def completed(self) -> int:
106        return sum(
107            1 for w in self.writes
108            if w.result in (WriteResult.DONE, WriteResult.ERROR, WriteResult.CANCELLED)
109        )
110
111    @property
112    def total(self) -> int:
113        return len(self.writes)
114
115    @property
116    def in_flight(self) -> BatchWrite | None:
117        for w in self.writes:
118            if w.result is WriteResult.WRITING:
119                return w
120        return None
121
122    def is_terminal(self) -> bool:
123        return self.state in (
124            BatchState.DONE, BatchState.CANCELLED, BatchState.ERROR,
125        )
126
127    def request_cancel(self) -> None:
128        """Set the cancel flag; the worker stops after the current write."""
129        self._cancel.set()
130
131    def to_status(self) -> dict[str, Any]:
132        """Serializable snapshot for the polling endpoint."""
133        in_flight = self.in_flight
134        return {
135            "batch_id": self.batch_id,
136            "loco_id": self.loco_id,
137            "state": self.state.value,
138            "completed": self.completed,
139            "total": self.total,
140            "in_flight": {
141                "cv": in_flight.cv, "item": in_flight.item,
142                "value": in_flight.value,
143            } if in_flight else None,
144            "writes": [
145                {
146                    "cv": w.cv, "value": w.value, "item": w.item,
147                    "result": w.result.value, "error": w.error,
148                }
149                for w in self.writes
150            ],
151            "error": self.error,
152            "elapsed_s": (
153                (self.finished_at or time.monotonic()) - self.started_at
154            ),
155        }
156
157
158class BatchInFlightError(Exception):
159    """Raised when a caller tries to start a new batch while one is running."""
160
161    def __init__(self, loco_id: str, batch_id: str) -> None:
162        super().__init__(
163            f"loco {loco_id!r} already has batch {batch_id!r} in flight; "
164            f"cancel or wait for it to finish first",
165        )
166        self.loco_id = loco_id
167        self.batch_id = batch_id
168
169
170class ProgrammingQueue:
171    """Holds the active :class:`Batch` for one loco.
172
173    A completed batch stays visible on the queue until the caller either
174    starts a new batch (which replaces it) or explicitly clears it. That
175    lets the UI show the final result without racing the completion
176    notification.
177
178    ``read_callback`` is optional; when omitted, only write batches can
179    run (a read op in the batch will fail with a NotImplementedError).
180    """
181
182    def __init__(
183        self, loco_id: str, write_callback: WriteCallback,
184        read_callback: ReadCallback | None = None,
185    ) -> None:
186        self._loco_id = loco_id
187        self._write_callback = write_callback
188        self._read_callback = read_callback
189        self._current: Batch | None = None
190
191    @property
192    def loco_id(self) -> str:
193        return self._loco_id
194
195    @property
196    def current(self) -> Batch | None:
197        return self._current
198
199    def is_busy(self) -> bool:
200        return (
201            self._current is not None
202            and self._current.state is BatchState.IN_PROGRESS
203        )
204
205    def clear(self) -> None:
206        """Drop a terminal-state batch. Raises if a batch is still running."""
207        if self._current is not None and not self._current.is_terminal():
208            raise BatchInFlightError(self._loco_id, self._current.batch_id)
209        self._current = None
210
211    def cancel(self) -> Batch | None:
212        """Signal cancel on the current batch. Returns it (or None)."""
213        if self._current is None or self._current.is_terminal():
214            return self._current
215        self._current.request_cancel()
216        return self._current
217
218    async def start(
219        self,
220        writes: list[BatchWrite],
221        *,
222        mode: str = "service",
223        address: int | None = None,
224        long_address: bool = False,
225        file_name: str | None = None,
226    ) -> Batch:
227        """Queue *writes* and start the worker.
228
229        Raises :class:`BatchInFlightError` if a batch is already running.
230        Callers should either await its completion, cancel it, or clear
231        it (if terminal) before starting a new one.
232        """
233        if self._current is not None and not self._current.is_terminal():
234            raise BatchInFlightError(self._loco_id, self._current.batch_id)
235        batch = Batch(
236            batch_id=uuid.uuid4().hex[:12],
237            loco_id=self._loco_id,
238            writes=list(writes),
239            mode=mode,
240            address=address,
241            long_address=long_address,
242            file_name=file_name,
243        )
244        self._current = batch
245        batch.state = BatchState.IN_PROGRESS
246        batch._task = asyncio.create_task(
247            self._run(batch), name=f"batch-{self._loco_id}-{batch.batch_id}",
248        )
249        return batch
250
251    async def _run(self, batch: Batch) -> None:
252        """Iterate ops serially, honouring the cancel flag between each."""
253        try:
254            for op in batch.writes:
255                if batch._cancel.is_set():
256                    op.result = WriteResult.CANCELLED
257                    continue
258                op.result = WriteResult.WRITING
259                extra: dict[str, Any] = {
260                    "mode": batch.mode,
261                    "file_name": batch.file_name,
262                }
263                if batch.address is not None:
264                    extra["address"] = batch.address
265                    extra["long_address"] = batch.long_address
266                try:
267                    if op.kind == "read":
268                        if self._read_callback is None:
269                            raise NotImplementedError(
270                                "read op requested but no read_callback configured",
271                            )
272                        value = await self._read_callback(
273                            op.cv, op.item, extra,
274                        )
275                        # Reads land the byte back into ``value`` so
276                        # the status snapshot's shape stays uniform
277                        # with writes.
278                        op.value = int(value)
279                    else:
280                        await self._write_callback(
281                            op.cv, op.value, op.item, extra,
282                        )
283                except Exception as e:  # noqa: BLE001
284                    op.result = WriteResult.ERROR
285                    op.error = str(e)
286                    _log.warning(
287                        "batch %s: CV %d %s failed: %s",
288                        batch.batch_id, op.cv, op.kind, e,
289                    )
290                    continue
291                op.result = WriteResult.DONE
292            # Terminal state: cancelled if flag was set at any point,
293            # error if all failed, otherwise done.
294            any_pending = any(w.result is WriteResult.PENDING for w in batch.writes)
295            any_cancelled = any(w.result is WriteResult.CANCELLED for w in batch.writes)
296            all_errored = batch.writes and all(
297                w.result is WriteResult.ERROR for w in batch.writes
298            )
299            if batch._cancel.is_set() or any_cancelled or any_pending:
300                batch.state = BatchState.CANCELLED
301            elif all_errored:
302                batch.state = BatchState.ERROR
303                batch.error = "all writes failed"
304            else:
305                batch.state = BatchState.DONE
306        except asyncio.CancelledError:
307            batch.state = BatchState.CANCELLED
308            raise
309        except Exception as e:  # noqa: BLE001
310            batch.state = BatchState.ERROR
311            batch.error = str(e)
312            _log.exception("batch %s crashed", batch.batch_id)
313        finally:
314            batch.finished_at = time.monotonic()
315
316
317def get_or_create_queue(
318    ctx: Any, loco_id: str, write_callback: WriteCallback,
319    read_callback: ReadCallback | None = None,
320) -> ProgrammingQueue:
321    """Fetch the queue for *loco_id* off *ctx*, creating on first use.
322
323    The queue registry lives on ``ctx._programming_queues`` (dict). One
324    queue per loco keeps batches scoped so cancelling loco A doesn't
325    stop loco B's operation. When an existing queue is returned and a
326    non-None *read_callback* is supplied, the queue's read callback is
327    updated in place — lets a route that didn't set up reads on first
328    use enable them later without swapping the queue instance.
329    """
330    queues: dict[str, ProgrammingQueue] = getattr(ctx, "_programming_queues", None)
331    if queues is None:
332        queues = {}
333        ctx._programming_queues = queues  # type: ignore[attr-defined]
334    if loco_id not in queues:
335        queues[loco_id] = ProgrammingQueue(
336            loco_id, write_callback, read_callback=read_callback,
337        )
338    elif read_callback is not None and queues[loco_id]._read_callback is None:
339        queues[loco_id]._read_callback = read_callback
340    return queues[loco_id]