datum_agent/dcp/

server.rs

use std::{
    collections::{BTreeMap, HashMap},
    future::Future,
    net::{IpAddr, Ipv4Addr, Ipv6Addr, SocketAddr},
    pin::Pin,
    sync::{Arc, Mutex},
    time::{Duration, Instant, SystemTime, UNIX_EPOCH},
};

use datum::{
    Sink, StreamInstrumentationRegistry, StreamInstrumentationSnapshot, StreamInstrumentationState,
};
use datum_net::quic::quinn;
use prost::Message as ProstMessage;
use tokio::{
    io::{AsyncRead, AsyncWrite, AsyncWriteExt},
    net::{TcpListener, TcpStream},
    sync::{mpsc, watch},
    task::JoinHandle,
};

use crate::{
    AgentError, AgentHandle, AgentResult, ClusterJobMetadata, ClusterPlacementHistory, JobEvent,
    JobEventKind, JobExitReason, JobRegistryHandle, JobSpec, JobStatus as RegistryJobStatus,
    PlacementSpec as RegistryPlacementSpec, PlacementStrategy as RegistryPlacementStrategy,
    dcp::{
        DcpError, DcpResult,
        frame::{read_frame, write_frame},
        proto::{
            ClusterJobList, ClusterJobStart, ClusterNodeList, CompleteShardingAsk, ConfigValue,
            DCP_PROTOCOL_MAJOR, DcpFrame, Event, ForwardShardEnvelopes, Hello, JobList,
            JobStatus as WireJobStatus, MetricSample, PlacementSpec, PlacementStrategy,
            RememberClusterAssignment, RememberShardAllocations, Request, Response, ResponseStatus,
            ShardAllocation, ShardAllocationRequest, ShardAllocationTable,
            ShardEnvelopeBatchResult, ShardPipeFrame, StreamMetric, SubmitClusterJob, dcp_frame,
            request,
        },
    },
};

type DcpJobFactory =
    dyn Fn(String, HashMap<String, String>) -> AgentResult<JobSpec> + Send + Sync + 'static;
pub type ClusterViewFuture<'a, T> = Pin<Box<dyn Future<Output = DcpResult<T>> + Send + 'a>>;
const DEFAULT_CLUSTER_REQUEST_TIMEOUT: Duration = Duration::from_millis(750);

/// Provider installed by cluster-aware agents to serve read-only cluster DCP
/// requests.
pub trait ClusterViewProvider: Send + Sync + 'static {
    /// Submit a cluster-scoped job to the deterministic placement coordinator.
    fn submit_cluster_job(
        &self,
        request: SubmitClusterJob,
        timeout: Duration,
    ) -> ClusterViewFuture<'_, WireJobStatus>;

    /// Aggregate job snapshots from local and peer nodes.
    fn list_cluster_jobs(&self, timeout: Duration) -> ClusterViewFuture<'_, ClusterJobList>;

    /// Return the local cluster membership and DCP node-session view.
    fn cluster_node_info(&self, timeout: Duration) -> ClusterViewFuture<'_, ClusterNodeList>;

    /// Return status for a cluster job, routed through the coordinator table.
    fn cluster_job_status(
        &self,
        name: String,
        timeout: Duration,
    ) -> ClusterViewFuture<'_, WireJobStatus>;

    /// Drain a cluster job on its current placement node.
    fn drain_cluster_job(
        &self,
        name: String,
        timeout: Duration,
    ) -> ClusterViewFuture<'_, WireJobStatus>;

    /// Stop a cluster job on its current placement node.
    fn stop_cluster_job(
        &self,
        name: String,
        timeout: Duration,
    ) -> ClusterViewFuture<'_, WireJobStatus>;

    /// Remember an assignment replicated by the active coordinator.
    fn remember_cluster_assignment(
        &self,
        request: RememberClusterAssignment,
    ) -> ClusterViewFuture<'_, ()>;
}

/// Provider installed by cluster sharding to serve shard allocation, region
/// forwarding, and remote ask replies over the existing node-to-node DCP
/// sessions.
pub trait ShardingViewProvider: Send + Sync + 'static {
    /// Allocate or return the cached owner for one shard. Only the deterministic
    /// coordinator should accept this request.
    fn allocate_shard(
        &self,
        request: ShardAllocationRequest,
        timeout: Duration,
    ) -> ClusterViewFuture<'_, ShardAllocation>;

    /// Remember allocation-table rows replicated by the active coordinator.
    fn remember_shard_allocations(
        &self,
        request: RememberShardAllocations,
    ) -> ClusterViewFuture<'_, ()>;

    /// Return this region's cached allocation table for coordinator takeover.
    fn get_shard_allocations(
        &self,
        type_name: String,
        timeout: Duration,
    ) -> ClusterViewFuture<'_, ShardAllocationTable>;

    /// Enqueue a batch of serialized envelopes at this region.
    fn forward_shard_envelopes(
        &self,
        request: ForwardShardEnvelopes,
        timeout: Duration,
    ) -> ClusterViewFuture<'_, ShardEnvelopeBatchResult>;

    /// Complete an ask whose reply port originated on this region.
    fn complete_sharding_ask(&self, request: CompleteShardingAsk) -> ClusterViewFuture<'_, ()>;
}

