dora-ulite6/dora_ulite6/control.py

"""The threads that own the arm, and the mailboxes the event loop drains.

**Nothing outside this file may call into the driver**, with two exceptions
written to be safe from another thread: :meth:`ULite6.stop`, and
:meth:`ULite6.read_report`, which reads the controller's report cache rather
than the wire.

There are two threads because a motion must not stop the arm being watched. The
xArm SDK blocks: a move waits out the motion, and a read from a controller whose
cable was pulled blocks until its own socket timeout. So:

- The **control thread** runs one job at a time and is allowed to block. It is
  where every service is carried out, including the motions that take seconds.
- The **reporter thread** never blocks. It reads the report cache at a fixed
  cadence, so `tcp_pose` and `joint_pose` keep flowing right through a ten-second
  move — which is the only way a caller can follow one, since the thread that
  started it is inside it until it ends.
- The **stopper thread** exists only to call `stop`. It cannot be the control
  thread (a stop must not queue behind the motion it aborts) and it must not be
  the event loop: the SDK's `emergency_stop` re-commands the controller until it
  confirms, for up to three seconds, and a dataflow router cannot go away for
  three seconds. Requests to it coalesce — several stops in a row are one stop.

The dora event loop is a third party to both: it drains mailboxes and publishes,
and never waits for the arm.
"""

import queue
import threading
import time
from dataclasses import dataclass, field
from typing import Any

from dora_ulite6 import services
from dora_ulite6.config import Config
from dora_ulite6.driver import ULite6
from dora_ulite6.errors import ArmNeedsHomingError, RobotUnavailableError
from dora_ulite6.log import log
from dora_ulite6.pose import RobotPose, with_tool_offset

# The states the node reports on `robot_state`, in the order a healthy start
# walks through them.
CONNECTING = "connecting"
READY = "ready"
MOVING = "moving"
RECOVERING = "recovering"
UNAVAILABLE = "unavailable"


# The motion lifecycle published on `motion_state`. `accepted` exists precisely
# to close the race a bare moving/idle flag leaves open: between a caller being
# told its request was taken and the control thread reaching it, a two-valued
# state still reads "not moving", and a caller polling for completion would
# conclude the motion it just asked for had already finished.
ACCEPTED = "accepted"
DONE = "done"
FAILED = "failed"
PREEMPTED = "preempted"


@dataclass
class ServiceCall:
    """One service request waiting on the control thread."""

    request_id: str
    row: dict
    # Whether the caller is held until the arm stops. False means the response
    # was already sent as an acknowledgement, and completion is only reported on
    # `motion_state` — which is what keeps it one response per request.
    wait: bool = True
    # Set when a newer motion took this one's place, before or during its run.
    preempted_by: str | None = None

    @property
    def service(self) -> str:
        """The service this request names."""
        return self.row.get(services.SERVICE, "")

    @property
    def moves_the_arm(self) -> bool:
        """Whether carrying this out drives the arm, and so preempts and is preemptable."""
        return self.service in services.MOVES_THE_ARM


@dataclass
class Jog:
    """One velocity command. Not a service: nobody waits for an answer."""

    velocity: list[float]


@dataclass
class Telemetry:
    """One reading of the arm: joint angles in degrees, tool-tip pose in mm and degrees."""

    joints: list[float]
    tip: RobotPose
    timestamp_ns: int


@dataclass
class Message:
    """Something to publish: an output id, its fields, and its correlation metadata."""

    output: str
    fields: dict[str, Any]
    metadata: dict[str, str] = field(default_factory=dict)


