datum-agent 0.10.11

use std::{
    collections::{BTreeMap, BTreeSet, HashMap},
    fmt,
    net::SocketAddr,
    sync::{Arc, Mutex},
    time::{Duration, SystemTime},
};

use datum_cluster::{ClusterConfig, ClusterNode, ClusterState, Member, MemberState, Signal};
use datum_net::quic::quinn;
use tokio::{
    sync::{mpsc, oneshot, watch},
    task::JoinHandle,
    time::{Instant, sleep},
};

use crate::{
    Agent, AgentConfig, AgentHandle, ClusterJobMetadata, ClusterPlacementHistory,
    JobRegistryHandle, PlacementSpec as RegistryPlacementSpec,
    PlacementStrategy as RegistryPlacementStrategy,
    dcp::{
        ClientKind, ClusterJobList, ClusterJobNode, ClusterJobStart, ClusterNodeError,
        ClusterNodeList, ClusterNodeStatus, ClusterViewProvider, CompleteShardingAsk, DcpClient,
        DcpError, DcpJobFactories, DcpServer, DcpServerConfig, DcpServerHandle,
        ForwardShardEnvelopes, Hello, PlacementSpec, PlacementStrategy, RememberClusterAssignment,
        ResponseStatus, ShardAllocation, ShardAllocationTable, ShardEnvelopeBatchResult,
        ShardPipeClient, ShardPipeFrame, SubmitClusterJob,
        client::MetricSubscription,
        proto::MetricSample,
        server::{
            cluster_metadata_from_wire, placement_spec_from_wire, wire_cluster_job_start,
            wire_job_status,
        },
    },
};

/// Role used by cluster-aware agents to discover peers.
pub const AGENT_ROLE: &str = "agent";

/// Result type for cluster-aware agent startup and shutdown.
pub type ClusterAgentResult<T> = Result<T, ClusterAgentError>;

/// Errors returned by the cluster-aware agent entrypoint.
#[derive(Debug, thiserror::Error)]
pub enum ClusterAgentError {
    #[error("invalid cluster-agent config: {0}")]
    InvalidConfig(String),
    #[error(transparent)]
    Agent(#[from] crate::AgentError),
    #[error(transparent)]
    Dcp(#[from] DcpError),
    #[error(transparent)]
    Cluster(#[from] datum_cluster::ClusterError),
}

/// Transport used for node-to-node DCP sessions.
#[derive(Clone)]
pub enum NodeSessionTransport {
    /// Plaintext TCP is accepted only for loopback development and tests.
    TcpLoopback,
    /// QUIC with mTLS. The same client config is used for every peer.
    QuicMtls {
        server_name: String,
        client_config: quinn::ClientConfig,
    },
}

impl fmt::Debug for NodeSessionTransport {
    fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
        match self {
            Self::TcpLoopback => formatter.write_str("TcpLoopback"),
            Self::QuicMtls { server_name, .. } => formatter
                .debug_struct("QuicMtls")
                .field("server_name", server_name)
                .finish_non_exhaustive(),
        }
    }
}

/// Configuration for DCP node sessions maintained by a cluster-aware agent.
#[derive(Clone, Debug)]
pub struct NodeSessionConfig {
    /// Role that marks a member as running a DCP agent endpoint.
    pub agent_role: String,
    /// Node-to-node DCP transport.
    pub transport: NodeSessionTransport,
    /// Initial reconnect backoff.
    pub reconnect_min_backoff: Duration,
    /// Maximum reconnect backoff.
    pub reconnect_max_backoff: Duration,
    /// Default per-peer fan-out timeout.
    pub request_timeout: Duration,
    /// Bounded per-session command queue capacity.
    pub command_buffer: usize,
}

impl Default for NodeSessionConfig {
    fn default() -> Self {
        Self {
            agent_role: AGENT_ROLE.to_owned(),
            transport: NodeSessionTransport::TcpLoopback,
            reconnect_min_backoff: Duration::from_millis(50),
            reconnect_max_backoff: Duration::from_secs(2),
            request_timeout: Duration::from_millis(750),
            command_buffer: 32,
        }
    }
}

/// Combined configuration for one cluster-aware Datum agent.
#[derive(Clone, Default)]
pub struct ClusterAgentConfig {
    /// Local job registry configuration.
    pub agent: AgentConfig,
    /// Membership configuration.
    pub cluster: ClusterConfig,
    /// DCP listener configuration.
    pub dcp: DcpServerConfig,
    /// Node-session configuration.
    pub sessions: NodeSessionConfig,
}

impl fmt::Debug for ClusterAgentConfig {
    fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
        formatter
            .debug_struct("ClusterAgentConfig")
            .field("agent", &self.agent)
            .field("cluster", &self.cluster)
            .field("dcp", &"<DcpServerConfig>")
            .field("sessions", &self.sessions)
            .finish()
    }
}

/// Entrypoint for starting membership, the local agent, DCP, and node sessions
/// together.
pub struct ClusterAgent;

impl ClusterAgent {
    pub async fn start(
        config: ClusterAgentConfig,
        factories: DcpJobFactories,
    ) -> ClusterAgentResult<ClusterAgentHandle> {
        let mut config = config;
        validate_cluster_agent_config(&config)?;
        ensure_agent_role(&mut config.cluster, &config.sessions.agent_role);
        config.dcp.node_id = config.cluster.node_id.clone();

        let agent = Agent::start_with_config(config.agent.clone())?;
        let server = DcpServer::from_agent(&agent, factories.clone(), config.dcp.clone());
        let server_handle = server.start().await?;
        let agent_addr = advertised_agent_addr(&config.sessions.transport, &server_handle)?;
        config.cluster.agent_addr = Some(agent_addr);

        let cluster = ClusterNode::start(config.cluster.clone()).await?;
        let sessions = NodeSessionManagerHandle::start(
            config.sessions.clone(),
            cluster.node_id().to_owned(),
            cluster.state(),
        )?;
        let placement = PlacementCoordinatorHandle::start(
            cluster.node_id().to_owned(),
            config.sessions.agent_role.clone(),
            cluster.state(),
            agent.registry().clone(),
            sessions.clone(),
            factories.clone(),
            config.sessions.request_timeout,
        )?;
        let provider = Arc::new(ClusterView {
            registry: agent.registry().clone(),
            state: cluster.state(),
            sessions: sessions.clone(),
            placement: placement.clone(),
            self_node: cluster.node_id().to_owned(),
            agent_role: config.sessions.agent_role.clone(),
        });
        server.set_cluster_view(provider);

        Ok(ClusterAgentHandle {
            agent,
            cluster,
            server,
            server_handle: Some(server_handle),
            sessions,
            placement,
        })
    }
}

/// Running cluster-aware agent.
pub struct ClusterAgentHandle {
    agent: AgentHandle,
    cluster: ClusterNode,
    server: DcpServer,
    server_handle: Option<DcpServerHandle>,
    sessions: NodeSessionManagerHandle,
    placement: PlacementCoordinatorHandle,
}

impl ClusterAgentHandle {
    #[must_use]
    pub fn agent(&self) -> &AgentHandle {
        &self.agent
    }

    #[must_use]
    pub fn cluster(&self) -> &ClusterNode {
        &self.cluster
    }

    #[must_use]
    pub fn sessions(&self) -> &NodeSessionManagerHandle {
        &self.sessions
    }

    #[must_use]
    pub fn server(&self) -> &DcpServer {
        &self.server
    }

    #[must_use]
    pub fn tcp_addr(&self) -> Option<SocketAddr> {
        self.server_handle
            .as_ref()
            .and_then(DcpServerHandle::tcp_addr)
    }

    #[must_use]
    pub fn quic_addr(&self) -> Option<SocketAddr> {
        self.server_handle
            .as_ref()
            .and_then(DcpServerHandle::quic_addr)
    }

    pub async fn shutdown(mut self) -> ClusterAgentResult<()> {
        let node_id = self.cluster.node_id().to_owned();
        self.server.clear_cluster_view();
        self.server.clear_sharding_view();
        self.placement.shutdown().await;
        eprintln!(
            "datum-agent INFO shutdown_component_stopped node_id={node_id} component=placement"
        );
        self.sessions.shutdown().await;
        eprintln!(
            "datum-agent INFO shutdown_component_stopped node_id={node_id} component=sessions"
        );
        let _ = self.cluster.abort().await;
        eprintln!(
            "datum-agent INFO shutdown_component_stopped node_id={node_id} component=cluster"
        );
        if let Some(handle) = self.server_handle.take() {
            handle.shutdown().await;
        }
        eprintln!(
            "datum-agent INFO shutdown_component_stopped node_id={node_id} component=dcp_server"
        );
        self.agent.registry().shutdown()?;
        eprintln!(
            "datum-agent INFO shutdown_component_stopped node_id={node_id} component=registry"
        );
        eprintln!("datum-agent INFO shutdown_done node_id={node_id}");
        Ok(())
    }
}

impl Drop for ClusterAgentHandle {
    fn drop(&mut self) {
        self.server.clear_cluster_view();
        self.server.clear_sharding_view();
    }
}

#[derive(Clone)]
pub struct NodeSessionManagerHandle {
    inner: Arc<NodeSessionManagerInner>,
}

struct NodeSessionManagerInner {
    config: NodeSessionConfig,
    self_node: String,
    state: Signal<ClusterState>,
    sessions: tokio::sync::Mutex<BTreeMap<String, PeerSession>>,
    tasks: Mutex<Vec<JoinHandle<()>>>,
}

struct PeerSession {
    member: Member,
    commands: mpsc::Sender<SessionCommand>,
    pipe_commands: mpsc::Sender<ShardPipeCommand>,
    stop: watch::Sender<bool>,
    state: watch::Receiver<PeerSessionState>,
    task: JoinHandle<()>,
}

#[derive(Debug, Clone, PartialEq, Eq)]
enum PeerSessionState {
    Connecting,
    Connected,
    BackingOff(String),
    Closed,
}

impl PeerSessionState {
    fn as_text(&self) -> String {
        match self {
            Self::Connecting => "connecting".to_owned(),
            Self::Connected => "connected".to_owned(),
            Self::BackingOff(error) => format!("backing_off:{error}"),
            Self::Closed => "closed".to_owned(),
        }
    }
}

enum SessionCommand {
    ListJobs {
        timeout: Duration,
        reply: oneshot::Sender<Result<Vec<crate::dcp::proto::JobStatus>, String>>,
    },
    SubmitClusterJob {
        request: SubmitClusterJob,
        timeout: Duration,
        reply: oneshot::Sender<Result<crate::dcp::proto::JobStatus, String>>,
    },
    StartClusterJob {
        factory_name: String,
        instance_name: String,
        params: HashMap<String, String>,
        assignment: ClusterJobStart,
        timeout: Duration,
        reply: oneshot::Sender<Result<crate::dcp::proto::JobStatus, String>>,
    },
    ClusterJobStatus {
        name: String,
        cluster: bool,
        timeout: Duration,
        reply: oneshot::Sender<Result<crate::dcp::proto::JobStatus, String>>,
    },
    DrainClusterJob {
        name: String,
        cluster: bool,
        timeout: Duration,
        reply: oneshot::Sender<Result<crate::dcp::proto::JobStatus, String>>,
    },
    StopClusterJob {
        name: String,
        cluster: bool,
        timeout: Duration,
        reply: oneshot::Sender<Result<crate::dcp::proto::JobStatus, String>>,
    },
    ListClusterJobs {
        timeout: Duration,
        reply: oneshot::Sender<Result<ClusterJobList, String>>,
    },
    SubscribeMetrics {
        interval: Duration,
        job_names: Vec<String>,
        timeout: Duration,
        reply: oneshot::Sender<Result<MetricSubscription, String>>,
    },
    RememberClusterAssignment {
        instance_name: String,
        assignment: Option<ClusterJobStart>,
        timeout: Duration,
        reply: oneshot::Sender<Result<(), String>>,
    },
    AllocateShard {
        type_name: String,
        shard_id: String,
        timeout: Duration,
        reply: oneshot::Sender<Result<ShardAllocation, String>>,
    },
    RememberShardAllocations {
        table: ShardAllocationTable,
        timeout: Duration,
        reply: oneshot::Sender<Result<(), String>>,
    },
    GetShardAllocations {
        type_name: String,
        timeout: Duration,
        reply: oneshot::Sender<Result<ShardAllocationTable, String>>,
    },
    ForwardShardEnvelopes {
        batch: ForwardShardEnvelopes,
        timeout: Duration,
        reply: oneshot::Sender<Result<ShardEnvelopeBatchResult, String>>,
    },
    CompleteShardingAsk {
        response: CompleteShardingAsk,
        timeout: Duration,
        reply: oneshot::Sender<Result<(), String>>,
    },
}

enum ShardPipeCommand {
    Forward { batch: ForwardShardEnvelopes },
    Reply { response: CompleteShardingAsk },
}

impl NodeSessionManagerHandle {
    fn start(
        config: NodeSessionConfig,
        self_node: String,
        state: Signal<ClusterState>,
    ) -> ClusterAgentResult<Self> {
        let inner = Arc::new(NodeSessionManagerInner {
            config,
            self_node,
            state,
            sessions: tokio::sync::Mutex::new(BTreeMap::new()),
            tasks: Mutex::new(Vec::new()),
        });

        let manager_task = {
            let inner = Arc::clone(&inner);
            tokio::spawn(async move {
                run_node_session_manager(inner).await;
            })
        };
        {
            let mut tasks = inner.tasks.lock().expect("node-session tasks poisoned");
            tasks.push(manager_task);
        }

        Ok(Self { inner })
    }