#[derive(Clone, Default)]
pub struct DcpJobFactories {
    factories: Arc<Mutex<HashMap<String, Arc<DcpJobFactory>>>>,
}

impl DcpJobFactories {
    #[must_use]
    pub fn new() -> Self {
        Self::default()
    }

    pub fn register<F>(&self, name: impl Into<String>, factory: F) -> AgentResult<()>
    where
        F: Fn(String, HashMap<String, String>) -> AgentResult<JobSpec> + Send + Sync + 'static,
    {
        let name = name.into();
        if name.trim().is_empty() {
            return Err(AgentError::InvalidJobName);
        }
        self.factories
            .lock()
            .expect("DCP job factories poisoned")
            .insert(name, Arc::new(factory));
        Ok(())
    }

    pub(crate) fn build(
        &self,
        factory_name: &str,
        instance_name: String,
        params: HashMap<String, String>,
    ) -> DcpResult<JobSpec> {
        let factory = self
            .factories
            .lock()
            .expect("DCP job factories poisoned")
            .get(factory_name)
            .cloned()
            .ok_or_else(|| {
                DcpError::response(
                    ResponseStatus::NotFound,
                    format!("job factory not found: {factory_name}"),
                )
            })?;
        Ok(factory(instance_name, params)?)
    }
}

#[derive(Clone)]
pub struct DcpTcpServerConfig {
    pub addr: SocketAddr,
}

#[derive(Clone)]
pub struct DcpQuicServerConfig {
    pub addr: SocketAddr,
    pub server_config: quinn::ServerConfig,
}

#[derive(Clone)]
pub struct DcpServerConfig {
    pub node_id: String,
    pub tcp: Option<DcpTcpServerConfig>,
    pub quic: Option<DcpQuicServerConfig>,
    pub auth_token: Option<String>,
    pub metrics_interval: Duration,
    pub frame_buffer: usize,
}

impl Default for DcpServerConfig {
    fn default() -> Self {
        Self {
            node_id: format!("datum-agent-{}", std::process::id()),
            tcp: Some(DcpTcpServerConfig {
                addr: SocketAddr::new(IpAddr::V4(Ipv4Addr::LOCALHOST), 0),
            }),
            quic: None,
            auth_token: None,
            metrics_interval: Duration::from_secs(1),
            frame_buffer: 256,
        }
    }
}

#[derive(Clone)]
pub struct DcpServer {
    state: Arc<DcpServerState>,
}

struct DcpServerState {
    registry: JobRegistryHandle,
    instrumentation: StreamInstrumentationRegistry,
    factories: DcpJobFactories,
    config: DcpServerConfig,
    config_store: Mutex<HashMap<String, String>>,
    cluster_view: Mutex<Option<Arc<dyn ClusterViewProvider>>>,
    sharding_view: Mutex<Option<Arc<dyn ShardingViewProvider>>>,
}

impl DcpServer {
    #[must_use]
    pub fn new(
        registry: JobRegistryHandle,
        instrumentation: StreamInstrumentationRegistry,
        factories: DcpJobFactories,
        config: DcpServerConfig,
    ) -> Self {
        Self {
            state: Arc::new(DcpServerState {
                registry,
                instrumentation,
                factories,
                config,
                config_store: Mutex::new(HashMap::new()),
                cluster_view: Mutex::new(None),
                sharding_view: Mutex::new(None),
            }),
        }
    }

    #[must_use]
    pub fn from_agent(
        agent: &AgentHandle,
        factories: DcpJobFactories,
        config: DcpServerConfig,
    ) -> Self {
        Self::new(
            agent.registry().clone(),
            agent.instrumentation_registry().clone(),
            factories,
            config,
        )
    }

    /// Install or replace the provider used for cluster-wide read-only DCP
    /// requests.
    pub fn set_cluster_view(&self, provider: Arc<dyn ClusterViewProvider>) {
        *self
            .state
            .cluster_view
            .lock()
            .expect("DCP cluster view provider poisoned") = Some(provider);
    }

    /// Remove the cluster view provider. Cluster DCP requests will fail until a
    /// provider is installed again.
    pub fn clear_cluster_view(&self) {
        *self
            .state
            .cluster_view
            .lock()
            .expect("DCP cluster view provider poisoned") = None;
    }

    /// Install or replace the provider used for cluster-sharding DCP requests.
    pub fn set_sharding_view(&self, provider: Arc<dyn ShardingViewProvider>) {
        *self
            .state
            .sharding_view
            .lock()
            .expect("DCP sharding view provider poisoned") = Some(provider);
    }

    /// Remove the sharding provider. Sharding DCP requests will fail until a
    /// provider is installed again.
    pub fn clear_sharding_view(&self) {
        *self
            .state
            .sharding_view
            .lock()
            .expect("DCP sharding view provider poisoned") = None;
    }