class Control:
    """Owns the arm and hands the event loop what there is to publish."""

    def __init__(self, config: Config, restarted: bool = False):
        self.config = config
        # A restarted node inherits an arm that is wherever it was left. Driving
        # it to INIT_MODE because the process came back is a motion nobody asked
        # for, so a restart skips it.
        self.restarted = restarted
        self.robot: ULite6 | None = None
        self._jobs: queue.Queue[ServiceCall | Jog] = queue.Queue()
        self._lock = threading.Lock()
        self._telemetry: Telemetry | None = None
        self._outbox: list[Message] = []
        self._state = CONNECTING
        self._state_changes: list[str] = [CONNECTING]
        # The motion the control thread is inside right now, and whether the
        # last preemption left a software stop for the next motion to clear.
        self._running_motion: ServiceCall | None = None
        self._stopped_for_preempt = False
        # Set by anyone who needs the arm halted; the stopper thread drains it.
        self._stop_wanted = threading.Event()
        # Whether a velocity command is currently keeping the arm moving, and
        # when the last one arrived — the controller expires them on its own.
        self._jogging = False
        self._last_jog = 0.0
        # Whether the arm has been reached once and put where INIT_MODE wants it.
        self._started = False
        # Set to ask the threads to finish. The sense matters: this is waited on
        # as the threads' sleep, and an event that is *set* while running makes
        # every wait() return at once — a sleep that never sleeps, which is a
        # busy loop rather than a cadence.
        self._stopping = threading.Event()
        self._control = threading.Thread(target=self._run_control, name="ulite6-control", daemon=True)
        self._reporter = threading.Thread(target=self._run_reporter, name="ulite6-reporter", daemon=True)
        self._stopper = threading.Thread(target=self._run_stopper, name="ulite6-stopper", daemon=True)

    def start(self) -> None:
        """Start the three threads; they reach the arm themselves, so this never blocks."""
        self._control.start()
        self._reporter.start()
        self._stopper.start()

    def shutdown(self) -> bool:
        """Halt the arm, stop the threads, and say whether they really finished.

        Halting first is the point. A motion lives in the *controller*, not in
        this process: dropping the connection does not abort it, so a node that
        exits mid-move leaves an arm travelling with nobody left to stop it. The
        stop is commanded here on the caller's thread and waited for, because by
        now the event loop has finished routing and there is nothing left to
        starve.

        Returns False if the control thread is still inside the SDK afterwards,
        which is the caller's cue to leave without waiting for it.
        """
        if self.robot is not None:
            self._attempt(self.robot.stop)
        self._stopping.set()
        self._stop_wanted.set()
        self._stopper.join(timeout=1.0)
        self._reporter.join(timeout=1.0)
        self._control.join(timeout=self.config.motion.ready_timeout_s)
        if self._control.is_alive():
            log("the control thread is still inside the SDK; leaving it behind")
            return False
        return True

    # --- what the event loop puts in ---

    def submit(self, job: ServiceCall | Jog) -> None:
        """Queue work for the control thread. Returns immediately."""
        self._jobs.put(job)

    def request_stop(self, request_id: str = "") -> None:
        """Ask for the arm to be halted, and come straight back.

        Answering before the arm has actually stopped is deliberate: the reply
        says the stop was commanded, and whether it took is reported where every
        other outcome is, on `motion_state` and `robot_state`. Waiting here
        instead would hand the caller a better answer at the price of the whole
        dataflow, for up to three seconds.
        """
        self._stop_wanted.set()
        if request_id:
            self._respond(
                ServiceCall(request_id=request_id, row={services.SERVICE: services.STOP}),
                (True, {"message": "stop commanded"}),
            )

    def _preempt_motions(self, by: str) -> None:
        """Mark the running motion and any still waiting, and halt the arm if one was moving."""
        with self._lock:
            running = self._running_motion
            waiting = [
                job
                for job in list(self._jobs.queue)
                if isinstance(job, ServiceCall) and job.moves_the_arm and job.preempted_by is None
            ]
            for job in [running, *waiting]:
                if job is not None:
                    job.preempted_by = by
            # The next motion to actually run clears the software stop this is
            # about to cause. Only ever set here, never unset: preempting a
            # motion that had not started yet stops nothing, and must not clear a
            # flag left by an earlier preemption that nothing has consumed — that
            # would send the next motion at an arm still in its stopped state,
            # where it would be refused the moment it was commanded.
            self._stopped_for_preempt = self._stopped_for_preempt or running is not None
        if running is not None:
            log(f"{running.service} ({running.request_id}) preempted by {by}")
            self.request_stop()

    def start_jog(self, velocity: list[float]) -> None:
        """Take one velocity command, preempting a motion the first time it moves the arm.

        A jog arrives many times a second, so everything here happens on the
        transitions only: a stream of them is one preemption and one state
        change, not hundreds.
        """
        moving = any(velocity)
        if moving and not self._jogging:
            self._preempt_motions("jog")
            self._set_state(MOVING)
        with self._lock:
            was_jogging, self._jogging = self._jogging, moving
            self._last_jog = time.monotonic()
        # Only a jog that was in fact jogging ends one. A zero velocity from a
        # teleop source that is merely idling says nothing about a motion that
        # owns the arm, and must not answer for it.
        if not moving and was_jogging:
            self._idle_unless_busy()
        self._jobs.put(Jog(velocity=velocity))

    def start_motion(self, call: ServiceCall) -> None:
        """Take a motion, preempting whatever motion was already in play.

        The arm is one resource: there is never more than one motion, so a new
        one replaces the old rather than queueing behind it. The running motion
        is cut short through the arm — the same stop a caller can send — and any
        motion still waiting its turn is dropped without ever moving anything.

        Each preempted request is answered exactly once, as every request is: if
        its caller is still waiting, by its response; if it was a `wait: false`
        request, it was acknowledged when it arrived, and the preemption shows
        only on `motion_state`.
        """
        self._preempt_motions(call.request_id)
        self._motion_state(call, ACCEPTED, "queued")
        if not call.wait:
            self._respond(call, (True, {"message": "accepted", "accepted": True}))
        self._jobs.put(call)

    def post(self, message: Message) -> None:
        """Add a message for the event loop to publish on the next tick.

        The loop uses this to answer a request it refused itself, so a malformed
        message is still answered under its own correlation id without the arm
        threads being troubled by it.
        """
        self._post(message)

    def call_now(self, request_id: str, row: dict) -> None:
        """Run a service on the calling thread, for the ones that must not queue.

        Only `stop` takes this path: queued, it would sit behind the very motion
        it exists to abort. It is safe here because the driver method it reaches
        is the one written for another thread.
        """
        call = ServiceCall(request_id=request_id, row=row)
        if self.robot is None:
            self._respond(call, (False, {"message": "the arm has not been reached yet"}))
            return
        self._respond(call, self._attempt(lambda: services.run(self.robot, row)))

    # --- what the event loop drains ---

    def take_telemetry(self) -> Telemetry | None:
        """Return the newest reading and clear it, so a tick never republishes a stale pose."""
        with self._lock:
            telemetry, self._telemetry = self._telemetry, None
            return telemetry

    def take_messages(self) -> list[Message]:
        """Return every response recorded since the last call."""
        with self._lock:
            outbox, self._outbox = self._outbox, []
            return outbox

    def take_state_changes(self) -> list[str]:
        """Return every state transition since the last call."""
        with self._lock:
            changes, self._state_changes = self._state_changes, []
            return changes

    def _post(self, message: Message) -> None:
        with self._lock:
            self._outbox.append(message)

    # --- the control thread ---

    def _run_control(self) -> None:
        while not self._stopping.is_set():
            if not self._started:
                self._start_up()
                continue
            try:
                job = self._jobs.get(timeout=0.1)
            except queue.Empty:
                continue
            if isinstance(job, Jog):
                self._run_jog(job)
            else:
                self._run_service(job)

    def _start_up(self) -> None:
        """Reach an arm that answers, put it where INIT_MODE says, and start serving.

        The driver is built once and then retried through its own ``connect()``,
        which reuses the SDK handle: rebuilding it on every attempt would leak a
        socket and its threads for every second an absent arm stays absent.
        """
        self._set_state(CONNECTING)
        if self.robot is None:
            self.robot = ULite6(self.config)
        try:
            if not self.robot.is_connected() and not self.robot.connect():
                raise RobotUnavailableError(f"robot at {self.config.robot_ip} is not reachable")
            self._start_position()
        except RobotUnavailableError as error:
            log(f"{error}; retrying in {self.config.reconnect_interval_s:.0f}s")
            self._set_state(UNAVAILABLE)
            self._stopping.wait(self.config.reconnect_interval_s)
            return
        self._started = True
        self._set_state(READY)

    def _start_position(self) -> None:
        """Put the arm where INIT_MODE says it should start, if anywhere."""
        if self.restarted and self.config.init_mode != "none":
            log(f"restarted: leaving the arm where it is rather than running INIT_MODE={self.config.init_mode}")
            return
        if self.config.init_mode == "home":
            log("moving to home")
            self.robot.go_home()
        elif self.config.init_mode == "pose":
            log(f"moving to the initial pose {self.config.init_pose.as_list()}")
            self.robot.move_to(
                with_tool_offset(self.config.init_pose, -self.config.tool_offset_mm),
                self.config.init_speed,
            )

    def _run_service(self, call: ServiceCall) -> None:
        """Carry out one service and post its response."""
        if call.moves_the_arm:
            self._run_motion(call)
            return
        self._respond(call, self._attempt(lambda: services.run(self.robot, call.row)))

    def _run_motion(self, call: ServiceCall) -> None:
        """Carry out one motion, reporting its whole life on `motion_state`.

        Blocks for as long as the arm takes. `robot_state` says `moving`
        meanwhile, so a caller watching it can tell a slow service from a stalled
        one, and `tcp_pose` keeps arriving because the reporter thread is not the
        one in here.
        """
        if call.preempted_by is not None:
            # Replaced while it waited its turn; it never moved anything.
            self._finish_motion(call, PREEMPTED, False, f"preempted by {call.preempted_by}")
            return
        with self._lock:
            self._running_motion = call
            clear_first = self._stopped_for_preempt
            self._stopped_for_preempt = False
        self._set_state(MOVING)
        self._motion_state(call, MOVING, "")
        if self._leave_jog():
            # A pose command would switch the controller out of velocity mode
            # behind the driver's back, leaving its own bookkeeping saying the
            # arm is still jogging. Leaving through the front door keeps them
            # agreed.
            self._attempt(lambda: self.robot.set_cartesian_velocity([0.0] * 6))
        if clear_first:
            # Preempting leaves the controller in its software-stopped state.
            # The driver's own set_state(0) would probably be enough; clearing
            # first is cheap and does not depend on that being true.
            self._attempt(self.robot.clear)
        outcome = self._attempt(lambda: services.run(self.robot, call.row))
        with self._lock:
            self._running_motion = None
            preempted_by = call.preempted_by
        ok, fields = outcome
        if preempted_by is not None:
            # A preempted move fails through MotionInterruptedError like any
            # aborted one; only the flag says who cut it short.
            self._finish_motion(call, PREEMPTED, False, f"preempted by {preempted_by}")
        else:
            self._finish_motion(call, DONE if ok else FAILED, ok, fields.get("message", ""), fields)
        self._idle_unless_busy()

    def _leave_jog(self) -> bool:
        """Clear the jogging flag, reporting whether the arm was in fact jogging."""
        with self._lock:
            was_jogging, self._jogging = self._jogging, False
        return was_jogging

    def _finish_motion(self, call: ServiceCall, state: str, ok: bool, message: str, fields=None) -> None:
        """Publish a motion's end, and answer its caller if one is still waiting."""
        self._motion_state(call, state, message)
        if call.wait:
            self._respond(call, (ok, dict(fields or {"message": message})))

    def _motion_state(self, call: ServiceCall, state: str, message: str) -> None:
        """Publish where a motion stands.

        The request_id rides in the *fields*, not the metadata: metadata carrying
        a request_id means "this is the answer to that request", and this is a
        broadcast of the arm's state, which anyone may read. Keeping it in the
        fields is also what lets a caller polling `get_state` see the same value
        as one watching the topic.
        """
        self._post(
            Message(
                output="motion_state",
                fields={
                    "request_id": call.request_id,
                    "service": call.service,
                    "state": state,
                    "message": message,
                },
            )
        )

    def _run_jog(self, jog: Jog) -> None:
        """Command one velocity. Nothing is published: a jog has no answer."""
        def command() -> dict:
            self.robot.set_cartesian_velocity(jog.velocity)
            return {"message": "jogging"}

        ok, fields = self._attempt(command)
        if not ok:
            log(f"jog refused: {fields['message']}")

    def _respond(self, call: ServiceCall, outcome: tuple[bool, dict]) -> None:
        """Post the one response this request gets."""
        ok, fields = outcome
        fields.setdefault("message", "")
        self._post(
            Message(
                output="response",
                fields={"ok": ok, "service": call.service, **fields},
                metadata={"request_id": call.request_id},
            )
        )

    def _attempt(self, work) -> tuple[bool, dict]:
        """Run ``work``, turning whatever it raises into a failed response."""
        try:
            return True, work()
        except ArmNeedsHomingError as error:
            return False, {"message": f"{error}; call `home_and_clear` once the area is clear"}
        except RobotUnavailableError as error:
            self._go_unavailable(error)
            return False, {"message": str(error)}
        except Exception as error:  # noqa: BLE001 — a failed request must never take the thread down
            return False, {"message": str(error)}

    def _go_unavailable(self, error: Exception) -> None:
        """Report the arm as unusable and try once to bring it back.

        A single attempt, not a retry loop: the thread has jobs to serve, and the
        reporter comes round again in milliseconds.
        """
        self._set_state(RECOVERING)
        log(f"the arm is unusable ({error}); trying to recover")
        try:
            self.robot.recover()
        except RobotUnavailableError as recovery_error:
            log(f"recovery failed: {recovery_error}")
            self._set_state(UNAVAILABLE)
            return
        self._set_state(READY)

    # --- the stopper thread ---

    def _run_stopper(self) -> None:
        """Halt the arm whenever anyone asks, and never make them wait for it.

        One stop per wake-up however many asked, which is what a stop means:
        several callers wanting the arm halted want the same single thing.
        """
        while not self._stopping.is_set():
            if not self._stop_wanted.wait(timeout=0.1):
                continue
            self._stop_wanted.clear()
            if self._stopping.is_set():
                return
            if self.robot is not None:
                self._attempt(self.robot.stop)

    def _set_state(self, state: str) -> None:
        """Record a state, queueing it for publication only when it is a change."""
        with self._lock:
            if state == self._state:
                return
            self._state = state
            self._state_changes.append(state)
        log(f"state: {state}")

    def _idle_unless_busy(self) -> None:
        """Report the arm idle, but only when nothing is still driving it.

        One thing ending is not the arm stopping. A motion may still be winding
        down inside the SDK after the jog that preempted it went quiet; a jog may
        have taken the arm over from the motion that just finished. Either way
        `ready` published over it is the false idle the state exists to rule out.

        The state also stops being ours to set once it has left `moving` — that
        is `unavailable` or `recovering` reporting something this does not know
        about, and it must not be overwritten with `ready`.
        """
        with self._lock:
            if self._running_motion is not None or self._jogging or self._state != MOVING:
                return
        self._set_state(READY)

    # --- the reporter thread ---

    def _expire_jog(self) -> None:
        """Stop calling the arm jogging once the controller has expired the command.

        A jog lives on the controller's own TTL, so a sender that goes quiet
        leaves an arm that stopped moving a moment later. Nothing tells us; the
        clock does.
        """
        with self._lock:
            if not self._jogging or time.monotonic() - self._last_jog < self.config.motion.velocity_command_ttl_s:
                return
            self._jogging = False
        self._idle_unless_busy()

    def _run_reporter(self) -> None:
        """Publish telemetry, whatever the control thread is doing.

        Never touches anything that can block, so it keeps running through a
        ten-second move — which is what lets a caller follow a motion whose
        response will not arrive until it ends.
        """
        while not self._stopping.is_set():
            self._stopping.wait(self.config.telemetry_interval_s)
            if self.robot is None:
                continue
            self._expire_jog()
            reading = self.robot.read_report()
            if reading is None:
                continue
            joints, flange = reading
            tip = with_tool_offset(flange, self.config.tool_offset_mm)
            with self._lock:
                self._telemetry = Telemetry(joints, tip, time.time_ns())