    pub async fn shutdown(&self) {
        let sessions = {
            let mut locked = self.inner.sessions.lock().await;
            std::mem::take(&mut *locked)
        };
        for (_, session) in sessions {
            let _ = session.stop.send(true);
            session.task.abort();
        }
        let tasks = {
            let mut locked = self
                .inner
                .tasks
                .lock()
                .expect("node-session tasks poisoned");
            std::mem::take(&mut *locked)
        };
        for task in tasks {
            task.abort();
        }
    }

    async fn list_jobs(
        &self,
        node_id: &str,
        timeout: Duration,
    ) -> Result<Vec<crate::dcp::proto::JobStatus>, String> {
        let sender = {
            let locked = self.inner.sessions.lock().await;
            locked.get(node_id).map(|session| session.commands.clone())
        }
        .ok_or_else(|| "node session unavailable".to_owned())?;
        let (reply, receiver) = oneshot::channel();
        sender
            .try_send(SessionCommand::ListJobs { timeout, reply })
            .map_err(|error| format!("node session command queue unavailable: {error}"))?;
        tokio::time::timeout(timeout, receiver)
            .await
            .map_err(|_| "node session request timed out".to_owned())?
            .map_err(|_| "node session closed".to_owned())?
    }

    async fn submit_cluster_job(
        &self,
        node_id: &str,
        request: SubmitClusterJob,
        timeout: Duration,
    ) -> Result<crate::dcp::proto::JobStatus, String> {
        self.request_status(node_id, timeout, |reply| SessionCommand::SubmitClusterJob {
            request,
            timeout,
            reply,
        })
        .await
    }

    async fn start_cluster_job(
        &self,
        node_id: &str,
        factory_name: String,
        instance_name: String,
        params: HashMap<String, String>,
        assignment: ClusterJobStart,
        timeout: Duration,
    ) -> Result<crate::dcp::proto::JobStatus, String> {
        self.request_status(node_id, timeout, |reply| SessionCommand::StartClusterJob {
            factory_name,
            instance_name,
            params,
            assignment,
            timeout,
            reply,
        })
        .await
    }

    async fn cluster_job_status(
        &self,
        node_id: &str,
        name: String,
        timeout: Duration,
    ) -> Result<crate::dcp::proto::JobStatus, String> {
        self.request_status(node_id, timeout, |reply| SessionCommand::ClusterJobStatus {
            name,
            cluster: false,
            timeout,
            reply,
        })
        .await
    }

    async fn drain_cluster_job(
        &self,
        node_id: &str,
        name: String,
        timeout: Duration,
    ) -> Result<crate::dcp::proto::JobStatus, String> {
        self.request_status(node_id, timeout, |reply| SessionCommand::DrainClusterJob {
            name,
            cluster: false,
            timeout,
            reply,
        })
        .await
    }

    async fn stop_cluster_job(
        &self,
        node_id: &str,
        name: String,
        timeout: Duration,
    ) -> Result<crate::dcp::proto::JobStatus, String> {
        self.request_status(node_id, timeout, |reply| SessionCommand::StopClusterJob {
            name,
            cluster: false,
            timeout,
            reply,
        })
        .await
    }

    async fn forward_cluster_job_status(
        &self,
        node_id: &str,
        name: String,
        timeout: Duration,
    ) -> Result<crate::dcp::proto::JobStatus, String> {
        self.request_status(node_id, timeout, |reply| SessionCommand::ClusterJobStatus {
            name,
            cluster: true,
            timeout,
            reply,
        })
        .await
    }

    async fn forward_drain_cluster_job(
        &self,
        node_id: &str,
        name: String,
        timeout: Duration,
    ) -> Result<crate::dcp::proto::JobStatus, String> {
        self.request_status(node_id, timeout, |reply| SessionCommand::DrainClusterJob {
            name,
            cluster: true,
            timeout,
            reply,
        })
        .await
    }

    async fn forward_stop_cluster_job(
        &self,
        node_id: &str,
        name: String,
        timeout: Duration,
    ) -> Result<crate::dcp::proto::JobStatus, String> {
        self.request_status(node_id, timeout, |reply| SessionCommand::StopClusterJob {
            name,
            cluster: true,
            timeout,
            reply,
        })
        .await
    }

    async fn list_cluster_jobs(
        &self,
        node_id: &str,
        timeout: Duration,
    ) -> Result<ClusterJobList, String> {
        let sender = self.session_sender(node_id).await?;
        let (reply, receiver) = oneshot::channel();
        sender
            .try_send(SessionCommand::ListClusterJobs { timeout, reply })
            .map_err(|error| format!("node session command queue unavailable: {error}"))?;
        wait_session_reply(timeout, receiver).await
    }

    async fn subscribe_metrics(
        &self,
        node_id: &str,
        interval: Duration,
        job_names: Vec<String>,
        timeout: Duration,
    ) -> Result<MetricSubscription, String> {
        let sender = self.session_sender(node_id).await?;
        let (reply, receiver) = oneshot::channel();
        sender
            .try_send(SessionCommand::SubscribeMetrics {
                interval,
                job_names,
                timeout,
                reply,
            })
            .map_err(|error| format!("node session command queue unavailable: {error}"))?;
        wait_session_reply(timeout, receiver).await
    }

    async fn remember_cluster_assignment(
        &self,
        node_id: &str,
        instance_name: String,
        assignment: Option<ClusterJobStart>,
        timeout: Duration,
    ) -> Result<(), String> {
        let sender = self.session_sender(node_id).await?;
        let (reply, receiver) = oneshot::channel();
        sender
            .try_send(SessionCommand::RememberClusterAssignment {
                instance_name,
                assignment,
                timeout,
                reply,
            })
            .map_err(|error| format!("node session command queue unavailable: {error}"))?;
        wait_session_reply(timeout, receiver).await
    }

    pub async fn allocate_shard(
        &self,
        node_id: &str,
        type_name: String,
        shard_id: String,
        timeout: Duration,
    ) -> Result<ShardAllocation, String> {
        let sender = self.session_sender(node_id).await?;
        let (reply, receiver) = oneshot::channel();
        sender
            .try_send(SessionCommand::AllocateShard {
                type_name,
                shard_id,
                timeout,
                reply,
            })
            .map_err(|error| format!("node session command queue unavailable: {error}"))?;
        wait_session_reply(timeout, receiver).await
    }

    pub async fn remember_shard_allocations(
        &self,
        node_id: &str,
        table: ShardAllocationTable,
        timeout: Duration,
    ) -> Result<(), String> {
        let sender = self.session_sender(node_id).await?;
        let (reply, receiver) = oneshot::channel();
        sender
            .try_send(SessionCommand::RememberShardAllocations {
                table,
                timeout,
                reply,
            })
            .map_err(|error| format!("node session command queue unavailable: {error}"))?;
        wait_session_reply(timeout, receiver).await
    }

    pub async fn get_shard_allocations(
        &self,
        node_id: &str,
        type_name: String,
        timeout: Duration,
    ) -> Result<ShardAllocationTable, String> {
        let sender = self.session_sender(node_id).await?;
        let (reply, receiver) = oneshot::channel();
        sender
            .try_send(SessionCommand::GetShardAllocations {
                type_name,
                timeout,
                reply,
            })
            .map_err(|error| format!("node session command queue unavailable: {error}"))?;
        wait_session_reply(timeout, receiver).await
    }

    pub async fn forward_shard_envelopes(
        &self,
        node_id: &str,
        batch: ForwardShardEnvelopes,
        timeout: Duration,
    ) -> Result<ShardEnvelopeBatchResult, String> {
        let sender = self.session_sender(node_id).await?;
        let (reply, receiver) = oneshot::channel();
        sender
            .try_send(SessionCommand::ForwardShardEnvelopes {
                batch,
                timeout,
                reply,
            })
            .map_err(|error| format!("node session command queue unavailable: {error}"))?;
        wait_session_reply(timeout, receiver).await
    }

    pub async fn complete_sharding_ask(
        &self,
        node_id: &str,
        response: CompleteShardingAsk,
        timeout: Duration,
    ) -> Result<(), String> {
        let sender = self.session_sender(node_id).await?;
        let (reply, receiver) = oneshot::channel();
        sender
            .try_send(SessionCommand::CompleteShardingAsk {
                response,
                timeout,
                reply,
            })
            .map_err(|error| format!("node session command queue unavailable: {error}"))?;
        wait_session_reply(timeout, receiver).await
    }

    pub async fn forward_shard_pipe_envelopes(
        &self,
        node_id: &str,
        batch: ForwardShardEnvelopes,
    ) -> Result<(), String> {
        let sender = self.pipe_sender(node_id).await?;
        sender
            .try_send(ShardPipeCommand::Forward { batch })
            .map_err(|error| format!("node session shard pipe queue unavailable: {error}"))
    }

    pub async fn complete_sharding_ask_pipe(
        &self,
        node_id: &str,
        response: CompleteShardingAsk,
    ) -> Result<(), String> {
        let sender = self.pipe_sender(node_id).await?;
        sender
            .try_send(ShardPipeCommand::Reply { response })
            .map_err(|error| format!("node session shard pipe queue unavailable: {error}"))
    }

    pub fn try_complete_sharding_ask_pipe(
        &self,
        node_id: &str,
        response: CompleteShardingAsk,
    ) -> Result<(), String> {
        let locked = self
            .inner
            .sessions
            .try_lock()
            .map_err(|_| "node session table busy".to_owned())?;
        let sender = locked
            .get(node_id)
            .map(|session| session.pipe_commands.clone())
            .ok_or_else(|| "node session unavailable".to_owned())?;
        sender
            .try_send(ShardPipeCommand::Reply { response })
            .map_err(|error| format!("node session shard pipe queue unavailable: {error}"))
    }

    async fn request_status<F>(
        &self,
        node_id: &str,
        timeout: Duration,
        make_command: F,
    ) -> Result<crate::dcp::proto::JobStatus, String>
    where
        F: FnOnce(oneshot::Sender<Result<crate::dcp::proto::JobStatus, String>>) -> SessionCommand,
    {
        let sender = self.session_sender(node_id).await?;
        let (reply, receiver) = oneshot::channel();
        sender
            .try_send(make_command(reply))
            .map_err(|error| format!("node session command queue unavailable: {error}"))?;
        wait_session_reply(timeout, receiver).await
    }

    async fn session_sender(&self, node_id: &str) -> Result<mpsc::Sender<SessionCommand>, String> {
        let locked = self.inner.sessions.lock().await;
        locked
            .get(node_id)
            .map(|session| session.commands.clone())
            .ok_or_else(|| "node session unavailable".to_owned())
    }

    async fn pipe_sender(&self, node_id: &str) -> Result<mpsc::Sender<ShardPipeCommand>, String> {
        let locked = self.inner.sessions.lock().await;
        locked
            .get(node_id)
            .map(|session| session.pipe_commands.clone())
            .ok_or_else(|| "node session unavailable".to_owned())
    }