    pub async fn start(&self) -> DcpResult<DcpServerHandle> {
        let (shutdown_sender, shutdown_receiver) = watch::channel(false);
        let mut tasks = Vec::new();
        let mut tcp_addr = None;
        let mut quic_addr = None;

        if let Some(tcp) = &self.state.config.tcp {
            ensure_loopback(tcp.addr)?;
            let listener = TcpListener::bind(tcp.addr).await?;
            tcp_addr = Some(listener.local_addr()?);
            let state = Arc::clone(&self.state);
            let shutdown = shutdown_receiver.clone();
            tasks.push(tokio::spawn(async move {
                run_tcp_listener(listener, state, shutdown).await;
            }));
        }

        if let Some(quic) = &self.state.config.quic {
            let endpoint = quinn::Endpoint::server(quic.server_config.clone(), quic.addr)?;
            quic_addr = Some(endpoint.local_addr()?);
            let state = Arc::clone(&self.state);
            let shutdown = shutdown_receiver.clone();
            tasks.push(tokio::spawn(async move {
                run_quic_listener(endpoint, state, shutdown).await;
            }));
        }

        if tcp_addr.is_none() && quic_addr.is_none() {
            return Err(DcpError::Protocol(
                "DCP server has no configured listeners".to_owned(),
            ));
        }

        Ok(DcpServerHandle {
            tcp_addr,
            quic_addr,
            shutdown: shutdown_sender,
            tasks,
        })
    }
}

#[must_use = "dropping DcpServerHandle immediately aborts the DCP listeners; hold it for the server's lifetime"]
pub struct DcpServerHandle {
    tcp_addr: Option<SocketAddr>,
    quic_addr: Option<SocketAddr>,
    shutdown: watch::Sender<bool>,
    tasks: Vec<JoinHandle<()>>,
}

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

    #[must_use]
    pub fn quic_addr(&self) -> Option<SocketAddr> {
        self.quic_addr
    }

    pub async fn shutdown(mut self) {
        let _ = self.shutdown.send(true);
        for task in self.tasks.drain(..) {
            task.abort();
            let _ = task.await;
        }
    }
}

impl Drop for DcpServerHandle {
    fn drop(&mut self) {
        let _ = self.shutdown.send(true);
        for task in &self.tasks {
            task.abort();
        }
    }
}

async fn run_tcp_listener(
    listener: TcpListener,
    state: Arc<DcpServerState>,
    mut shutdown: watch::Receiver<bool>,
) {
    loop {
        tokio::select! {
            changed = shutdown.changed() => {
                if changed.is_err() || *shutdown.borrow() {
                    break;
                }
            }
            accepted = listener.accept() => {
                let Ok((stream, _peer)) = accepted else {
                    break;
                };
                let state = Arc::clone(&state);
                tokio::spawn(async move {
                    let _ = run_tcp_connection(stream, state).await;
                });
            }
        }
    }
}

async fn run_tcp_connection(stream: TcpStream, state: Arc<DcpServerState>) -> DcpResult<()> {
    stream.set_nodelay(true)?;
    let (reader, writer) = stream.into_split();
    run_connection(reader, writer, state).await
}

async fn run_quic_listener(
    endpoint: quinn::Endpoint,
    state: Arc<DcpServerState>,
    mut shutdown: watch::Receiver<bool>,
) {
    loop {
        tokio::select! {
            changed = shutdown.changed() => {
                if changed.is_err() || *shutdown.borrow() {
                    endpoint.close(quinn::VarInt::from_u32(0), b"DCP shutdown");
                    break;
                }
            }
            incoming = endpoint.accept() => {
                let Some(incoming) = incoming else {
                    break;
                };
                let state = Arc::clone(&state);
                let shutdown = shutdown.clone();
                tokio::spawn(async move {
                    if let Ok(connection) = incoming.await {
                        run_quic_connection(connection, state, shutdown).await;
                    }
                });
            }
        }
    }
}

async fn run_quic_connection(
    connection: quinn::Connection,
    state: Arc<DcpServerState>,
    mut shutdown: watch::Receiver<bool>,
) {
    loop {
        tokio::select! {
            changed = shutdown.changed() => {
                if changed.is_err() || *shutdown.borrow() {
                    connection.close(quinn::VarInt::from_u32(0), b"DCP shutdown");
                    break;
                }
            }
            accepted = connection.accept_bi() => {
                let Ok((send, recv)) = accepted else {
                    break;
                };
                let state = Arc::clone(&state);
                tokio::spawn(async move {
                    let _ = run_connection(recv, send, state).await;
                });
            }
        }
    }
}

async fn run_connection<R, W>(mut reader: R, writer: W, state: Arc<DcpServerState>) -> DcpResult<()>
where
    R: AsyncRead + Unpin + Send + 'static,
    W: AsyncWrite + Unpin + Send + 'static,
{
    let (outbound, outbound_receiver) = mpsc::channel(state.config.frame_buffer.max(1));
    let writer_task = tokio::spawn(write_loop(writer, outbound_receiver));
    let mut subscriptions = Vec::new();

    let Some(first) = read_frame(&mut reader).await? else {
        return Err(DcpError::Closed);
    };
    let hello = match first.frame {
        Some(dcp_frame::Frame::Hello(hello)) => hello,
        _ => {
            let response =
                Response::error(0, ResponseStatus::BadRequest, "first frame must be Hello");
            send_frame(&outbound, DcpFrame::response(response)).await?;
            return Err(DcpError::Protocol("first frame must be Hello".to_owned()));
        }
    };

    let hello_response = negotiate_hello(&state, &hello);
    let accepted = hello_response.response_status() == ResponseStatus::Ok;
    send_frame(&outbound, DcpFrame::response(hello_response)).await?;
    if !accepted {
        return Ok(());
    }

    while let Some(frame) = read_frame(&mut reader).await? {
        let request = match frame.frame {
            Some(dcp_frame::Frame::Request(request)) => request,
            _ => {
                return Err(DcpError::Protocol(
                    "client sent non-request frame after hello".to_owned(),
                ));
            }
        };
        if matches!(&request.command, Some(request::Command::OpenShardPipe(_))) {
            send_frame(
                &outbound,
                DcpFrame::response(Response::ok(request.request_id, Vec::new())),
            )
            .await?;
            run_shard_pipe_connection(reader, outbound.clone(), Arc::clone(&state)).await?;
            break;
        }
        let (response, subscription) =
            dispatch_request(Arc::clone(&state), request, outbound.clone()).await;
        send_frame(&outbound, DcpFrame::response(response)).await?;
        if let Some(subscription) = subscription {
            subscriptions.push(subscription);
        }
    }

    for subscription in subscriptions {
        subscription.abort();
    }
    drop(outbound);
    writer_task.await??;
    Ok(())
}

