clusterflux-public/crates/clusterflux-node/src/assignment_runner/control_watcher.rs

use std::collections::HashMap;
use std::sync::atomic::{AtomicBool, Ordering};
use std::sync::{Arc, Mutex};
use std::thread;
use std::time::Duration;

use clusterflux_core::TaskSpec;
use clusterflux_node::WasmDebugControl;

use crate::coordinator_session::CoordinatorSession;
use crate::daemon::{worker_shutdown_requested, Args};
use crate::node_identity::signed_node_request_json;

const DEFAULT_DEBUG_FREEZE_TIMEOUT_MILLIS: u64 = 5_000;

fn debug_freeze_timeout() -> Duration {
    std::env::var("CLUSTERFLUX_DEBUG_FREEZE_TIMEOUT_MS")
        .ok()
        .and_then(|value| value.parse::<u64>().ok())
        .filter(|millis| *millis > 0)
        .map(Duration::from_millis)
        .unwrap_or_else(|| Duration::from_millis(DEFAULT_DEBUG_FREEZE_TIMEOUT_MILLIS))
}

fn debug_handle_snapshot(handles: &Arc<Mutex<HashMap<u64, TaskSpec>>>) -> Vec<(String, String)> {
    let Ok(handles) = handles.lock() else {
        return vec![(
            "handle-registry-diagnostic".to_owned(),
            "runtime task handle registry was unavailable".to_owned(),
        )];
    };
    let mut snapshot = handles
        .iter()
        .map(|(handle_id, spec)| {
            (
                format!("task_handle_{handle_id}"),
                format!(
                    "definition={} instance={} state=active",
                    spec.task_definition, spec.task_instance
                ),
            )
        })
        .collect::<Vec<_>>();
    snapshot.sort_by(|left, right| left.0.cmp(&right.0));
    snapshot
}

#[allow(clippy::too_many_arguments)]
pub(super) fn spawn_task_control_watcher(
    args: Args,
    process: String,
    task: String,
    task_definition: String,
    node_private_key: String,
    cancellation_requested: Arc<AtomicBool>,
    abort_requested: Arc<AtomicBool>,
    debug_control: Arc<WasmDebugControl>,
    task_args: Vec<(String, String)>,
    handles: Arc<Mutex<HashMap<u64, TaskSpec>>>,
    command_status: Arc<Mutex<Option<String>>>,
    stop: Arc<AtomicBool>,
) {
    thread::spawn(move || {
        let Ok(mut session) = CoordinatorSession::connect(&args.coordinator) else {
            return;
        };
        while !stop.load(Ordering::Acquire) {
            let request = signed_node_request_json(
                &args,
                &node_private_key,
                "poll_task_control",
                serde_json::json!({
                    "type": "poll_task_control",
                    "tenant": &args.tenant,
                    "project": &args.project,
                    "process": &process,
                    "node": &args.node,
                    "task": &task,
                }),
            );
            let Ok(request) = request else {
                return;
            };
            let Ok(response) = session.request(request) else {
                return;
            };
            cancellation_requested.store(
                response
                    .get("cancel_requested")
                    .and_then(serde_json::Value::as_bool)
                    .unwrap_or(false),
                Ordering::Release,
            );
            if response
                .get("abort_requested")
                .and_then(serde_json::Value::as_bool)
                .unwrap_or(false)
            {
                abort_requested.store(true, Ordering::Release);
                return;
            }
            let debug_request = signed_node_request_json(
                &args,
                &node_private_key,
                "poll_debug_command",
                serde_json::json!({
                    "type": "poll_debug_command",
                    "tenant": &args.tenant,
                    "project": &args.project,
                    "process": &process,
                    "node": &args.node,
                    "task": &task,
                }),
            );
            let Ok(debug_request) = debug_request else {
                return;
            };
            let Ok(debug_response) = session.request(debug_request) else {
                return;
            };
            if let (Some(epoch), Some(command)) = (
                debug_response
                    .get("epoch")
                    .and_then(serde_json::Value::as_u64),
                debug_response
                    .get("command")
                    .and_then(serde_json::Value::as_str),
            ) {
                let freeze_timeout = debug_freeze_timeout();
                let (state, message) = match command {
                    "freeze" => {
                        if std::env::var_os("CLUSTERFLUX_DEBUG_CONTROL_TRACE").is_some() {
                            eprintln!(
                                "clusterflux debug control: node received freeze for epoch {epoch} task {task} (debug={:?})",
                                Arc::as_ptr(&debug_control)
                            );
                        }
                        debug_control.request_freeze(epoch);
                        if debug_control.wait_until_frozen(epoch, freeze_timeout) {
                            ("frozen", None)
                        } else {
                            debug_control.request_resume(epoch);
                            (
                                "failed",
                                Some(format!(
                                    "node execution did not reach a freezeable Wasm safepoint or verified native/Podman boundary within {} ms",
                                    freeze_timeout.as_millis()
                                )),
                            )
                        }
                    }
                    "resume" => {
                        debug_control.request_resume(epoch);
                        if debug_control.wait_until_running(epoch, freeze_timeout) {
                            ("running", None)
                        } else {
                            (
                                "failed",
                                Some(format!(
                                    "node execution did not leave its verified frozen state within {} ms",
                                    freeze_timeout.as_millis()
                                )),
                            )
                        }
                    }
                    _ => (
                        "failed",
                        Some(format!("node received unknown debug command `{command}`")),
                    ),
                };
                let report = signed_node_request_json(
                    &args,
                    &node_private_key,
                    "report_debug_state",
                    serde_json::json!({
                        "type": "report_debug_state",
                        "tenant": &args.tenant,
                        "project": &args.project,
                        "process": &process,
                        "node": &args.node,
                        "task": &task,
                        "epoch": epoch,
                        "state": state,
                        "stack_frames": if state == "frozen" {
                            let mut frames = debug_control.stack_frames();
                            if let Some(frame) = frames.first_mut() {
                                *frame = format!("{task_definition}::wasm / {frame}");
                            } else {
                                frames.push(format!("{task_definition}::wasm"));
                            }
                            frames
                        } else {
                            Vec::new()
                        },
                        "local_values": [],
                        "task_args": &task_args,
                        "handles": debug_handle_snapshot(&handles),
                        "command_status": command_status
                            .lock()
                            .ok()
                            .and_then(|status| status.clone()),
                        "recent_output": [],
                        "message": message,
                    }),
                );
                let Ok(report) = report else {
                    return;
                };
                if session.request(report).is_err() {
                    return;
                }
            }
            thread::sleep(Duration::from_millis(50));
        }
    });
}

pub(super) fn spawn_worker_shutdown_watcher(
    abort_requested: Arc<AtomicBool>,
    debug_control: Arc<WasmDebugControl>,
    stop: Arc<AtomicBool>,
) {
    thread::spawn(move || {
        while !stop.load(Ordering::Acquire) {
            if worker_shutdown_requested() {
                abort_requested.store(true, Ordering::Release);
                if let Some(epoch) = debug_control.requested_epoch() {
                    debug_control.request_resume(epoch);
                }
                return;
            }
            thread::sleep(Duration::from_millis(10));
        }
    });
}