    async fn session_states(&self) -> BTreeMap<String, String> {
        let locked = self.inner.sessions.lock().await;
        locked
            .iter()
            .map(|(node_id, session)| (node_id.clone(), session.state.borrow().clone().as_text()))
            .collect()
    }
}

async fn wait_session_reply<T>(
    timeout: Duration,
    receiver: oneshot::Receiver<Result<T, String>>,
) -> Result<T, String> {
    tokio::time::timeout(timeout, receiver)
        .await
        .map_err(|_| "node session request timed out".to_owned())?
        .map_err(|_| "node session closed".to_owned())?
}

async fn run_node_session_manager(inner: Arc<NodeSessionManagerInner>) {
    let interval = inner
        .config
        .reconnect_min_backoff
        .min(Duration::from_millis(100))
        .max(Duration::from_millis(10));
    loop {
        reconcile_sessions(&inner).await;
        tokio::time::sleep(interval).await;
    }
}

async fn reconcile_sessions(inner: &Arc<NodeSessionManagerInner>) {
    let state = inner.state.get();
    let mut wanted = BTreeSet::new();
    for member in state.members.values() {
        if eligible_member(inner, member) {
            wanted.insert(member.node_id.clone());
            ensure_session(inner, member.clone()).await;
        }
    }
    close_unwanted_sessions(inner, &wanted).await;
}

fn eligible_member(inner: &NodeSessionManagerInner, member: &Member) -> bool {
    member.node_id != inner.self_node
        && member.state == MemberState::Up
        && member.has_role(&inner.config.agent_role)
        && member.agent_addr.is_some()
}

async fn ensure_session(inner: &Arc<NodeSessionManagerInner>, member: Member) {
    let mut locked = inner.sessions.lock().await;
    if let Some(existing) = locked.get(&member.node_id)
        && existing.member.agent_addr == member.agent_addr
        && existing.member.incarnation == member.incarnation
    {
        return;
    }

    if let Some(existing) = locked.remove(&member.node_id) {
        let _ = existing.stop.send(true);
        existing.task.abort();
    }

    let (commands, receiver) = mpsc::channel(inner.config.command_buffer.max(1));
    let pipe_buffer = inner.config.command_buffer.saturating_mul(4).max(1);
    let (pipe_commands, pipe_receiver) = mpsc::channel(pipe_buffer);
    let (stop, stop_receiver) = watch::channel(false);
    let (state_sender, state_receiver) = watch::channel(PeerSessionState::Connecting);
    let config = inner.config.clone();
    let local_node = inner.self_node.clone();
    let task_member = member.clone();
    let task = tokio::spawn(async move {
        run_peer_session(
            config,
            local_node,
            task_member,
            receiver,
            pipe_receiver,
            stop_receiver,
            state_sender,
        )
        .await;
    });

    locked.insert(
        member.node_id.clone(),
        PeerSession {
            member,
            commands,
            pipe_commands,
            stop,
            state: state_receiver,
            task,
        },
    );
}

async fn close_unwanted_sessions(inner: &Arc<NodeSessionManagerInner>, wanted: &BTreeSet<String>) {
    let to_close = {
        let locked = inner.sessions.lock().await;
        locked
            .keys()
            .filter(|node_id| !wanted.contains(*node_id))
            .cloned()
            .collect::<Vec<_>>()
    };
    for node_id in to_close {
        close_session(inner, &node_id).await;
    }
}

async fn close_session(inner: &Arc<NodeSessionManagerInner>, node_id: &str) {
    let removed = {
        let mut locked = inner.sessions.lock().await;
        locked.remove(node_id)
    };
    if let Some(session) = removed {
        let _ = session.stop.send(true);
        session.task.abort();
    }
}

async fn run_peer_session(
    config: NodeSessionConfig,
    local_node: String,
    member: Member,
    mut commands: mpsc::Receiver<SessionCommand>,
    mut pipe_commands: mpsc::Receiver<ShardPipeCommand>,
    mut stop: watch::Receiver<bool>,
    state: watch::Sender<PeerSessionState>,
) {
    let mut backoff = config.reconnect_min_backoff;
    loop {
        if *stop.borrow() {
            break;
        }
        let _ = state.send(PeerSessionState::Connecting);
        match connect_peer(&config.transport, &local_node, &member).await {
            Ok(client) => match connect_shard_pipe(&config.transport, &local_node, &member).await {
                Ok(pipe) => {
                    let _ = state.send(PeerSessionState::Connected);
                    eprintln!(
                        "datum-agent INFO peer_session_connected node_id={local_node} peer={}",
                        member.node_id
                    );
                    backoff = config.reconnect_min_backoff;
                    if !run_connected_peer(
                        client,
                        pipe,
                        &mut commands,
                        &mut pipe_commands,
                        &mut stop,
                    )
                    .await
                    {
                        break;
                    }
                    eprintln!(
                        "datum-agent WARN peer_session_lost node_id={local_node} peer={}",
                        member.node_id
                    );
                }
                Err(error) => {
                    log_peer_connect_failed(&local_node, &member.node_id, &error);
                    let _ = state.send(PeerSessionState::BackingOff(error.clone()));
                    if !backoff_or_stop(
                        backoff,
                        error,
                        &mut commands,
                        &mut pipe_commands,
                        &mut stop,
                    )
                    .await
                    {
                        break;
                    }
                    backoff = next_backoff(backoff, config.reconnect_max_backoff);
                    log_peer_reconnecting(&local_node, &member.node_id, backoff);
                }
            },
            Err(error) => {
                log_peer_connect_failed(&local_node, &member.node_id, &error);
                let _ = state.send(PeerSessionState::BackingOff(error.clone()));
                if !backoff_or_stop(backoff, error, &mut commands, &mut pipe_commands, &mut stop)
                    .await
                {
                    break;
                }
                backoff = next_backoff(backoff, config.reconnect_max_backoff);
                log_peer_reconnecting(&local_node, &member.node_id, backoff);
            }
        }
    }
    let _ = state.send(PeerSessionState::Closed);
    eprintln!(
        "datum-agent INFO peer_session_closed node_id={local_node} peer={}",
        member.node_id
    );
}

fn log_peer_connect_failed(local_node: &str, peer: &str, error: &str) {
    eprintln!(
        "datum-agent WARN peer_session_connect_failed node_id={local_node} peer={peer} error={error}"
    );
}

fn log_peer_reconnecting(local_node: &str, peer: &str, backoff: Duration) {
    eprintln!(
        "datum-agent INFO peer_session_reconnecting node_id={local_node} peer={peer} backoff_ms={}",
        backoff.as_millis()
    );
}

async fn connect_peer(
    transport: &NodeSessionTransport,
    local_node: &str,
    member: &Member,
) -> Result<DcpClient, String> {
    let addr = member
        .agent_addr
        .ok_or_else(|| "member did not advertise a DCP agent address".to_owned())?;
    let mut hello = Hello::new(local_node.to_owned(), ClientKind::ClusterNode);
    hello.capabilities.push("cluster-node-sessions".to_owned());
    match transport {
        NodeSessionTransport::TcpLoopback => {
            if !addr.ip().is_loopback() {
                return Err(format!(
                    "plaintext node-session TCP requires loopback, got {addr}"
                ));
            }
            DcpClient::connect_tcp(addr, hello)
                .await
                .map_err(|error| error.to_string())
        }
        NodeSessionTransport::QuicMtls {
            server_name,
            client_config,
        } => DcpClient::connect_quic(addr, server_name, client_config.clone(), hello)
            .await
            .map_err(|error| error.to_string()),
    }
}

async fn connect_shard_pipe(
    transport: &NodeSessionTransport,
    local_node: &str,
    member: &Member,
) -> Result<ShardPipeClient, String> {
    let addr = member
        .agent_addr
        .ok_or_else(|| "member did not advertise a DCP agent address".to_owned())?;
    let mut hello = Hello::new(local_node.to_owned(), ClientKind::ClusterNode);
    hello.capabilities.push("cluster-node-sessions".to_owned());
    hello.capabilities.push("cluster-sharding-pipe".to_owned());
    match transport {
        NodeSessionTransport::TcpLoopback => {
            if !addr.ip().is_loopback() {
                return Err(format!(
                    "plaintext node-session TCP requires loopback, got {addr}"
                ));
            }
            ShardPipeClient::connect_tcp(addr, hello)
                .await
                .map_err(|error| error.to_string())
        }
        NodeSessionTransport::QuicMtls {
            server_name,
            client_config,
        } => ShardPipeClient::connect_quic(addr, server_name, client_config.clone(), hello)
            .await
            .map_err(|error| error.to_string()),
    }
}

async fn run_connected_peer(
    client: DcpClient,
    pipe: ShardPipeClient,
    commands: &mut mpsc::Receiver<SessionCommand>,
    pipe_commands: &mut mpsc::Receiver<ShardPipeCommand>,
    stop: &mut watch::Receiver<bool>,
) -> bool {
    loop {
        tokio::select! {
            changed = stop.changed() => {
                return changed.is_ok() && !*stop.borrow();
            }
            command = pipe_commands.recv() => {
                let Some(command) = command else {
                    return false;
                };
                let frame = build_shard_pipe_frame(command, pipe_commands);
                if pipe.send(frame).await.is_err() {
                    return true;
                }
            }
            command = commands.recv() => {
                let Some(command) = command else {
                    return false;
                };
                match command {
                    SessionCommand::ListJobs { timeout, reply } => {
                        let result = tokio::time::timeout(timeout, client.list_jobs()).await
                            .map_err(|_| "peer ListJobs timed out".to_owned())
                            .and_then(|result| result.map_err(|error| error.to_string()));
                        let reconnect = result.is_err();
                        let _ = reply.send(result);
                        if reconnect {
                            return true;
                        }
                    }
                    SessionCommand::SubmitClusterJob {
                        request,
                        timeout,
                        reply,
                    } => {
                        let timeout_ms = duration_millis_u64(timeout);
                        let result = tokio::time::timeout(
                            timeout,
                            client.submit_cluster_job(
                                request.factory_name,
                                request.instance_name,
                                request.params,
                                request.placement.unwrap_or_else(default_wire_placement),
                                timeout_ms,
                            ),
                        )
                        .await
                        .map_err(|_| "peer SubmitClusterJob timed out".to_owned())
                        .and_then(|result| result.map_err(|error| error.to_string()));
                        let reconnect = result.is_err();
                        let _ = reply.send(result);
                        if reconnect {
                            return true;
                        }
                    }
                    SessionCommand::StartClusterJob {
                        factory_name,
                        instance_name,
                        params,
                        assignment,
                        timeout,
                        reply,
                    } => {
                        let result = tokio::time::timeout(
                            timeout,
                            client.start_cluster_job_on_node(
                                factory_name,
                                instance_name,
                                params,
                                assignment,
                            ),
                        )
                        .await
                        .map_err(|_| "peer StartJob timed out".to_owned())
                        .and_then(|result| result.map_err(|error| error.to_string()));
                        let reconnect = result.is_err();
                        let _ = reply.send(result);
                        if reconnect {
                            return true;
                        }
                    }
                    SessionCommand::ClusterJobStatus {
                        name,
                        cluster,
                        timeout,
                        reply,
                    } => {
                        let request = async {
                            if cluster {
                                client
                                    .cluster_job_status(name, duration_millis_u64(timeout))
                                    .await
                            } else {
                                client.job_status(name).await
                            }
                        };
                        let result = tokio::time::timeout(timeout, request).await
                        .map_err(|_| "peer JobStatus timed out".to_owned())
                        .and_then(|result| result.map_err(|error| error.to_string()));
                        let reconnect = result.is_err();
                        let _ = reply.send(result);
                        if reconnect {
                            return true;
                        }
                    }
                    SessionCommand::DrainClusterJob {
                        name,
                        cluster,
                        timeout,
                        reply,
                    } => {
                        let request = async {
                            if cluster {
                                client
                                    .drain_cluster_job(name, duration_millis_u64(timeout))
                                    .await
                            } else {
                                client.drain_job(name).await
                            }
                        };
                        let result = tokio::time::timeout(timeout, request).await
                        .map_err(|_| "peer DrainJob timed out".to_owned())
                        .and_then(|result| result.map_err(|error| error.to_string()));
                        let reconnect = result.is_err();
                        let _ = reply.send(result);
                        if reconnect {
                            return true;
                        }
                    }
                    SessionCommand::StopClusterJob {
                        name,
                        cluster,
                        timeout,
                        reply,
                    } => {
                        let request = async {
                            if cluster {
                                client
                                    .stop_cluster_job(name, duration_millis_u64(timeout))
                                    .await
                            } else {
                                client.stop_job(name).await
                            }
                        };
                        let result = tokio::time::timeout(timeout, request).await
                        .map_err(|_| "peer StopJob timed out".to_owned())
                        .and_then(|result| result.map_err(|error| error.to_string()));
                        let reconnect = result.is_err();
                        let _ = reply.send(result);
                        if reconnect {
                            return true;
                        }
                    }
                    SessionCommand::ListClusterJobs { timeout, reply } => {
                        let result = tokio::time::timeout(
                            timeout,
                            client.list_cluster_jobs(duration_millis_u64(timeout)),
                        )
                        .await
                        .map_err(|_| "peer ListClusterJobs timed out".to_owned())
                        .and_then(|result| result.map_err(|error| error.to_string()));
                        let reconnect = result.is_err();
                        let _ = reply.send(result);
                        if reconnect {
                            return true;
                        }
                    }
                    SessionCommand::SubscribeMetrics {
                        interval,
                        job_names,
                        timeout,
                        reply,
                    } => {
                        let result = tokio::time::timeout(
                            timeout,
                            client.subscribe_local_metrics(
                                duration_millis_u64(interval),
                                job_names,
                            ),
                        )
                        .await
                        .map_err(|_| "peer SubscribeMetrics timed out".to_owned())
                        .and_then(|result| result.map_err(|error| error.to_string()));
                        let reconnect = result.is_err();
                        let _ = reply.send(result);
                        if reconnect {
                            return true;
                        }
                    }
                    SessionCommand::RememberClusterAssignment {
                        instance_name,
                        assignment,
                        timeout,
                        reply,
                    } => {
                        let result = tokio::time::timeout(
                            timeout,
                            client.remember_cluster_assignment(instance_name, assignment),
                        )
                        .await
                        .map_err(|_| "peer RememberClusterAssignment timed out".to_owned())
                        .and_then(|result| result.map_err(|error| error.to_string()));
                        let reconnect = result.is_err();
                        let _ = reply.send(result);
                        if reconnect {
                            return true;
                        }
                    }
                    SessionCommand::AllocateShard {
                        type_name,
                        shard_id,
                        timeout,
                        reply,
                    } => {
                        let timeout_ms = duration_millis_u64(timeout);
                        let result = tokio::time::timeout(
                            timeout,
                            client.allocate_shard(type_name, shard_id, timeout_ms),
                        )
                        .await
                        .map_err(|_| "peer AllocateShard timed out".to_owned())
                        .and_then(|result| result.map_err(|error| error.to_string()));
                        let reconnect = result.is_err();
                        let _ = reply.send(result);
                        if reconnect {
                            return true;
                        }
                    }
                    SessionCommand::RememberShardAllocations {
                        table,
                        timeout,
                        reply,
                    } => {
                        let result = tokio::time::timeout(
                            timeout,
                            client.remember_shard_allocations(table),
                        )
                        .await
                        .map_err(|_| "peer RememberShardAllocations timed out".to_owned())
                        .and_then(|result| result.map_err(|error| error.to_string()));
                        let reconnect = result.is_err();
                        let _ = reply.send(result);
                        if reconnect {
                            return true;
                        }
                    }
                    SessionCommand::GetShardAllocations {
                        type_name,
                        timeout,
                        reply,
                    } => {
                        let timeout_ms = duration_millis_u64(timeout);
                        let result = tokio::time::timeout(
                            timeout,
                            client.get_shard_allocations(type_name, timeout_ms),
                        )
                        .await
                        .map_err(|_| "peer GetShardAllocations timed out".to_owned())
                        .and_then(|result| result.map_err(|error| error.to_string()));
                        let reconnect = result.is_err();
                        let _ = reply.send(result);
                        if reconnect {
                            return true;
                        }
                    }
                    SessionCommand::ForwardShardEnvelopes {
                        batch,
                        timeout,
                        reply,
                    } => {
                        let timeout_ms = duration_millis_u64(timeout);
                        let result = tokio::time::timeout(
                            timeout,
                            client.forward_shard_envelopes(batch, timeout_ms),
                        )
                        .await
                        .map_err(|_| "peer ForwardShardEnvelopes timed out".to_owned())
                        .and_then(|result| result.map_err(|error| error.to_string()));
                        let reconnect = result.is_err();
                        let _ = reply.send(result);
                        if reconnect {
                            return true;
                        }
                    }
                    SessionCommand::CompleteShardingAsk {
                        response,
                        timeout,
                        reply,
                    } => {
                        let result =
                            tokio::time::timeout(timeout, client.complete_sharding_ask(response))
                                .await
                                .map_err(|_| "peer CompleteShardingAsk timed out".to_owned())
                                .and_then(|result| result.map_err(|error| error.to_string()));
                        let reconnect = result.is_err();
                        let _ = reply.send(result);
                        if reconnect {
                            return true;
                        }
                    }
                }
            }
        }
    }
}

const SHARD_PIPE_MAX_BATCH_COMMANDS: usize = 1024;

fn build_shard_pipe_frame(
    first: ShardPipeCommand,
    pipe_commands: &mut mpsc::Receiver<ShardPipeCommand>,
) -> ShardPipeFrame {
    let mut frame = ShardPipeFrame {
        forwards: Vec::new(),
        replies: Vec::new(),
    };
    push_shard_pipe_command(&mut frame, first);
    for _ in 1..SHARD_PIPE_MAX_BATCH_COMMANDS {
        match pipe_commands.try_recv() {
            Ok(command) => push_shard_pipe_command(&mut frame, command),
            Err(_) => break,
        }
    }
    frame
}

fn push_shard_pipe_command(frame: &mut ShardPipeFrame, command: ShardPipeCommand) {
    match command {
        ShardPipeCommand::Forward { batch } => {
            if let Some(existing) = frame
                .forwards
                .iter_mut()
                .find(|existing| existing.type_name == batch.type_name)
            {
                existing.envelopes.extend(batch.envelopes);
            } else {
                frame.forwards.push(batch);
            }
        }
        ShardPipeCommand::Reply { response } => {
            frame.replies.push(response);
        }
    }
}

async fn backoff_or_stop(
    backoff: Duration,
    error: String,
    commands: &mut mpsc::Receiver<SessionCommand>,
    pipe_commands: &mut mpsc::Receiver<ShardPipeCommand>,
    stop: &mut watch::Receiver<bool>,
) -> bool {
    let sleep = sleep(backoff);
    tokio::pin!(sleep);
    loop {
        tokio::select! {
            () = &mut sleep => return true,
            changed = stop.changed() => {
                return changed.is_ok() && !*stop.borrow();
            }
            command = pipe_commands.recv() => {
                if command.is_none() {
                    return false;
                }
            }
            command = commands.recv() => {
                let Some(command) = command else {
                    return false;
                };
                match command {
                    SessionCommand::ListJobs { reply, .. } => {
                        let _ = reply.send(Err(format!("node session unreachable: {error}")));
                    }
                    SessionCommand::SubmitClusterJob { reply, .. }
                    | SessionCommand::StartClusterJob { reply, .. }
                    | SessionCommand::ClusterJobStatus { reply, .. }
                    | SessionCommand::DrainClusterJob { reply, .. }
                    | SessionCommand::StopClusterJob { reply, .. } => {
                        let _ = reply.send(Err(format!("node session unreachable: {error}")));
                    }
                    SessionCommand::ListClusterJobs { reply, .. } => {
                        let _ = reply.send(Err(format!("node session unreachable: {error}")));
                    }
                    SessionCommand::SubscribeMetrics { reply, .. } => {
                        let _ = reply.send(Err(format!("node session unreachable: {error}")));
                    }
                    SessionCommand::RememberClusterAssignment { reply, .. } => {
                        let _ = reply.send(Err(format!("node session unreachable: {error}")));
                    }
                    SessionCommand::AllocateShard { reply, .. } => {
                        let _ = reply.send(Err(format!("node session unreachable: {error}")));
                    }
                    SessionCommand::RememberShardAllocations { reply, .. } => {
                        let _ = reply.send(Err(format!("node session unreachable: {error}")));
                    }
                    SessionCommand::GetShardAllocations { reply, .. } => {
                        let _ = reply.send(Err(format!("node session unreachable: {error}")));
                    }
                    SessionCommand::ForwardShardEnvelopes { reply, .. } => {
                        let _ = reply.send(Err(format!("node session unreachable: {error}")));
                    }
                    SessionCommand::CompleteShardingAsk { reply, .. } => {
                        let _ = reply.send(Err(format!("node session unreachable: {error}")));
                    }
                }
            }
        }
    }
}

fn default_wire_placement() -> PlacementSpec {
    PlacementSpec {
        role_constraint: String::new(),
        strategy: PlacementStrategy::LeastJobs as i32,
        pinned_node_id: String::new(),
    }
}

fn duration_millis_u64(duration: Duration) -> u64 {
    duration.as_millis().min(u128::from(u64::MAX)) as u64
}

fn next_backoff(current: Duration, max: Duration) -> Duration {
    current.saturating_mul(2).min(max)
}

#[derive(Clone)]
struct PlacementCoordinatorHandle {
    commands: mpsc::Sender<PlacementCommand>,
    assignments: watch::Receiver<Arc<BTreeMap<String, String>>>,
    tasks: Arc<Mutex<Vec<JoinHandle<()>>>>,
}

enum PlacementCommand {
    Submit {
        request: SubmitClusterJob,
        timeout: Duration,
        reply: oneshot::Sender<Result<crate::dcp::proto::JobStatus, DcpError>>,
    },
    Status {
        name: String,
        timeout: Duration,
        reply: oneshot::Sender<Result<crate::dcp::proto::JobStatus, DcpError>>,
    },
    Drain {
        name: String,
        timeout: Duration,
        reply: oneshot::Sender<Result<crate::dcp::proto::JobStatus, DcpError>>,
    },
    Stop {
        name: String,
        timeout: Duration,
        reply: oneshot::Sender<Result<crate::dcp::proto::JobStatus, DcpError>>,
    },
    Remember {
        instance_name: String,
        assignment: Option<ClusterJobStart>,
        reply: oneshot::Sender<Result<(), DcpError>>,
    },
    RegisterRestarted {
        instance_name: String,
        assignment: ClusterJobStart,
        timeout: Duration,
        reply: oneshot::Sender<Result<(), DcpError>>,
    },
    Tick,
    Shutdown,
}

struct PlacementActor {
    self_node: String,
    agent_role: String,
    state: Signal<ClusterState>,
    registry: JobRegistryHandle,
    sessions: NodeSessionManagerHandle,
    factories: DcpJobFactories,
    request_timeout: Duration,
    active_coordinator: bool,
    last_known_coordinator: Option<String>,
    assignments: BTreeMap<String, ClusterJobMetadata>,
    assignment_updates: watch::Sender<Arc<BTreeMap<String, String>>>,
}

impl PlacementCoordinatorHandle {
    fn start(
        self_node: String,
        agent_role: String,
        state: Signal<ClusterState>,
        registry: JobRegistryHandle,
        sessions: NodeSessionManagerHandle,
        factories: DcpJobFactories,
        request_timeout: Duration,
    ) -> ClusterAgentResult<Self> {
        let (commands, receiver) = mpsc::channel(128);
        let (assignment_updates, assignments) = watch::channel(Arc::new(BTreeMap::new()));
        let tasks = Arc::new(Mutex::new(Vec::new()));
        let actor = PlacementActor {
            self_node,
            agent_role,
            state,
            registry,
            sessions,
            factories,
            request_timeout,
            active_coordinator: false,
            last_known_coordinator: None,
            assignments: BTreeMap::new(),
            assignment_updates,
        };
        let actor_task = tokio::spawn(run_placement_actor(actor, receiver));
        let tick_commands = commands.clone();
        let tick_task = tokio::spawn(async move {
            let mut interval = tokio::time::interval(Duration::from_millis(50));
            loop {
                interval.tick().await;
                if tick_commands.send(PlacementCommand::Tick).await.is_err() {
                    break;
                }
            }
        });
        {
            let mut locked = tasks.lock().expect("placement tasks poisoned");
            locked.push(actor_task);
            locked.push(tick_task);
        }
        Ok(Self {
            commands,
            assignments,
            tasks,
        })
    }