async fn run_shard_pipe_connection<R>(
    mut reader: R,
    outbound: mpsc::Sender<DcpFrame>,
    state: Arc<DcpServerState>,
) -> DcpResult<()>
where
    R: AsyncRead + Unpin,
{
    while let Some(frame) = read_frame(&mut reader).await? {
        let pipe = match frame.frame {
            Some(dcp_frame::Frame::ShardPipe(pipe)) => pipe,
            _ => {
                return Err(DcpError::Protocol(
                    "shard pipe connection received a non-pipe frame".to_owned(),
                ));
            }
        };
        dispatch_shard_pipe(Arc::clone(&state), pipe).await?;
    }
    drop(outbound);
    Ok(())
}

async fn dispatch_shard_pipe(state: Arc<DcpServerState>, pipe: ShardPipeFrame) -> DcpResult<()> {
    let provider = sharding_view_provider(&state)?;
    for batch in pipe.forwards {
        provider
            .forward_shard_envelopes(batch, DEFAULT_CLUSTER_REQUEST_TIMEOUT)
            .await?;
    }
    for reply in pipe.replies {
        provider.complete_sharding_ask(reply).await?;
    }
    Ok(())
}

async fn write_loop<W>(mut writer: W, mut outbound: mpsc::Receiver<DcpFrame>) -> DcpResult<()>
where
    W: AsyncWrite + Unpin,
{
    while let Some(frame) = outbound.recv().await {
        write_frame(&mut writer, &frame).await?;
    }
    let _ = writer.shutdown().await;
    Ok(())
}

async fn send_frame(sender: &mpsc::Sender<DcpFrame>, frame: DcpFrame) -> DcpResult<()> {
    sender.send(frame).await.map_err(|_| DcpError::Closed)
}

fn negotiate_hello(state: &DcpServerState, hello: &Hello) -> Response {
    match protocol_major(&hello.protocol_version) {
        Some(DCP_PROTOCOL_MAJOR) => {}
        Some(other) => {
            return Response::error(
                0,
                ResponseStatus::ProtocolMismatch,
                format!("unsupported DCP major version: {other}"),
            );
        }
        None => {
            return Response::error(
                0,
                ResponseStatus::ProtocolMismatch,
                format!("invalid DCP protocol version: {}", hello.protocol_version),
            );
        }
    }

    if let Some(expected) = &state.config.auth_token {
        let actual = hello
            .auth
            .as_ref()
            .map(|auth| auth.bearer_token.as_str())
            .unwrap_or_default();
        if actual != expected {
            return Response::error(0, ResponseStatus::Unauthorized, "invalid DCP token");
        }
    }

    Response::ok(0, state.config.node_id.as_bytes().to_vec())
}

fn protocol_major(version: &str) -> Option<u32> {
    version.split('.').next()?.parse().ok()
}

async fn dispatch_request(
    state: Arc<DcpServerState>,
    request: Request,
    outbound: mpsc::Sender<DcpFrame>,
) -> (Response, Option<JoinHandle<()>>) {
    let request_id = request.request_id;
    let deadline = request.deadline_ms;
    let dispatch = async {
        let command = request.command.ok_or_else(|| {
            DcpError::response(ResponseStatus::BadRequest, "request missing command")
        })?;
        dispatch_command(state, request_id, command, outbound).await
    };

    let result = if deadline == 0 {
        dispatch.await
    } else {
        match tokio::time::timeout(Duration::from_millis(deadline), dispatch).await {
            Ok(result) => result,
            Err(_) => Err(DcpError::response(
                ResponseStatus::DeadlineExceeded,
                "request deadline exceeded",
            )),
        }
    };

    match result {
        Ok(result) => result,
        Err(error) => (response_for_error(request_id, error), None),
    }
}

