wowlab_node/

realtime.rs

//! Realtime subscriptions for node updates and chunk assignments.

use std::time::Duration;

use serde::Deserialize;
use tokio::{runtime::Handle, sync::mpsc};
use tokio_util::sync::CancellationToken;
use wowlab_centrifuge::{Client, ClientConfig, ClientEvent, SubscriptionConfig, SubscriptionEvent};
use wowlab_common::{NodePublicKey, NodeRealtimeMessage, RuntimeChunkPayload};
use wowlab_types::sensitive::Sensitive;

use crate::sentinel::SentinelClient;

const CONNECT_TIMEOUT: Duration =
    Duration::from_secs(wowlab_types::constants::HTTP_CONNECT_TIMEOUT_SECS);
const EVENT_CHANNEL_SIZE: usize = 32;

#[derive(Clone, Debug)]
#[non_exhaustive]
pub(crate) enum RealtimeEvent {
    NodeUpdated {
        name: String,
        total_cores: i32,
        max_parallel: i32,
    },
    ChunkAssigned(RuntimeChunkPayload),
    Connected,
    Disconnected,
    Error(String),
}

#[derive(Debug)]
pub(crate) struct NodeRealtime {
    config: ClientConfig,
}

pub(crate) struct RealtimeConfig {
    pub url: String,
    pub token: Sensitive<String>,
    pub name: String,
    pub version: String,
    pub sentinel: SentinelClient,
}

fn log_connection_error(error: &str) {
    tracing::warn!(error, "Realtime connection error");
}

impl NodeRealtime {
    pub(crate) fn new(settings: RealtimeConfig) -> Self {
        let config = ClientConfig::new(settings.url, settings.token)
            .name(settings.name)
            .version(settings.version)
            .get_token(move || {
                let sentinel = settings.sentinel.clone();

                async move {
                    sentinel
                        .refresh_token()
                        .await
                        .map_err(|error| wowlab_centrifuge::Error::protocol(error.to_string()))
                }
            });

        Self { config }
    }

    pub(crate) fn subscribe(
        &self,
        public_key: &NodePublicKey,
        handle: &Handle,
        shutdown: CancellationToken,
    ) -> mpsc::Receiver<RealtimeEvent> {
        let (tx, rx) = mpsc::channel(EVENT_CHANNEL_SIZE);
        let config = self.config.clone();
        let pk_str = public_key.to_string();

        handle.spawn(async move {
            let client = Client::new(config);
            let mut events = client.events().await;

            client.connect();

            let connected = tokio::time::timeout(CONNECT_TIMEOUT, async {
                while let Some(event) = events.recv().await {
                    if matches!(event, ClientEvent::Connected(_)) {
                        return true;
                    }

                    if let ClientEvent::Error(error) = event {
                        log_connection_error(&error);
                    }

                    tokio::task::yield_now().await;
                }

                false
            })
            .await;

            if connected != Ok(true) {
                let _ = tx
                    .send(RealtimeEvent::Error("Connection timeout".into()))
                    .await;

                client.disconnect();

                return;
            }

            if let Err(e) = run_subscriptions(&client, &pk_str, &tx, shutdown).await {
                let _ = tx.send(RealtimeEvent::Error(e.to_string())).await;
            }

            let _ = tx.send(RealtimeEvent::Disconnected).await;

            client.disconnect();
            tracing::debug!("Realtime subscription shut down");
        });

        rx
    }
}

async fn run_subscriptions(
    client: &Client,
    public_key: &str,
    tx: &mpsc::Sender<RealtimeEvent>,
    shutdown: CancellationToken,
) -> Result<(), wowlab_centrifuge::Error> {
    // docref:start realtime-node-subscribe
    let mut node_sub = client
        .subscribe(SubscriptionConfig::new(format!("nodes:{public_key}")))
        .await?;
    let mut chunks_sub = client
        .subscribe(SubscriptionConfig::new(format!("chunks:{public_key}")))
        .await?;
    let mut presence_sub = client
        .subscribe(SubscriptionConfig::new("nodes:online").join_leave(true))
        .await?;
    // docref:end realtime-node-subscribe

    tracing::debug!(public_key, "Subscribed to node realtime channels");
    let _ = tx.send(RealtimeEvent::Connected).await;

    loop {
        tokio::select! {
            () = shutdown.cancelled() => break,
            Some(event) = node_sub.recv() => {
                if let Some(event) = extract_node_publication(&event) {
                    let _ = tx.send(event).await;
                }
            }
            Some(event) = chunks_sub.recv() => {
                if let Some(payload) = extract_publication::<RuntimeChunkPayload>(&event) {
                    let _ = tx.send(RealtimeEvent::ChunkAssigned(payload)).await;
                }
            }
            Some(event) = presence_sub.recv() => {
                log_presence_event(&event);
            }
            else => break,
        }
    }

    Ok(())
}

fn extract_node_publication(event: &SubscriptionEvent) -> Option<RealtimeEvent> {
    let SubscriptionEvent::Publication(publication) = event else {
        return None;
    };

    parse_node_publication(&publication.data)
}

fn parse_node_publication(data: &[u8]) -> Option<RealtimeEvent> {
    match parse_json::<NodeRealtimeMessage>(data)? {
        NodeRealtimeMessage::ConfigurationUpdated {
            name,
            total_cores,
            max_parallel,
        } => Some(RealtimeEvent::NodeUpdated {
            name,
            total_cores,
            max_parallel,
        }),
        NodeRealtimeMessage::RefinementUpdated { ids } => {
            tracing::debug!(?ids, "Ignoring irrelevant node refinement update");

            None
        }
        _ => {
            tracing::debug!("Ignoring irrelevant node realtime message");

            None
        }
    }
}