    async fn submit_cluster_job(
        &self,
        request: SubmitClusterJob,
        timeout: Duration,
    ) -> Result<crate::dcp::proto::JobStatus, DcpError> {
        self.request(timeout, |reply| PlacementCommand::Submit {
            request,
            timeout,
            reply,
        })
        .await
    }

    async fn cluster_job_status(
        &self,
        name: String,
        timeout: Duration,
    ) -> Result<crate::dcp::proto::JobStatus, DcpError> {
        self.request(timeout, |reply| PlacementCommand::Status {
            name,
            timeout,
            reply,
        })
        .await
    }

    async fn drain_cluster_job(
        &self,
        name: String,
        timeout: Duration,
    ) -> Result<crate::dcp::proto::JobStatus, DcpError> {
        self.request(timeout, |reply| PlacementCommand::Drain {
            name,
            timeout,
            reply,
        })
        .await
    }

    async fn stop_cluster_job(
        &self,
        name: String,
        timeout: Duration,
    ) -> Result<crate::dcp::proto::JobStatus, DcpError> {
        self.request(timeout, |reply| PlacementCommand::Stop {
            name,
            timeout,
            reply,
        })
        .await
    }

    async fn remember_cluster_assignment(
        &self,
        instance_name: String,
        assignment: Option<ClusterJobStart>,
    ) -> Result<(), DcpError> {
        self.request(Duration::from_millis(250), |reply| {
            PlacementCommand::Remember {
                instance_name,
                assignment,
                reply,
            }
        })
        .await
    }