async fn dispatch_command(
    state: Arc<DcpServerState>,
    request_id: u64,
    command: request::Command,
    outbound: mpsc::Sender<DcpFrame>,
) -> DcpResult<(Response, Option<JoinHandle<()>>)> {
    match command {
        request::Command::ListJobs(_) => {
            let registry = state.registry.clone();
            let jobs = registry_call(move || registry.list()).await?;
            let payload = JobList {
                jobs: jobs.iter().map(wire_job_status).collect(),
            }
            .encode_to_vec();
            Ok((Response::ok(request_id, payload), None))
        }
        request::Command::StartJob(start) => {
            if start.factory_name.trim().is_empty() || start.instance_name.trim().is_empty() {
                return Err(DcpError::response(
                    ResponseStatus::BadRequest,
                    "StartJob requires factory_name and instance_name",
                ));
            }
            let mut spec = state.factories.build(
                &start.factory_name,
                start.instance_name.clone(),
                start.params,
            )?;
            if let Some(cluster) = start.cluster {
                spec = spec.with_cluster_metadata(cluster_metadata_from_wire(cluster)?);
            }
            let registry = state.registry.clone();
            let name = spec.name.clone();
            let status = registry_call(move || {
                registry.submit(spec)?;
                registry.start(name)
            })
            .await?;
            Ok((status_response(request_id, &status), None))
        }
        request::Command::DrainJob(drain) => {
            if drain.cluster {
                let provider = cluster_view_provider(&state)?;
                let timeout = request_timeout(0);
                let payload = provider
                    .drain_cluster_job(drain.name, timeout)
                    .await?
                    .encode_to_vec();
                return Ok((Response::ok(request_id, payload), None));
            }
            let registry = state.registry.clone();
            let status = registry_call(move || registry.drain(drain.name)).await?;
            Ok((status_response(request_id, &status), None))
        }
        request::Command::StopJob(stop) => {
            if stop.cluster {
                let provider = cluster_view_provider(&state)?;
                let timeout = request_timeout(0);
                let payload = provider
                    .stop_cluster_job(stop.name, timeout)
                    .await?
                    .encode_to_vec();
                return Ok((Response::ok(request_id, payload), None));
            }
            let registry = state.registry.clone();
            let status = registry_call(move || registry.stop(stop.name)).await?;
            Ok((status_response(request_id, &status), None))
        }
        request::Command::RestartJob(restart) => {
            if restart.cluster {
                return Err(DcpError::response(
                    ResponseStatus::BadRequest,
                    "cluster restart is not implemented in v0.10; drain/stop and submit again",
                ));
            }
            let registry = state.registry.clone();
            let status = registry_call(move || registry.restart(restart.name)).await?;
            Ok((status_response(request_id, &status), None))
        }
        request::Command::JobStatus(status) => {
            if status.cluster {
                let provider = cluster_view_provider(&state)?;
                let timeout = request_timeout(0);
                let payload = provider
                    .cluster_job_status(status.name, timeout)
                    .await?
                    .encode_to_vec();
                return Ok((Response::ok(request_id, payload), None));
            }
            let registry = state.registry.clone();
            let status = registry_call(move || registry.status(status.name)).await?;
            Ok((status_response(request_id, &status), None))
        }
        request::Command::SubscribeEvents(_) => {
            let registry = state.registry.clone();
            let subscription = spawn_event_subscription(request_id, registry, outbound);
            Ok((Response::ok(request_id, Vec::new()), Some(subscription)))
        }
        request::Command::SubscribeMetrics(metrics) => {
            let interval = if metrics.interval_ms == 0 {
                state.config.metrics_interval
            } else {
                Duration::from_millis(metrics.interval_ms)
            };
            let subscription = spawn_metrics_subscription(
                request_id,
                state.instrumentation.clone(),
                interval.max(Duration::from_millis(1)),
                outbound,
            );
            Ok((Response::ok(request_id, Vec::new()), Some(subscription)))
        }
        request::Command::GetConfig(get) => {
            let value = state
                .config_store
                .lock()
                .expect("DCP config store poisoned")
                .get(&get.key)
                .cloned();
            let payload = ConfigValue {
                key: get.key,
                value: value.clone().unwrap_or_default(),
                existed: value.is_some(),
            }
            .encode_to_vec();
            Ok((Response::ok(request_id, payload), None))
        }
        request::Command::PutConfig(put) => {
            let existed = state
                .config_store
                .lock()
                .expect("DCP config store poisoned")
                .insert(put.key.clone(), put.value.clone())
                .is_some();
            let payload = ConfigValue {
                key: put.key,
                value: put.value,
                existed,
            }
            .encode_to_vec();
            Ok((Response::ok(request_id, payload), None))
        }
        request::Command::ListClusterJobs(list) => {
            let provider = cluster_view_provider(&state)?;
            let timeout = request_timeout(list.timeout_ms);
            let payload = provider.list_cluster_jobs(timeout).await?.encode_to_vec();
            Ok((Response::ok(request_id, payload), None))
        }
        request::Command::ClusterNodeInfo(info) => {
            let provider = cluster_view_provider(&state)?;
            let timeout = request_timeout(info.timeout_ms);
            let payload = provider.cluster_node_info(timeout).await?.encode_to_vec();
            Ok((Response::ok(request_id, payload), None))
        }
        request::Command::SubmitClusterJob(submit) => {
            let provider = cluster_view_provider(&state)?;
            let timeout = request_timeout(submit.timeout_ms);
            let payload = provider
                .submit_cluster_job(submit, timeout)
                .await?
                .encode_to_vec();
            Ok((Response::ok(request_id, payload), None))
        }
        request::Command::RememberClusterAssignment(remember) => {
            let provider = cluster_view_provider(&state)?;
            provider.remember_cluster_assignment(remember).await?;
            Ok((Response::ok(request_id, Vec::new()), None))
        }
        request::Command::AllocateShard(allocate) => {
            let provider = sharding_view_provider(&state)?;
            let timeout = request_timeout(allocate.timeout_ms);
            let payload = provider
                .allocate_shard(allocate, timeout)
                .await?
                .encode_to_vec();
            Ok((Response::ok(request_id, payload), None))
        }
        request::Command::RememberShardAllocations(remember) => {
            let provider = sharding_view_provider(&state)?;
            provider.remember_shard_allocations(remember).await?;
            Ok((Response::ok(request_id, Vec::new()), None))
        }
        request::Command::GetShardAllocations(get) => {
            let provider = sharding_view_provider(&state)?;
            let timeout = request_timeout(0);
            let payload = provider
                .get_shard_allocations(get.type_name, timeout)
                .await?
                .encode_to_vec();
            Ok((Response::ok(request_id, payload), None))
        }
        request::Command::ForwardShardEnvelopes(batch) => {
            let provider = sharding_view_provider(&state)?;
            let timeout = request_timeout(0);
            let payload = provider
                .forward_shard_envelopes(batch, timeout)
                .await?
                .encode_to_vec();
            Ok((Response::ok(request_id, payload), None))
        }
        request::Command::CompleteShardingAsk(reply) => {
            let provider = sharding_view_provider(&state)?;
            provider.complete_sharding_ask(reply).await?;
            Ok((Response::ok(request_id, Vec::new()), None))
        }
        request::Command::OpenShardPipe(_) => Err(DcpError::response(
            ResponseStatus::BadRequest,
            "OpenShardPipe must be handled by the DCP connection loop",
        )),
    }
}