fn extract_publication<T>(event: &SubscriptionEvent) -> Option<T>
where
    T: for<'de> Deserialize<'de>,
{
    let SubscriptionEvent::Publication(pub_) = event else {
        return None;
    };

    parse_json(&pub_.data)
}

fn log_presence_event(event: &SubscriptionEvent) {
    match event {
        SubscriptionEvent::Join(info) => {
            tracing::debug!(user = %info.user, client = %info.client, "Node joined");
        }
        SubscriptionEvent::Leave(info) => {
            tracing::debug!(user = %info.user, client = %info.client, "Node left");
        }
        _ => {}
    }
}

fn parse_json<T>(data: &[u8]) -> Option<T>
where
    T: for<'de> Deserialize<'de>,
{
    serde_json::from_slice(data)
        .inspect_err(|e| {
            tracing::warn!(
                error = %e,
                payload_bytes = data.len(),
                "Failed to parse publication"
            );
        })
        .ok()
}

#[cfg(test)]
mod tests {
    use googletest::prelude::*;
    use wowlab_common::{RuntimeChunkId, RuntimeWorkFidelity, RuntimeWorkItemKind};

    use super::*;

    #[derive(Clone, Debug, Eq, PartialEq)]
    struct ObservedNodeState {
        name: String,
        total_cores: u32,
        max_parallel: u32,
    }

    fn apply_node_publication(data: &[u8], state: &mut ObservedNodeState) -> bool {
        let Some(RealtimeEvent::NodeUpdated {
            name,
            total_cores,
            max_parallel,
        }) = parse_node_publication(data)
        else {
            return false;
        };

        state.name = name;
        state.total_cores = total_cores.unsigned_abs();
        state.max_parallel = max_parallel.unsigned_abs();

        true
    }

    fn configured_state() -> ObservedNodeState {
        ObservedNodeState {
            name: "Configured Worker".to_string(),
            total_cores: 16,
            max_parallel: 8,
        }
    }

    fn valid_configuration() -> &'static [u8] {
        br#"{
            "type": "configuration_updated",
            "payload": {
                "name": "Updated Worker",
                "totalCores": 32,
                "maxParallel": 12
            }
        }"#
    }

    #[gtest]
    fn parses_runtime_chunk_payload_from_centrifuge_json() -> Result<()> {
        let json = br#"{
            "jobId": "job-abc",
            "chunkId": 42,
            "workContextHash": "ab12cd34ab12cd34ab12cd34ab12cd34ab12cd34ab12cd34ab12cd34ab12cd34",
            "claimToken": "claim-1",
            "workItems": [
                {
                    "itemId": 1,
                    "kind": { "type": "tournament", "permIdx": 9 },
                    "tag": 9,
                    "iterations": 2000,
                    "seedOffset": 18000,
                    "fidelity": "dpsOnly"
                },
                {
                    "itemId": 2,
                    "kind": { "type": "base" },
                    "tag": 0,
                    "iterations": 10000,
                    "seedOffset": 0,
                    "fidelity": "full"
                }
            ]
        }"#;

        let payload: RuntimeChunkPayload = parse_json(json).or_fail()?;

        verify_that!(payload.job_id, eq("job-abc"))?;
        verify_that!(payload.chunk_id, eq(RuntimeChunkId::new(42)))?;
        verify_that!(payload.claim_token.as_str(), eq("claim-1"))?;
        verify_that!(payload.work_items, len(eq(2)))?;

        let first = &payload.work_items[0];

        verify_true!(matches!(
            first.kind,
            RuntimeWorkItemKind::Tournament { perm_idx: 9 }
        ))?;
        verify_that!(first.fidelity, eq(RuntimeWorkFidelity::DpsOnly))?;

        let second = &payload.work_items[1];

        verify_true!(matches!(second.kind, RuntimeWorkItemKind::Base))?;

        verify_that!(second.fidelity, eq(RuntimeWorkFidelity::Full))
    }

    #[gtest]
    fn rejects_malformed_chunk_payload() -> Result<()> {
        let json = br#"{ "jobId": "x" }"#;
        let parsed: Option<RuntimeChunkPayload> = parse_json(json);

        verify_that!(parsed, none())
    }

    #[gtest]
    fn rejects_sentinel_refinement_payload_as_node_configuration() -> Result<()> {
        let json = br#"{
            "type": "updated",
            "payload": { "ids": ["node-public-key"] }
        }"#;

        let parsed = parse_node_publication(json);

        verify_that!(parsed, none())
    }

    #[gtest]
    fn irrelevant_variant_preserves_state_and_next_publication_is_processed() -> Result<()> {
        let refinement = br#"{
            "type": "refinement_updated",
            "payload": { "ids": ["node-public-key"] }
        }"#;
        let mut state = configured_state();
        let before = state.clone();

        verify_false!(apply_node_publication(refinement, &mut state))?;
        verify_that!(state, eq(&before))?;
        verify_true!(apply_node_publication(valid_configuration(), &mut state))?;

        verify_that!(state.name, eq("Updated Worker"))
    }

    #[gtest]
    fn malformed_payload_preserves_state_and_next_publication_is_processed() -> Result<()> {
        let malformed = br#"{
            "type": "configuration_updated",
            "payload": { "name": "Incomplete Worker" }
        }"#;
        let mut state = configured_state();
        let before = state.clone();

        verify_false!(apply_node_publication(malformed, &mut state))?;
        verify_that!(state, eq(&before))?;
        verify_true!(apply_node_publication(valid_configuration(), &mut state))?;

        verify_that!(state.name, eq("Updated Worker"))
    }
}