    async fn register_restarted_cluster_assignment(
        &self,
        instance_name: String,
        assignment: ClusterJobStart,
        timeout: Duration,
    ) -> Result<(), DcpError> {
        self.request(timeout, |reply| PlacementCommand::RegisterRestarted {
            instance_name,
            assignment,
            timeout,
            reply,
        })
        .await
    }

    fn subscribe_assignments(&self) -> watch::Receiver<Arc<BTreeMap<String, String>>> {
        self.assignments.clone()
    }

    async fn shutdown(&self) {
        let _ = self.commands.send(PlacementCommand::Shutdown).await;
        let tasks = {
            let mut locked = self.tasks.lock().expect("placement tasks poisoned");
            std::mem::take(&mut *locked)
        };
        for task in tasks {
            task.abort();
        }
    }

    async fn request<T, F>(&self, timeout: Duration, make: F) -> Result<T, DcpError>
    where
        T: Send + 'static,
        F: FnOnce(oneshot::Sender<Result<T, DcpError>>) -> PlacementCommand,
    {
        let (reply, receiver) = oneshot::channel();
        self.commands
            .send(make(reply))
            .await
            .map_err(|_| DcpError::response(ResponseStatus::Failed, "placement actor stopped"))?;
        tokio::time::timeout(timeout.saturating_add(Duration::from_millis(250)), receiver)
            .await
            .map_err(|_| {
                DcpError::response(
                    ResponseStatus::DeadlineExceeded,
                    "placement request timed out",
                )
            })?
            .map_err(|_| DcpError::response(ResponseStatus::Failed, "placement actor stopped"))?
    }
}

async fn run_placement_actor(
    mut actor: PlacementActor,
    mut receiver: mpsc::Receiver<PlacementCommand>,
) {
    while let Some(command) = receiver.recv().await {
        match command {
            PlacementCommand::Submit {
                request,
                timeout,
                reply,
            } => {
                let result = actor.submit(request, timeout).await;
                let _ = reply.send(result);
            }
            PlacementCommand::Status {
                name,
                timeout,
                reply,
            } => {
                let result = actor.cluster_job_status(name, timeout).await;
                let _ = reply.send(result);
            }
            PlacementCommand::Drain {
                name,
                timeout,
                reply,
            } => {
                let result = actor.drain_cluster_job(name, timeout).await;
                let _ = reply.send(result);
            }
            PlacementCommand::Stop {
                name,
                timeout,
                reply,
            } => {
                let result = actor.stop_cluster_job(name, timeout).await;
                let _ = reply.send(result);
            }
            PlacementCommand::Remember {
                instance_name,
                assignment,
                reply,
            } => {
                let result = actor.remember(instance_name, assignment).await;
                let _ = reply.send(result);
            }
            PlacementCommand::RegisterRestarted {
                instance_name,
                assignment,
                timeout,
                reply,
            } => {
                let result = actor
                    .register_restarted(instance_name, assignment, timeout)
                    .await;
                let _ = reply.send(result);
            }
            PlacementCommand::Tick => {
                let _ = actor.reconcile().await;
            }
            PlacementCommand::Shutdown => break,
        }
    }
}

impl PlacementActor {
    async fn submit(
        &mut self,
        request: SubmitClusterJob,
        timeout: Duration,
    ) -> Result<crate::dcp::proto::JobStatus, DcpError> {
        self.ensure_active_coordinator(timeout).await?;
        validate_submit(&request)?;
        if self.assignments.contains_key(&request.instance_name) {
            return Err(DcpError::response(
                ResponseStatus::Conflict,
                format!("cluster job already exists: {}", request.instance_name),
            ));
        }

        let placement = placement_spec_from_wire(request.placement.clone())?;
        let target = self.choose_submit_target(&placement)?;
        let history = ClusterPlacementHistory {
            generation: 1,
            from_node: None,
            to_node: target.clone(),
            reason: "submitted".to_owned(),
            timestamp: SystemTime::now(),
        };
        let metadata = ClusterJobMetadata {
            factory_name: request.factory_name.clone(),
            params: request.params.clone().into_iter().collect(),
            placement,
            coordinator_node: self.self_node.clone(),
            assigned_node: target.clone(),
            placement_generation: 1,
            history: vec![history],
        };
        let status = self
            .start_on_node(
                &target,
                request.factory_name,
                request.instance_name.clone(),
                request.params,
                metadata.clone(),
                timeout,
            )
            .await?;
        self.assignments
            .insert(request.instance_name.clone(), metadata.clone());
        self.publish_assignments();
        self.publish_assignment(request.instance_name, metadata, timeout)
            .await?;
        Ok(status)
    }

    async fn cluster_job_status(
        &mut self,
        name: String,
        timeout: Duration,
    ) -> Result<crate::dcp::proto::JobStatus, DcpError> {
        self.ensure_active_coordinator(timeout).await?;
        let metadata = self.assignment_or_inactive(&name, timeout).await?;
        let mut status = self
            .status_on_node(&metadata.assigned_node, name, timeout)
            .await?;
        status.coordinator_node_id.clone_from(&self.self_node);
        Ok(status)
    }

    async fn drain_cluster_job(
        &mut self,
        name: String,
        timeout: Duration,
    ) -> Result<crate::dcp::proto::JobStatus, DcpError> {
        self.ensure_active_coordinator(timeout).await?;
        let metadata = self.assignment_or_rebuild(&name, timeout).await?;
        let status = self
            .drain_on_node(&metadata.assigned_node, name.clone(), timeout)
            .await?;
        self.assignments.remove(&name);
        self.publish_assignments();
        self.replicate_assignment_mutation(name, None, timeout)
            .await;
        Ok(status)
    }

    async fn stop_cluster_job(
        &mut self,
        name: String,
        timeout: Duration,
    ) -> Result<crate::dcp::proto::JobStatus, DcpError> {
        self.ensure_active_coordinator(timeout).await?;
        let metadata = self.assignment_or_rebuild(&name, timeout).await?;
        let status = self
            .stop_on_node(&metadata.assigned_node, name.clone(), timeout)
            .await?;
        self.assignments.remove(&name);
        self.publish_assignments();
        self.replicate_assignment_mutation(name, None, timeout)
            .await;
        Ok(status)
    }

    async fn remember(
        &mut self,
        instance_name: String,
        assignment: Option<ClusterJobStart>,
    ) -> Result<(), DcpError> {
        let Some(assignment) = assignment else {
            if self.assignments.remove(&instance_name).is_some() {
                self.publish_assignments();
            }
            return Ok(());
        };
        let metadata = cluster_metadata_from_wire(assignment)?;
        if !remember_assignment(
            &mut self.assignments,
            instance_name.clone(),
            metadata.clone(),
        ) {
            return Ok(());
        }
        self.publish_assignments();
        self.update_local_assignment(&instance_name, &metadata)
            .await?;
        Ok(())
    }

    async fn register_restarted(
        &mut self,
        instance_name: String,
        assignment: ClusterJobStart,
        timeout: Duration,
    ) -> Result<(), DcpError> {
        let mut metadata = cluster_metadata_from_wire(assignment)?;
        if metadata.assigned_node != self.self_node {
            return Err(DcpError::response(
                ResponseStatus::Conflict,
                "only the owning node can re-register restarted cluster intent",
            ));
        }
        let snapshot = self.state.get();
        metadata.coordinator_node = snapshot
            .placement_coordinator()
            .map(|member| member.node_id.clone())
            .ok_or_else(|| {
                DcpError::response(
                    ResponseStatus::Failed,
                    "no placement coordinator is available",
                )
            })?;
        drop(snapshot);
        self.assignments
            .insert(instance_name.clone(), metadata.clone());
        self.publish_assignments();
        self.publish_assignment(instance_name, metadata, timeout)
            .await
    }

    async fn reconcile(&mut self) -> Result<(), DcpError> {
        let snapshot = self.state.get();
        let coordinator_id = snapshot
            .placement_coordinator()
            .map(|member| member.node_id.clone());
        let is_coordinator = snapshot.is_placement_coordinator(&self.self_node);
        drop(snapshot);

        if coordinator_id != self.last_known_coordinator {
            if let Some(coordinator_id) = &coordinator_id {
                eprintln!(
                    "datum-agent INFO coordinator_changed node_id={} new_coordinator={coordinator_id}",
                    self.self_node
                );
            }
            self.last_known_coordinator = coordinator_id;
        }

        if !is_coordinator {
            if self.active_coordinator {
                eprintln!(
                    "datum-agent INFO placement_coordinator_inactive node_id={}",
                    self.self_node
                );
            }
            self.active_coordinator = false;
            return Ok(());
        }
        self.ensure_active_coordinator(self.request_timeout).await?;
        self.replace_down_members(self.request_timeout).await
    }

    async fn ensure_active_coordinator(&mut self, timeout: Duration) -> Result<(), DcpError> {
        let snapshot = self.state.get();
        if !snapshot.is_placement_coordinator(&self.self_node) {
            return Err(DcpError::response(
                ResponseStatus::Failed,
                "this node is not the placement coordinator",
            ));
        }
        drop(snapshot);
        if !self.active_coordinator {
            self.rebuild_from_registries(timeout).await?;
            self.active_coordinator = true;
            eprintln!(
                "datum-agent INFO placement_coordinator_active node_id={}",
                self.self_node
            );
        }
        Ok(())
    }

    async fn assignment_or_rebuild(
        &mut self,
        name: &str,
        timeout: Duration,
    ) -> Result<ClusterJobMetadata, DcpError> {
        if let Some(metadata) = self.assignments.get(name).cloned() {
            return Ok(metadata);
        }
        self.rebuild_from_registries(timeout).await?;
        self.assignments.get(name).cloned().ok_or_else(|| {
            DcpError::response(
                ResponseStatus::NotFound,
                format!("cluster job not found: {name}"),
            )
        })
    }

    async fn assignment_or_inactive(
        &mut self,
        name: &str,
        timeout: Duration,
    ) -> Result<ClusterJobMetadata, DcpError> {
        if let Some(metadata) = self.assignments.get(name).cloned() {
            return Ok(metadata);
        }
        let inactive = self.rebuild_from_registries(timeout).await?;
        self.assignments
            .get(name)
            .or_else(|| inactive.get(name))
            .cloned()
            .ok_or_else(|| {
                DcpError::response(
                    ResponseStatus::NotFound,
                    format!("cluster job not found: {name}"),
                )
            })
    }