fn cluster_view_provider(state: &DcpServerState) -> DcpResult<Arc<dyn ClusterViewProvider>> {
    state
        .cluster_view
        .lock()
        .expect("DCP cluster view provider poisoned")
        .clone()
        .ok_or_else(|| {
            DcpError::response(
                ResponseStatus::Failed,
                "cluster view is not configured on this datum-agent",
            )
        })
}

fn sharding_view_provider(state: &DcpServerState) -> DcpResult<Arc<dyn ShardingViewProvider>> {
    state
        .sharding_view
        .lock()
        .expect("DCP sharding view provider poisoned")
        .clone()
        .ok_or_else(|| {
            DcpError::response(
                ResponseStatus::Failed,
                "cluster sharding is not configured on this datum-agent",
            )
        })
}

fn request_timeout(timeout_ms: u64) -> Duration {
    if timeout_ms == 0 {
        DEFAULT_CLUSTER_REQUEST_TIMEOUT
    } else {
        Duration::from_millis(timeout_ms)
    }
}

async fn registry_call<T, F>(call: F) -> DcpResult<T>
where
    T: Send + 'static,
    F: FnOnce() -> AgentResult<T> + Send + 'static,
{
    Ok(tokio::task::spawn_blocking(call).await??)
}

fn status_response(request_id: u64, status: &RegistryJobStatus) -> Response {
    Response::ok(request_id, wire_job_status(status).encode_to_vec())
}

fn response_for_error(request_id: u64, error: DcpError) -> Response {
    match error {
        DcpError::Response { status, message } => Response::error(request_id, status, message),
        DcpError::Agent(error) => {
            let status = match &error {
                AgentError::InvalidJobName => ResponseStatus::BadRequest,
                AgentError::JobNotFound(_) => ResponseStatus::NotFound,
                AgentError::JobAlreadyExists(_) | AgentError::JobAlreadyRunning(_) => {
                    ResponseStatus::Conflict
                }
                AgentError::DrainUnsupported(_)
                | AgentError::JobNotRunning(_)
                | AgentError::RestartLimitExceeded(_) => ResponseStatus::Failed,
                AgentError::RegistryClosed | AgentError::Actor(_) | AgentError::Stream(_) => {
                    ResponseStatus::Failed
                }
            };
            Response::error(request_id, status, error.to_string())
        }
        other => Response::error(request_id, ResponseStatus::Failed, other.to_string()),
    }
}

fn spawn_event_subscription(
    subscription_id: u64,
    registry: JobRegistryHandle,
    outbound: mpsc::Sender<DcpFrame>,
) -> JoinHandle<()> {
    tokio::task::spawn_blocking(move || {
        let Ok(queue) = registry.events().run_with(Sink::queue()) else {
            return;
        };
        while let Ok(Some(event)) = queue.pull() {
            let frame = DcpFrame::event(subscription_id, wire_event(&event));
            if outbound.blocking_send(frame).is_err() {
                break;
            }
        }
    })
}

fn spawn_metrics_subscription(
    subscription_id: u64,
    instrumentation: StreamInstrumentationRegistry,
    interval: Duration,
    outbound: mpsc::Sender<DcpFrame>,
) -> JoinHandle<()> {
    tokio::spawn(async move {
        let mut ticker = tokio::time::interval(interval);
        loop {
            ticker.tick().await;
            let sample = MetricSample {
                timestamp_ms: system_time_ms(SystemTime::now()),
                streams: instrumentation
                    .snapshots()
                    .iter()
                    .map(wire_stream_metric)
                    .collect(),
            };
            let frame = DcpFrame::metric(subscription_id, sample);
            match outbound.try_send(frame) {
                Ok(()) => {}
                Err(mpsc::error::TrySendError::Full(_)) => {}
                Err(mpsc::error::TrySendError::Closed(_)) => break,
            }
        }
    })
}

