aion-server 0.31.0

//! The gRPC stream's task-refusal arm: re-parking one dispatch a worker
//! declined for want of a free execution slot.
//!
//! Split out of `worker_grpc.rs`: the ordered three-act re-park (fence, clear
//! tracking, resolve attempt-neutrally) is a self-contained decision, and
//! keeping it beside the stream handler pushed that file past the per-file
//! length budget.

use aion_proto::generated;

use crate::worker::dispatch::ActivityCompletion;

use super::WorkerSession;

/// Re-park one dispatch a worker declined for want of a free execution slot.
///
/// Three acts, in this order and for this reason:
///
/// 1. **Prove the generation FIRST.** The execution-generation token goes to the
///    completion fence exactly as a result's does, and everything else happens
///    inside the acceptance. A refusal from a superseded generation changes
///    nothing at all.
/// 2. **Clear the in-flight tracking, once accepted.** This is the count
///    correction. It frees the slot the server believed this worker was using,
///    so the re-selection does not choose the same full worker again, and it
///    releases the tracker's liveness entry for an activity that is not running.
///    It runs in `after_accept`, which is before the waiter is resolved, so the
///    correction still precedes any re-selection.
/// 3. **Resolve the waiter attempt-neutrally.** The dispatch is resolved with a
///    transport-domain reason, which the engine's retry loop re-dispatches on
///    the SAME attempt — a fresh selection, a fresh lease, a fresh per-attempt
///    clock. Nothing reaches the workflow and nothing is recorded.
///
/// The first two used to be the other way round, and it was wrong in a way the
/// fence could not save. `clear_completed_task_tracking` is keyed by
/// `(worker, workflow, activity)` and carries no token, so it retires whatever
/// attempt is tracked under that key: a late refusal from a superseded
/// generation untracked the LIVE attempt and returned a slot the worker was
/// still using, and the fence then rejected the refusal far too late to undo it.
/// The result arm sixteen lines up never had that inversion.
///
/// A malformed refusal is logged and dropped: with no ids there is nothing to
/// re-park, and the dispatch's other resolutions (the worker's eventual result,
/// the stream teardown sweep) are untouched by the frame being unreadable.
pub(super) fn handle_task_refusal(session: &WorkerSession<'_>, refusal: generated::TaskRefused) {
    let (Some(workflow_id), Some(activity_id)) = (refusal.workflow_id, refusal.activity_id) else {
        tracing::error!(
            worker_id = ?session.worker_id,
            reason = refusal.reason,
            "malformed task-refusal frame: it names no activity, so there is nothing to re-park"
        );
        return;
    };
    let workflow_id = match aion_core::WorkflowId::try_from(aion_proto::ProtoWorkflowId {
        uuid: workflow_id.uuid,
    }) {
        Ok(workflow_id) => workflow_id,
        Err(error) => {
            tracing::error!(
                worker_id = ?session.worker_id,
                %error,
                "malformed task-refusal frame: its workflow id could not be decoded"
            );
            return;
        }
    };
    let activity_id = aion_core::ActivityId::from(aion_proto::ProtoActivityId {
        sequence_position: activity_id.sequence_position,
    });
    let completion_token = match crate::worker::CompletionToken::from_wire(
        &workflow_id,
        &activity_id,
        refusal.completion_token,
    ) {
        Ok(token) => token,
        Err(error) => {
            tracing::error!(
                worker_id = ?session.worker_id,
                workflow_id = %workflow_id,
                activity_id = %activity_id,
                %error,
                "task refusal carries no usable execution-generation token; leaving the dispatch                  to its other resolutions rather than re-parking it unfenced"
            );
            return;
        }
    };
    // INSIDE the closure, exactly as the result arm at `:413` does it, and for a
    // reason the refusal arm originally got backwards.
    //
    // `clear_completed_task_tracking` is keyed by `(worker, workflow, activity)`
    // and takes no completion token, so it retires WHATEVER attempt is tracked
    // under that key. Run before the fence, a late refusal from a superseded
    // generation therefore untracked the LIVE attempt that had taken the site
    // over and handed back a slot the worker was still using — the expiry sweep
    // could no longer see that dispatch, drain accounting was off by one, and
    // the worker's selection capacity was over-reported. The fence rejected the
    // refusal a few lines later, far too late to undo any of it.
    //
    // `after_accept` runs before the waiter is resolved, so clearing here still
    // precedes any re-selection — the property the original ordering was
    // reaching for — while a refusal the fence rejects now touches nothing.
    let after_accept = || {
        // Fail open: the tracking clear is bookkeeping, and the refusal's
        // re-park is the load-bearing half. A poisoned lock here must not
        // withhold the re-park.
        let _ = crate::worker::bridge::clear_completed_task_tracking(
            session.heartbeat,
            session.registry,
            session.worker_id,
            &workflow_id,
            &activity_id,
        );
        Ok(())
    };
    if let Err(error) = session.pending.complete_activity_after_accept(
        ActivityCompletion {
            workflow_id: workflow_id.clone(),
            activity_id: activity_id.clone(),
            run_id: None,
            completion_token,
            outcome: crate::worker::ActivityCompletionOutcome::Refused {
                worker_id: session.worker_id,
                reason: refusal.reason,
            },
        },
        after_accept,
    ) {
        tracing::error!(
            worker_id = ?session.worker_id,
            workflow_id = %workflow_id,
            activity_id = %activity_id,
            %error,
            "task refusal rejected by execution-generation proof; it is from a generation this              site has already moved past and re-parking it would displace the attempt that              superseded it"
        );
    }
}