    async fn rebuild_from_registries(
        &mut self,
        timeout: Duration,
    ) -> Result<BTreeMap<String, ClusterJobMetadata>, DcpError> {
        // Live registries are authoritative for both running intent and stopped tombstones. The
        // new coordinator's replicated table is retained for jobs whose owning node is already
        // down and therefore cannot participate in this handoff scan. Peer replica maps are not
        // exchanged here, so a dead-owner job is recoverable only when its upsert reached the
        // specific node that adopts coordination.
        let statuses = self.collect_job_statuses(timeout).await?;
        let mut running = BTreeSet::new();
        let mut inactive = BTreeSet::new();
        let mut inactive_metadata = BTreeMap::new();
        for status in statuses {
            if !status.cluster_job {
                continue;
            }
            if status.desired_state == "Running"
                && let Some(mut metadata) = metadata_from_status(&status)?
            {
                metadata.coordinator_node.clone_from(&self.self_node);
                running.insert(status.name.clone());
                inactive.remove(&status.name);
                inactive_metadata.remove(&status.name);
                let replace = self
                    .assignments
                    .get(&status.name)
                    .map(|existing| {
                        metadata.placement_generation > existing.placement_generation
                            || (metadata.placement_generation == existing.placement_generation
                                && metadata.assigned_node < existing.assigned_node)
                    })
                    .unwrap_or(true);
                if replace {
                    self.assignments.insert(status.name.clone(), metadata);
                    self.publish_assignments();
                }
            } else if !running.contains(&status.name) {
                if let Some(mut metadata) = metadata_from_status(&status)? {
                    metadata.coordinator_node.clone_from(&self.self_node);
                    inactive_metadata.insert(status.name.clone(), metadata);
                }
                inactive.insert(status.name);
            }
        }
        for name in &inactive {
            if self.assignments.remove(name).is_some() {
                self.publish_assignments();
            }
        }
        for metadata in self.assignments.values_mut() {
            metadata.coordinator_node.clone_from(&self.self_node);
        }
        self.publish_assignments();

        let assignments = self
            .assignments
            .iter()
            .map(|(name, metadata)| (name.clone(), metadata.clone()))
            .collect::<Vec<_>>();
        for (name, metadata) in assignments {
            self.publish_assignment(name, metadata, timeout).await?;
        }
        for name in inactive {
            if let Some(metadata) = inactive_metadata.get(&name).cloned() {
                self.publish_assignment(name.clone(), metadata, timeout)
                    .await?;
            }
            self.replicate_assignment_mutation(name, None, timeout)
                .await;
        }
        Ok(inactive_metadata)
    }

    async fn replace_down_members(&mut self, timeout: Duration) -> Result<(), DcpError> {
        let snapshot = self.state.get();
        let down_nodes = snapshot
            .members
            .values()
            .filter(|member| member.state == MemberState::Down)
            .map(|member| member.node_id.clone())
            .collect::<BTreeSet<_>>();
        drop(snapshot);
        if down_nodes.is_empty() {
            return Ok(());
        }

        let to_replace = self
            .assignments
            .iter()
            .filter(|(_name, metadata)| down_nodes.contains(&metadata.assigned_node))
            .map(|(name, metadata)| (name.clone(), metadata.clone()))
            .collect::<Vec<_>>();

        for (name, mut metadata) in to_replace {
            let old_node = metadata.assigned_node.clone();
            let target = self.choose_replacement_target(&metadata.placement, &old_node)?;
            metadata.placement_generation = metadata.placement_generation.saturating_add(1);
            metadata.assigned_node = target.clone();
            metadata.coordinator_node = self.self_node.clone();
            let reason = format!("node_down:{old_node}");
            metadata.history.push(ClusterPlacementHistory {
                generation: metadata.placement_generation,
                from_node: Some(old_node.clone()),
                to_node: target.clone(),
                reason: reason.clone(),
                timestamp: SystemTime::now(),
            });
            eprintln!(
                "datum-agent INFO job_replaced job={name} from_node={old_node} to_node={target} reason={reason} generation={}",
                metadata.placement_generation
            );
            let params = metadata.params.clone().into_iter().collect();
            let factory_name = metadata.factory_name.clone();
            self.start_on_node(
                &target,
                factory_name,
                name.clone(),
                params,
                metadata.clone(),
                timeout,
            )
            .await?;
            self.assignments.insert(name.clone(), metadata.clone());
            self.publish_assignments();
            self.publish_assignment(name, metadata, timeout).await?;
        }
        Ok(())
    }

    fn publish_assignments(&self) {
        let assignments = self
            .assignments
            .iter()
            .map(|(name, metadata)| (name.clone(), metadata.assigned_node.clone()))
            .collect();
        self.assignment_updates.send_replace(Arc::new(assignments));
    }

    fn choose_submit_target(&self, placement: &RegistryPlacementSpec) -> Result<String, DcpError> {
        let snapshot = self.state.get();
        let candidates = eligible_members(
            &snapshot,
            &self.agent_role,
            placement.role_constraint.as_deref(),
        );
        if candidates.is_empty() {
            return Err(no_eligible_node_error(placement));
        }
        match &placement.strategy {
            RegistryPlacementStrategy::LeastJobs => self.least_loaded(candidates),
            RegistryPlacementStrategy::Pinned { node_id } => candidates
                .into_iter()
                .find(|member| member.node_id == *node_id)
                .map(|member| member.node_id.clone())
                .ok_or_else(|| {
                    DcpError::response(
                        ResponseStatus::NotFound,
                        format!("pinned node is not eligible for placement: {node_id}"),
                    )
                }),
        }
    }

    fn choose_replacement_target(
        &self,
        placement: &RegistryPlacementSpec,
        old_node: &str,
    ) -> Result<String, DcpError> {
        let snapshot = self.state.get();
        let candidates = eligible_members(
            &snapshot,
            &self.agent_role,
            placement.role_constraint.as_deref(),
        )
        .into_iter()
        .filter(|member| member.node_id != old_node)
        .collect::<Vec<_>>();
        if candidates.is_empty() {
            return Err(no_eligible_node_error(placement));
        }
        self.least_loaded(candidates)
    }

    fn least_loaded(&self, candidates: Vec<&Member>) -> Result<String, DcpError> {
        candidates
            .into_iter()
            .min_by(|left, right| {
                let left_count = self.jobs_on_node(&left.node_id);
                let right_count = self.jobs_on_node(&right.node_id);
                left_count
                    .cmp(&right_count)
                    .then_with(|| left.node_id.cmp(&right.node_id))
            })
            .map(|member| member.node_id.clone())
            .ok_or_else(|| DcpError::response(ResponseStatus::NotFound, "no eligible node"))
    }

    fn jobs_on_node(&self, node_id: &str) -> usize {
        self.assignments
            .values()
            .filter(|metadata| metadata.assigned_node == node_id)
            .count()
    }

    async fn start_on_node(
        &self,
        node_id: &str,
        factory_name: String,
        instance_name: String,
        params: HashMap<String, String>,
        metadata: ClusterJobMetadata,
        timeout: Duration,
    ) -> Result<crate::dcp::proto::JobStatus, DcpError> {
        if node_id == self.self_node {
            let mut spec = self
                .factories
                .build(&factory_name, instance_name.clone(), params)?;
            spec = spec.with_cluster_metadata(metadata);
            let registry = self.registry.clone();
            let status = tokio::task::spawn_blocking(move || {
                registry.submit(spec)?;
                registry.start(instance_name)
            })
            .await??;
            return Ok(wire_job_status(&status));
        }

        self.sessions
            .start_cluster_job(
                node_id,
                factory_name,
                instance_name,
                params,
                wire_cluster_job_start(&metadata),
                timeout,
            )
            .await
            .map_err(session_error)
    }

    async fn status_on_node(
        &self,
        node_id: &str,
        name: String,
        timeout: Duration,
    ) -> Result<crate::dcp::proto::JobStatus, DcpError> {
        if node_id == self.self_node {
            let registry = self.registry.clone();
            let status = tokio::task::spawn_blocking(move || registry.status(name)).await??;
            return Ok(wire_job_status(&status));
        }
        self.sessions
            .cluster_job_status(node_id, name, timeout)
            .await
            .map_err(session_error)
    }

    async fn drain_on_node(
        &self,
        node_id: &str,
        name: String,
        timeout: Duration,
    ) -> Result<crate::dcp::proto::JobStatus, DcpError> {
        if node_id == self.self_node {
            let registry = self.registry.clone();
            let status = tokio::task::spawn_blocking(move || registry.drain(name)).await??;
            return Ok(wire_job_status(&status));
        }
        self.sessions
            .drain_cluster_job(node_id, name, timeout)
            .await
            .map_err(session_error)
    }

    async fn stop_on_node(
        &self,
        node_id: &str,
        name: String,
        timeout: Duration,
    ) -> Result<crate::dcp::proto::JobStatus, DcpError> {
        if node_id == self.self_node {
            let registry = self.registry.clone();
            let status = tokio::task::spawn_blocking(move || registry.stop(name)).await??;
            return Ok(wire_job_status(&status));
        }
        self.sessions
            .stop_cluster_job(node_id, name, timeout)
            .await
            .map_err(session_error)
    }

    async fn collect_job_statuses(
        &self,
        timeout: Duration,
    ) -> Result<Vec<crate::dcp::proto::JobStatus>, DcpError> {
        let snapshot = self.state.get();
        let mut statuses = Vec::new();
        let registry = self.registry.clone();
        let local_jobs = tokio::task::spawn_blocking(move || registry.list()).await??;
        statuses.extend(local_jobs.iter().map(wire_job_status));

        let mut pending = Vec::new();
        for member in snapshot.members.values() {
            if member.node_id == self.self_node
                || !member.has_role(&self.agent_role)
                || member.state != MemberState::Up
                || member.unreachable
            {
                continue;
            }
            let node_id = member.node_id.clone();
            let sessions = self.sessions.clone();
            pending.push(tokio::spawn(async move {
                sessions.list_jobs(&node_id, timeout).await
            }));
        }
        drop(snapshot);

        // Adoption is deliberately fail-closed until every member still considered Up responds.
        // An unresponsive former coordinator or peer delays adoption until the failure detector
        // marks it Down (bounded by suspect_timeout), which prevents a surviving job from being
        // omitted. A genuinely Up but slow peer, such as one in a GC pause, defers adoption too.
        let deadline = Instant::now() + timeout;
        for task in pending {
            let remaining = deadline.saturating_duration_since(Instant::now());
            if remaining.is_zero() {
                return Err(DcpError::response(
                    ResponseStatus::DeadlineExceeded,
                    "registry reconciliation timed out",
                ));
            }
            match tokio::time::timeout(remaining, task).await {
                Ok(Ok(Ok(mut peer_statuses))) => statuses.append(&mut peer_statuses),
                Ok(Ok(Err(message))) => return Err(session_error(message)),
                Ok(Err(error)) => {
                    return Err(DcpError::response(
                        ResponseStatus::Failed,
                        format!("registry reconciliation task failed: {error}"),
                    ));
                }
                Err(_) => {
                    return Err(DcpError::response(
                        ResponseStatus::DeadlineExceeded,
                        "registry reconciliation timed out",
                    ));
                }
            }
        }
        Ok(statuses)
    }

    async fn publish_assignment(
        &self,
        instance_name: String,
        metadata: ClusterJobMetadata,
        timeout: Duration,
    ) -> Result<(), DcpError> {
        self.update_local_assignment(&instance_name, &metadata)
            .await?;
        self.replicate_assignment_mutation(
            instance_name,
            Some(wire_cluster_job_start(&metadata)),
            timeout,
        )
        .await;
        Ok(())
    }

    async fn update_local_assignment(
        &self,
        instance_name: &str,
        metadata: &ClusterJobMetadata,
    ) -> Result<(), DcpError> {
        if metadata.assigned_node != self.self_node {
            return Ok(());
        }
        let registry = self.registry.clone();
        let instance_name = instance_name.to_owned();
        let metadata = metadata.clone();
        tokio::task::spawn_blocking(move || {
            registry.update_cluster_metadata(instance_name, metadata)
        })
        .await??;
        Ok(())
    }

    async fn replicate_assignment_mutation(
        &self,
        instance_name: String,
        assignment: Option<ClusterJobStart>,
        timeout: Duration,
    ) {
        let snapshot = self.state.get();
        let peers = snapshot
            .members
            .values()
            .filter(|member| {
                member.node_id != self.self_node
                    && member.has_role(&self.agent_role)
                    && member.state == MemberState::Up
                    && !member.unreachable
            })
            .map(|member| member.node_id.clone())
            .collect::<Vec<_>>();
        drop(snapshot);
        let timeout = timeout.min(Duration::from_millis(250));
        for node_id in peers {
            let _ = self
                .sessions
                .remember_cluster_assignment(
                    &node_id,
                    instance_name.clone(),
                    assignment.clone(),
                    timeout,
                )
                .await;
        }
    }
}

fn remember_assignment(
    assignments: &mut BTreeMap<String, ClusterJobMetadata>,
    instance_name: String,
    metadata: ClusterJobMetadata,
) -> bool {
    let accept = assignments
        .get(&instance_name)
        .is_none_or(|existing| metadata.placement_generation >= existing.placement_generation);
    if accept {
        assignments.insert(instance_name, metadata);
    }
    accept
}

fn validate_submit(request: &SubmitClusterJob) -> Result<(), DcpError> {
    if request.factory_name.trim().is_empty() || request.instance_name.trim().is_empty() {
        return Err(DcpError::response(
            ResponseStatus::BadRequest,
            "SubmitClusterJob requires factory_name and instance_name",
        ));
    }
    Ok(())
}

fn eligible_members<'a>(
    snapshot: &'a ClusterState,
    agent_role: &str,
    role_constraint: Option<&str>,
) -> Vec<&'a Member> {
    snapshot
        .members
        .values()
        .filter(|member| member.state == MemberState::Up && !member.unreachable)
        .filter(|member| member.has_role(agent_role))
        .filter(|member| role_constraint.is_none_or(|role| member.has_role(role)))
        .collect()
}