pub(crate) fn wire_job_status(status: &RegistryJobStatus) -> WireJobStatus {
    let now = Instant::now();
    let cluster = status.cluster.as_ref();
    WireJobStatus {
        name: status.name.clone(),
        job_id: status.job_id.0,
        state: format!("{:?}", status.state),
        desired_state: format!("{:?}", status.desired_state),
        generation: status.generation,
        starts_total: status.starts_total,
        restarts_total: status.restarts_total,
        last_start_at_ms: status.last_start_at.map(system_time_ms),
        last_exit_at_ms: status.last_exit_at.map(system_time_ms),
        last_exit_reason: status
            .last_exit_reason
            .as_ref()
            .map(exit_reason_text)
            .unwrap_or_default(),
        backoff_remaining_ms: status
            .backoff_until
            .map(|deadline| duration_ms(deadline.saturating_duration_since(now))),
        drain_remaining_ms: status
            .drain_deadline
            .map(|deadline| duration_ms(deadline.saturating_duration_since(now))),
        drain_supported: status.drain_supported,
        active_streams: status.active_streams.map(|streams| streams as u64),
        cluster_job: cluster.is_some(),
        factory_name: cluster
            .map(|metadata| metadata.factory_name.clone())
            .unwrap_or_default(),
        placement: cluster.map(|metadata| wire_placement_spec(&metadata.placement)),
        coordinator_node_id: cluster
            .map(|metadata| metadata.coordinator_node.clone())
            .unwrap_or_default(),
        placement_node_id: cluster
            .map(|metadata| metadata.assigned_node.clone())
            .unwrap_or_default(),
        placement_generation: cluster
            .map(|metadata| metadata.placement_generation)
            .unwrap_or(0),
        placement_history: cluster
            .map(|metadata| {
                metadata
                    .history
                    .iter()
                    .map(wire_placement_history)
                    .collect()
            })
            .unwrap_or_default(),
        params: cluster
            .map(|metadata| metadata.params.clone().into_iter().collect())
            .unwrap_or_default(),
    }
}

pub(crate) fn wire_placement_spec(spec: &RegistryPlacementSpec) -> PlacementSpec {
    let (strategy, pinned_node_id) = match &spec.strategy {
        RegistryPlacementStrategy::LeastJobs => (PlacementStrategy::LeastJobs, String::new()),
        RegistryPlacementStrategy::Pinned { node_id } => {
            (PlacementStrategy::Pinned, node_id.clone())
        }
    };
    PlacementSpec {
        role_constraint: spec.role_constraint.clone().unwrap_or_default(),
        strategy: strategy as i32,
        pinned_node_id,
    }
}

pub(crate) fn placement_spec_from_wire(
    spec: Option<PlacementSpec>,
) -> DcpResult<RegistryPlacementSpec> {
    let spec = spec.unwrap_or(PlacementSpec {
        role_constraint: String::new(),
        strategy: PlacementStrategy::LeastJobs as i32,
        pinned_node_id: String::new(),
    });
    let role_constraint = normalize_optional_string(spec.role_constraint);
    let strategy =
        match PlacementStrategy::try_from(spec.strategy).unwrap_or(PlacementStrategy::LeastJobs) {
            PlacementStrategy::LeastJobs => RegistryPlacementStrategy::LeastJobs,
            PlacementStrategy::Pinned => {
                let node_id = spec.pinned_node_id.trim().to_owned();
                if node_id.is_empty() {
                    return Err(DcpError::response(
                        ResponseStatus::BadRequest,
                        "Pinned placement requires pinned_node_id",
                    ));
                }
                RegistryPlacementStrategy::Pinned { node_id }
            }
        };
    Ok(RegistryPlacementSpec {
        role_constraint,
        strategy,
    })
}

fn wire_placement_history(
    history: &ClusterPlacementHistory,
) -> crate::dcp::proto::ClusterPlacementHistory {
    crate::dcp::proto::ClusterPlacementHistory {
        generation: history.generation,
        from_node_id: history.from_node.clone().unwrap_or_default(),
        to_node_id: history.to_node.clone(),
        reason: history.reason.clone(),
        timestamp_ms: system_time_ms(history.timestamp),
    }
}

fn placement_history_from_wire(
    history: crate::dcp::proto::ClusterPlacementHistory,
) -> ClusterPlacementHistory {
    ClusterPlacementHistory {
        generation: history.generation,
        from_node: normalize_optional_string(history.from_node_id),
        to_node: history.to_node_id,
        reason: history.reason,
        timestamp: UNIX_EPOCH + Duration::from_millis(history.timestamp_ms),
    }
}

pub(crate) fn cluster_metadata_from_wire(
    cluster: ClusterJobStart,
) -> DcpResult<ClusterJobMetadata> {
    Ok(ClusterJobMetadata {
        factory_name: cluster.factory_name,
        params: cluster.params.into_iter().collect::<BTreeMap<_, _>>(),
        placement: placement_spec_from_wire(cluster.placement)?,
        coordinator_node: cluster.coordinator_node_id,
        assigned_node: cluster.assigned_node_id,
        placement_generation: cluster.placement_generation,
        history: cluster
            .history
            .into_iter()
            .map(placement_history_from_wire)
            .collect(),
    })
}