fn no_eligible_node_error(placement: &RegistryPlacementSpec) -> DcpError {
    let role = placement
        .role_constraint
        .as_deref()
        .map(|role| format!(" with role '{role}'"))
        .unwrap_or_default();
    DcpError::response(
        ResponseStatus::NotFound,
        format!("no eligible Up placement nodes{role}"),
    )
}

fn metadata_from_status(
    status: &crate::dcp::proto::JobStatus,
) -> Result<Option<ClusterJobMetadata>, DcpError> {
    if !status.cluster_job {
        return Ok(None);
    }
    let placement = placement_spec_from_wire(status.placement.clone())?;
    Ok(Some(ClusterJobMetadata {
        factory_name: status.factory_name.clone(),
        params: status.params.clone().into_iter().collect(),
        placement,
        coordinator_node: status.coordinator_node_id.clone(),
        assigned_node: status.placement_node_id.clone(),
        placement_generation: status.placement_generation,
        history: status
            .placement_history
            .iter()
            .map(|history| ClusterPlacementHistory {
                generation: history.generation,
                from_node: if history.from_node_id.is_empty() {
                    None
                } else {
                    Some(history.from_node_id.clone())
                },
                to_node: history.to_node_id.clone(),
                reason: history.reason.clone(),
                timestamp: std::time::UNIX_EPOCH + Duration::from_millis(history.timestamp_ms),
            })
            .collect(),
    }))
}

fn session_error(message: String) -> DcpError {
    let status = if message.contains("DCP response NotFound") {
        ResponseStatus::NotFound
    } else if message.contains("DCP response Conflict") {
        ResponseStatus::Conflict
    } else if message.contains("DCP response BadRequest") {
        ResponseStatus::BadRequest
    } else if message.contains("DCP response DeadlineExceeded") {
        ResponseStatus::DeadlineExceeded
    } else {
        ResponseStatus::Failed
    };
    DcpError::response(status, message)
}

struct ClusterView {
    registry: JobRegistryHandle,
    state: Signal<ClusterState>,
    sessions: NodeSessionManagerHandle,
    placement: PlacementCoordinatorHandle,
    self_node: String,
    agent_role: String,
}

impl ClusterViewProvider for ClusterView {
    fn subscribe_cluster_metrics(
        &self,
        interval: Duration,
    ) -> Option<mpsc::Receiver<MetricSample>> {
        Some(spawn_cluster_metrics_proxy(
            self.self_node.clone(),
            self.sessions.clone(),
            self.placement.subscribe_assignments(),
            interval,
        ))
    }

    fn submit_cluster_job(
        &self,
        request: SubmitClusterJob,
        timeout: Duration,
    ) -> crate::dcp::server::ClusterViewFuture<'_, crate::dcp::proto::JobStatus> {
        Box::pin(async move { self.submit_cluster_job_inner(request, timeout).await })
    }

    fn list_cluster_jobs(
        &self,
        timeout: Duration,
    ) -> crate::dcp::server::ClusterViewFuture<'_, ClusterJobList> {
        Box::pin(async move { self.list_cluster_jobs_inner(timeout).await })
    }

    fn cluster_node_info(
        &self,
        timeout: Duration,
    ) -> crate::dcp::server::ClusterViewFuture<'_, ClusterNodeList> {
        Box::pin(async move { self.cluster_node_info_inner(timeout).await })
    }

    fn cluster_job_status(
        &self,
        name: String,
        timeout: Duration,
    ) -> crate::dcp::server::ClusterViewFuture<'_, crate::dcp::proto::JobStatus> {
        Box::pin(async move { self.cluster_job_status_inner(name, timeout).await })
    }

    fn drain_cluster_job(
        &self,
        name: String,
        timeout: Duration,
    ) -> crate::dcp::server::ClusterViewFuture<'_, crate::dcp::proto::JobStatus> {
        Box::pin(async move { self.drain_cluster_job_inner(name, timeout).await })
    }

    fn stop_cluster_job(
        &self,
        name: String,
        timeout: Duration,
    ) -> crate::dcp::server::ClusterViewFuture<'_, crate::dcp::proto::JobStatus> {
        Box::pin(async move { self.stop_cluster_job_inner(name, timeout).await })
    }

    fn remember_cluster_assignment(
        &self,
        request: RememberClusterAssignment,
    ) -> crate::dcp::server::ClusterViewFuture<'_, ()> {
        Box::pin(async move { self.remember_cluster_assignment_inner(request).await })
    }

    fn register_restarted_cluster_assignment(
        &self,
        instance_name: String,
        assignment: ClusterJobStart,
        timeout: Duration,
    ) -> crate::dcp::server::ClusterViewFuture<'_, ()> {
        Box::pin(async move {
            self.placement
                .register_restarted_cluster_assignment(instance_name, assignment, timeout)
                .await
        })
    }
}

struct RemoteMetricRoute {
    job_names: BTreeSet<String>,
    cancel: oneshot::Sender<()>,
    task: JoinHandle<()>,
}

struct RemoteMetricUpdate {
    node_id: String,
    sample: MetricSample,
}

struct RemoteMetricSubscriptionAttempt {
    node_id: String,
    job_names: BTreeSet<String>,
    result: Result<MetricSubscription, String>,
}

#[derive(Default)]
struct RemoteMetricProxyState {
    routes: BTreeMap<String, RemoteMetricRoute>,
    latest: BTreeMap<String, MetricSample>,
    pending: BTreeMap<String, BTreeSet<String>>,
    wanted: BTreeMap<String, BTreeSet<String>>,
}

fn spawn_cluster_metrics_proxy(
    self_node: String,
    sessions: NodeSessionManagerHandle,
    assignments: watch::Receiver<Arc<BTreeMap<String, String>>>,
    interval: Duration,
) -> mpsc::Receiver<MetricSample> {
    let (output, receiver) = mpsc::channel(8);
    tokio::spawn(run_cluster_metrics_proxy(
        self_node,
        sessions,
        assignments,
        interval,
        output,
    ));
    receiver
}

async fn run_cluster_metrics_proxy(
    self_node: String,
    sessions: NodeSessionManagerHandle,
    mut assignments: watch::Receiver<Arc<BTreeMap<String, String>>>,
    interval: Duration,
    output: mpsc::Sender<MetricSample>,
) {
    let (updates, mut update_receiver) = mpsc::channel(32);
    let (attempts, mut attempt_receiver) = mpsc::channel(32);
    let mut proxy = RemoteMetricProxyState::default();
    let retry_interval = interval
        .max(Duration::from_millis(50))
        .min(Duration::from_secs(1));
    let mut retry = tokio::time::interval(retry_interval);
    let mut emit = tokio::time::interval(interval);

    let initial_assignments = assignments.borrow_and_update().clone();
    reconcile_metric_routes(
        &self_node,
        &sessions,
        &initial_assignments,
        interval,
        &attempts,
        &mut proxy,
    )
    .await;

    loop {
        tokio::select! {
            _ = output.closed() => break,
            changed = assignments.changed() => {
                if changed.is_err() {
                    break;
                }
                let current_assignments = assignments.borrow_and_update().clone();
                reconcile_metric_routes(
                    &self_node,
                    &sessions,
                    &current_assignments,
                    interval,
                    &attempts,
                    &mut proxy,
                ).await;
            }
            _ = retry.tick() => {
                let current_assignments = assignments.borrow().clone();
                reconcile_metric_routes(
                    &self_node,
                    &sessions,
                    &current_assignments,
                    interval,
                    &attempts,
                    &mut proxy,
                ).await;
            }
            attempt = attempt_receiver.recv() => {
                if let Some(attempt) = attempt {
                    install_metric_route(
                        attempt,
                        &mut proxy,
                        &updates,
                    ).await;
                }
            }
            update = update_receiver.recv() => {
                if let Some(update) = update
                    && proxy.routes.contains_key(&update.node_id)
                {
                    proxy.latest.insert(update.node_id, update.sample);
                }
            }
            _ = emit.tick() => {
                let streams = proxy.latest
                    .values()
                    .flat_map(|sample| sample.streams.iter().cloned())
                    .collect();
                let timestamp_ms = proxy.latest
                    .values()
                    .map(|sample| sample.timestamp_ms)
                    .max()
                    .unwrap_or(0);
                match output.try_send(MetricSample { timestamp_ms, streams }) {
                    Ok(()) | Err(mpsc::error::TrySendError::Full(_)) => {}
                    Err(mpsc::error::TrySendError::Closed(_)) => break,
                }
            }
        }
    }

    drop(attempt_receiver);
    drop(attempts);
    cancel_metric_routes(proxy.routes).await;
}

async fn reconcile_metric_routes(
    self_node: &str,
    sessions: &NodeSessionManagerHandle,
    assignments: &BTreeMap<String, String>,
    interval: Duration,
    attempts: &mpsc::Sender<RemoteMetricSubscriptionAttempt>,
    proxy: &mut RemoteMetricProxyState,
) {
    let mut next_wanted = BTreeMap::<String, BTreeSet<String>>::new();
    for (job_name, node_id) in assignments {
        if node_id != self_node {
            next_wanted
                .entry(node_id.clone())
                .or_default()
                .insert(job_name.clone());
        }
    }
    proxy.wanted = next_wanted;

    let stale = proxy
        .routes
        .iter()
        .filter(|(node_id, route)| {
            proxy.wanted.get(*node_id) != Some(&route.job_names) || route.task.is_finished()
        })
        .map(|(node_id, _)| node_id.clone())
        .collect::<Vec<_>>();
    for node_id in stale {
        proxy.latest.remove(&node_id);
        if let Some(route) = proxy.routes.remove(&node_id) {
            cancel_metric_route(route).await;
        }
    }

    for (node_id, job_names) in &proxy.wanted {
        if proxy.routes.contains_key(node_id) || proxy.pending.get(node_id) == Some(job_names) {
            continue;
        }
        proxy.pending.insert(node_id.clone(), job_names.clone());
        let timeout = sessions.inner.config.request_timeout;
        let sessions = sessions.clone();
        let attempts = attempts.clone();
        let node_id = node_id.clone();
        let job_names = job_names.clone();
        let _task = tokio::spawn(async move {
            let result = sessions
                .subscribe_metrics(
                    &node_id,
                    interval,
                    job_names.iter().cloned().collect(),
                    timeout,
                )
                .await;
            let attempt = RemoteMetricSubscriptionAttempt {
                node_id,
                job_names,
                result,
            };
            if let Err(error) = attempts.send(attempt).await
                && let Ok(subscription) = error.0.result
            {
                let _ = subscription.cancel().await;
            }
        });
    }
}

async fn install_metric_route(
    attempt: RemoteMetricSubscriptionAttempt,
    proxy: &mut RemoteMetricProxyState,
    updates: &mpsc::Sender<RemoteMetricUpdate>,
) {
    let RemoteMetricSubscriptionAttempt {
        node_id,
        job_names,
        result,
    } = attempt;
    if proxy.pending.get(&node_id) == Some(&job_names) {
        proxy.pending.remove(&node_id);
    }
    let is_current =
        proxy.wanted.get(&node_id) == Some(&job_names) && !proxy.routes.contains_key(&node_id);
    let Ok(subscription) = result else {
        return;
    };
    if !is_current {
        let _ = subscription.cancel().await;
        return;
    }

    let (cancel, cancel_receiver) = oneshot::channel();
    let task = tokio::spawn(pump_remote_metrics(
        node_id.clone(),
        subscription,
        job_names.clone(),
        updates.clone(),
        cancel_receiver,
    ));
    proxy.routes.insert(
        node_id,
        RemoteMetricRoute {
            job_names,
            cancel,
            task,
        },
    );
}

async fn pump_remote_metrics(
    node_id: String,
    mut subscription: MetricSubscription,
    job_names: BTreeSet<String>,
    updates: mpsc::Sender<RemoteMetricUpdate>,
    mut cancel: oneshot::Receiver<()>,
) {
    loop {
        tokio::select! {
            _ = &mut cancel => {
                let _ = subscription.cancel().await;
                break;
            }
            sample = subscription.recv() => {
                let Some(mut sample) = sample else {
                    break;
                };
                sample.streams.retain(|metric| {
                    job_names.iter().any(|job_name| {
                        metric.name == *job_name
                            || metric
                                .name
                                .strip_prefix(job_name)
                                .is_some_and(|suffix| suffix.starts_with(':'))
                    })
                });
                match updates.try_send(RemoteMetricUpdate {
                    node_id: node_id.clone(),
                    sample,
                }) {
                    Ok(()) | Err(mpsc::error::TrySendError::Full(_)) => {}
                    Err(mpsc::error::TrySendError::Closed(_)) => {
                        let _ = subscription.cancel().await;
                        break;
                    }
                }
            }
        }
    }
}

async fn cancel_metric_route(route: RemoteMetricRoute) {
    let _ = route.cancel.send(());
    let _ = route.task.await;
}

async fn cancel_metric_routes(routes: BTreeMap<String, RemoteMetricRoute>) {
    for (_, route) in routes {
        cancel_metric_route(route).await;
    }
}

impl ClusterView {
    async fn submit_cluster_job_inner(
        &self,
        request: SubmitClusterJob,
        timeout: Duration,
    ) -> Result<crate::dcp::proto::JobStatus, DcpError> {
        if self.is_local_coordinator() {
            return self.placement.submit_cluster_job(request, timeout).await;
        }
        let coordinator = self.coordinator_node_id()?;
        self.sessions
            .submit_cluster_job(&coordinator, request, timeout)
            .await
            .map_err(session_error)
    }

    async fn list_cluster_jobs_inner(&self, timeout: Duration) -> Result<ClusterJobList, DcpError> {
        if !self.is_local_coordinator() {
            let coordinator = self.coordinator_node_id()?;
            return self
                .sessions
                .list_cluster_jobs(&coordinator, timeout)
                .await
                .map_err(session_error);
        }

        let timeout = if timeout.is_zero() {
            self.sessions.inner.config.request_timeout
        } else {
            timeout
        };
        let snapshot = self.state.get();
        let local_member = snapshot.member(&self.self_node).cloned();
        let local_jobs = list_local_jobs(self.registry.clone()).await?;
        let mut nodes = Vec::new();
        nodes.push(ClusterJobNode {
            node_id: self.self_node.clone(),
            address: local_member
                .as_ref()
                .map(|member| member.address.to_string())
                .unwrap_or_default(),
            local: true,
            jobs: local_jobs,
        });

        let mut errors = Vec::new();
        let mut pending = Vec::new();
        for member in snapshot.members.values() {
            if member.node_id == self.self_node || !member.has_role(&self.agent_role) {
                continue;
            }
            if member.state == MemberState::Removed {
                continue;
            }
            if member.state != MemberState::Up || member.unreachable {
                errors.push(ClusterNodeError {
                    node_id: member.node_id.clone(),
                    message: if member.unreachable {
                        "member is unreachable".to_owned()
                    } else {
                        format!("member is {:?}", member.state)
                    },
                });
                continue;
            }
            let node_id = member.node_id.clone();
            let address = member.address.to_string();
            let sessions = self.sessions.clone();
            pending.push(tokio::spawn(async move {
                let result = sessions.list_jobs(&node_id, timeout).await;
                (node_id, address, result)
            }));
        }

        let deadline = Instant::now() + timeout;
        for task in pending {
            let remaining = deadline.saturating_duration_since(Instant::now());
            match tokio::time::timeout(remaining, task).await {
                Ok(Ok((node_id, address, Ok(jobs)))) => nodes.push(ClusterJobNode {
                    node_id,
                    address,
                    local: false,
                    jobs,
                }),
                Ok(Ok((node_id, _address, Err(message)))) => {
                    errors.push(ClusterNodeError { node_id, message });
                }
                Ok(Err(error)) => errors.push(ClusterNodeError {
                    node_id: "unknown".to_owned(),
                    message: format!("cluster fan-out task failed: {error}"),
                }),
                Err(_) => errors.push(ClusterNodeError {
                    node_id: "unknown".to_owned(),
                    message: "cluster fan-out timed out".to_owned(),
                }),
            }
        }

        nodes.sort_by(|left, right| left.node_id.cmp(&right.node_id));
        errors.sort_by(|left, right| left.node_id.cmp(&right.node_id));
        Ok(ClusterJobList {
            partial: !errors.is_empty(),
            nodes,
            errors,
        })
    }

    async fn cluster_node_info_inner(
        &self,
        _timeout: Duration,
    ) -> Result<ClusterNodeList, DcpError> {
        let snapshot = self.state.get();
        let coordinator_node_id = snapshot
            .placement_coordinator()
            .map(|member| member.node_id.clone())
            .unwrap_or_default();
        let session_states = self.sessions.session_states().await;
        let mut nodes = snapshot
            .members
            .values()
            .map(|member| ClusterNodeStatus {
                node_id: member.node_id.clone(),
                member_state: format!("{:?}", member.state),
                address: member.address.to_string(),
                agent_addr: member
                    .agent_addr
                    .map(|addr| addr.to_string())
                    .unwrap_or_default(),
                roles: member.roles.clone(),
                unreachable: member.unreachable,
                local: member.node_id == self.self_node,
                session_state: if member.node_id == self.self_node {
                    "local".to_owned()
                } else {
                    session_states
                        .get(&member.node_id)
                        .cloned()
                        .unwrap_or_else(|| "not_connected".to_owned())
                },
            })
            .collect::<Vec<_>>();
        nodes.sort_by(|left, right| left.node_id.cmp(&right.node_id));

        let mut errors = nodes
            .iter()
            .filter(|node| !node.local)
            .filter(|node| node.roles.iter().any(|role| role == &self.agent_role))
            .filter(|node| {
                node.member_state == "Up" && !node.unreachable && node.session_state != "connected"
            })
            .map(|node| ClusterNodeError {
                node_id: node.node_id.clone(),
                message: format!("node session {}", node.session_state),
            })
            .collect::<Vec<_>>();
        errors.sort_by(|left, right| left.node_id.cmp(&right.node_id));

        Ok(ClusterNodeList {
            partial: !errors.is_empty(),
            nodes,
            errors,
            coordinator_node_id,
        })
    }

    async fn cluster_job_status_inner(
        &self,
        name: String,
        timeout: Duration,
    ) -> Result<crate::dcp::proto::JobStatus, DcpError> {
        if self.is_local_coordinator() {
            return self.placement.cluster_job_status(name, timeout).await;
        }
        let coordinator = self.coordinator_node_id()?;
        self.sessions
            .forward_cluster_job_status(&coordinator, name, timeout)
            .await
            .map_err(session_error)
    }

    async fn drain_cluster_job_inner(
        &self,
        name: String,
        timeout: Duration,
    ) -> Result<crate::dcp::proto::JobStatus, DcpError> {
        if self.is_local_coordinator() {
            return self.placement.drain_cluster_job(name, timeout).await;
        }
        let coordinator = self.coordinator_node_id()?;
        self.sessions
            .forward_drain_cluster_job(&coordinator, name, timeout)
            .await
            .map_err(session_error)
    }

    async fn stop_cluster_job_inner(
        &self,
        name: String,
        timeout: Duration,
    ) -> Result<crate::dcp::proto::JobStatus, DcpError> {
        if self.is_local_coordinator() {
            return self.placement.stop_cluster_job(name, timeout).await;
        }
        let coordinator = self.coordinator_node_id()?;
        self.sessions
            .forward_stop_cluster_job(&coordinator, name, timeout)
            .await
            .map_err(session_error)
    }

    async fn remember_cluster_assignment_inner(
        &self,
        request: RememberClusterAssignment,
    ) -> Result<(), DcpError> {
        self.placement
            .remember_cluster_assignment(request.instance_name, request.assignment)
            .await
    }

    fn is_local_coordinator(&self) -> bool {
        self.state.get().is_placement_coordinator(&self.self_node)
    }

    fn coordinator_node_id(&self) -> Result<String, DcpError> {
        self.state
            .get()
            .placement_coordinator()
            .map(|member| member.node_id.clone())
            .ok_or_else(|| {
                DcpError::response(
                    ResponseStatus::Failed,
                    "no placement coordinator is available",
                )
            })
    }
}

async fn list_local_jobs(
    registry: JobRegistryHandle,
) -> Result<Vec<crate::dcp::proto::JobStatus>, DcpError> {
    let statuses = tokio::task::spawn_blocking(move || registry.list()).await??;
    Ok(statuses.iter().map(wire_job_status).collect())
}

fn validate_cluster_agent_config(config: &ClusterAgentConfig) -> ClusterAgentResult<()> {
    if config.sessions.agent_role.trim().is_empty() {
        return Err(ClusterAgentError::InvalidConfig(
            "sessions.agent_role must not be empty".to_owned(),
        ));
    }
    if config.sessions.reconnect_min_backoff.is_zero()
        || config.sessions.reconnect_max_backoff < config.sessions.reconnect_min_backoff
    {
        return Err(ClusterAgentError::InvalidConfig(
            "node-session reconnect backoff must be non-zero and ordered".to_owned(),
        ));
    }
    match config.sessions.transport {
        NodeSessionTransport::TcpLoopback if config.dcp.tcp.is_none() => {
            Err(ClusterAgentError::InvalidConfig(
                "node-session TCP transport requires a DCP TCP listener".to_owned(),
            ))
        }
        NodeSessionTransport::QuicMtls { .. } if config.dcp.quic.is_none() => {
            Err(ClusterAgentError::InvalidConfig(
                "node-session QUIC transport requires a DCP QUIC listener".to_owned(),
            ))
        }
        _ => Ok(()),
    }
}

fn advertised_agent_addr(
    transport: &NodeSessionTransport,
    handle: &DcpServerHandle,
) -> ClusterAgentResult<SocketAddr> {
    match transport {
        NodeSessionTransport::TcpLoopback => handle.tcp_addr().ok_or_else(|| {
            ClusterAgentError::InvalidConfig("DCP TCP listener did not start".to_owned())
        }),
        NodeSessionTransport::QuicMtls { .. } => handle.quic_addr().ok_or_else(|| {
            ClusterAgentError::InvalidConfig("DCP QUIC listener did not start".to_owned())
        }),
    }
}

fn ensure_agent_role(config: &mut ClusterConfig, role: &str) {
    if !config.roles.iter().any(|candidate| candidate == role) {
        config.roles.push(role.to_owned());
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn remembered_assignment_generation_is_monotonic() {
        let mut assignments = BTreeMap::new();
        assert!(remember_assignment(
            &mut assignments,
            "job".to_owned(),
            assignment(2, "current-owner"),
        ));

        assert!(!remember_assignment(
            &mut assignments,
            "job".to_owned(),
            assignment(1, "stale-owner"),
        ));
        assert_eq!(assignments["job"].assigned_node, "current-owner");
        assert_eq!(assignments["job"].placement_generation, 2);

        assert!(remember_assignment(
            &mut assignments,
            "job".to_owned(),
            assignment(2, "equal-generation-owner"),
        ));
        assert_eq!(assignments["job"].assigned_node, "equal-generation-owner");

        assert!(remember_assignment(
            &mut assignments,
            "job".to_owned(),
            assignment(3, "new-owner"),
        ));
        assert_eq!(assignments["job"].assigned_node, "new-owner");
        assert_eq!(assignments["job"].placement_generation, 3);
    }

    fn assignment(placement_generation: u64, assigned_node: &str) -> ClusterJobMetadata {
        ClusterJobMetadata {
            factory_name: "ticker".to_owned(),
            params: BTreeMap::new(),
            placement: RegistryPlacementSpec::least_jobs(None),
            coordinator_node: "coordinator".to_owned(),
            assigned_node: assigned_node.to_owned(),
            placement_generation,
            history: Vec::new(),
        }
    }
}