pub(crate) fn wire_cluster_job_start(metadata: &ClusterJobMetadata) -> ClusterJobStart {
    ClusterJobStart {
        factory_name: metadata.factory_name.clone(),
        params: metadata.params.clone().into_iter().collect(),
        placement: Some(wire_placement_spec(&metadata.placement)),
        coordinator_node_id: metadata.coordinator_node.clone(),
        assigned_node_id: metadata.assigned_node.clone(),
        placement_generation: metadata.placement_generation,
        history: metadata
            .history
            .iter()
            .map(wire_placement_history)
            .collect(),
    }
}

fn wire_event(event: &JobEvent) -> Event {
    let (kind, detail) = event_kind_text(&event.kind);
    Event {
        sequence: event.sequence,
        timestamp_ms: system_time_ms(event.timestamp),
        name: event.name.clone(),
        job_id: event.job_id.0,
        generation: event.generation,
        kind,
        detail,
    }
}

fn wire_stream_metric(snapshot: &StreamInstrumentationSnapshot) -> StreamMetric {
    StreamMetric {
        id: snapshot.id.get(),
        name: snapshot.name.clone(),
        elements_through: snapshot.elements_through,
        restarts: snapshot.restarts,
        state: instrumentation_state_text(snapshot.state).to_owned(),
        started_at_ms: system_time_ms(snapshot.started_at),
        state_changed_at_ms: system_time_ms(snapshot.state_changed_at),
        finished_at_ms: snapshot.finished_at.map(system_time_ms),
        uptime_ms: duration_ms(snapshot.uptime),
    }
}

fn instrumentation_state_text(state: StreamInstrumentationState) -> &'static str {
    match state {
        StreamInstrumentationState::Running => "Running",
        StreamInstrumentationState::Draining => "Draining",
        StreamInstrumentationState::Completed => "Completed",
        StreamInstrumentationState::Failed => "Failed",
    }
}

fn event_kind_text(kind: &JobEventKind) -> (String, String) {
    match kind {
        JobEventKind::Submitted => ("Submitted".to_owned(), String::new()),
        JobEventKind::Started => ("Started".to_owned(), String::new()),
        JobEventKind::Failed { reason } => ("Failed".to_owned(), exit_reason_text(reason)),
        JobEventKind::RestartScheduled { delay } => (
            "RestartScheduled".to_owned(),
            format!("delay_ms={}", duration_ms(*delay)),
        ),
        JobEventKind::Restarted {
            previous_generation,
        } => (
            "Restarted".to_owned(),
            format!("previous_generation={previous_generation}"),
        ),
        JobEventKind::Draining => ("Draining".to_owned(), String::new()),
        JobEventKind::Drained => ("Drained".to_owned(), String::new()),
        JobEventKind::Stopped { reason } => ("Stopped".to_owned(), exit_reason_text(reason)),
        JobEventKind::Completed => ("Completed".to_owned(), String::new()),
    }
}

fn exit_reason_text(reason: &JobExitReason) -> String {
    match reason {
        JobExitReason::Completed => "Completed".to_owned(),
        JobExitReason::Failed(error) => format!("Failed({error})"),
        JobExitReason::Drained => "Drained".to_owned(),
        JobExitReason::Stopped => "Stopped".to_owned(),
        JobExitReason::DrainTimedOut => "DrainTimedOut".to_owned(),
    }
}

fn system_time_ms(time: SystemTime) -> u64 {
    time.duration_since(UNIX_EPOCH)
        .map(duration_ms)
        .unwrap_or(0)
}

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

fn normalize_optional_string(value: String) -> Option<String> {
    let value = value.trim().to_owned();
    if value.is_empty() { None } else { Some(value) }
}

fn ensure_loopback(addr: SocketAddr) -> DcpResult<()> {
    if addr.ip().is_loopback() {
        return Ok(());
    }
    Err(DcpError::Protocol(format!(
        "plaintext DCP TCP listener must bind loopback, got {addr}"
    )))
}

fn client_bind_addr(remote_addr: SocketAddr) -> SocketAddr {
    if remote_addr.is_ipv6() {
        SocketAddr::new(IpAddr::V6(Ipv6Addr::UNSPECIFIED), 0)
    } else {
        SocketAddr::new(IpAddr::V4(Ipv4Addr::UNSPECIFIED), 0)
    }
}

pub(crate) async fn connect_quic_stream(
    addr: SocketAddr,
    server_name: &str,
    client_config: quinn::ClientConfig,
) -> DcpResult<(
    quinn::Endpoint,
    quinn::Connection,
    quinn::RecvStream,
    quinn::SendStream,
)> {
    let mut endpoint = quinn::Endpoint::client(client_bind_addr(addr))?;
    endpoint.set_default_client_config(client_config);
    let connection = endpoint
        .connect(addr, server_name)
        .map_err(|error| DcpError::Protocol(error.to_string()))?
        .await
        .map_err(|error| DcpError::Protocol(error.to_string()))?;
    let (send, recv) = connection
        .open_bi()
        .await
        .map_err(|error| DcpError::Protocol(error.to_string()))?;
    Ok((endpoint, connection, recv, send))
}