motorcortex_rust/client/

subscribe.rs

use crate::client::subscription::ReadOnlySubscription;
use crate::client::{receive_message, Parameters, Subscription};
use crate::connection::{ConnectionManager, ConnectionOptions};
use crate::error::{MotorcortexError, Result};
use crate::{Request, StatusCode};
use std::collections::HashMap;
use std::sync::atomic::{AtomicBool, Ordering};
use std::sync::{Arc, RwLock};
use std::thread;

/// Real-time streaming client that receives parameter updates from the server.
///
/// # Thread Safety
///
/// `Subscribe` is `Send` but not `Sync`. You can move it to another thread,
/// but you cannot share it between threads by reference. The internal receive
/// thread is managed automatically — no external synchronization is needed
/// for subscription reads via `ReadOnlySubscription`, which is independently
/// thread-safe.
pub struct Subscribe {
    connection_data: ConnectionManager,
    receive_thread: Option<thread::JoinHandle<()>>,
    stop_signal: Arc<AtomicBool>,
    active_subscriptions: Arc<RwLock<HashMap<u32, Arc<RwLock<Subscription>>>>>,
}

impl Subscribe {
    const ID_BYTE_SIZE: usize = 3;

    /// Creates a new `Subscribe` instance.
    ///
    /// No socket is opened until `connect()` is called.
    pub fn new() -> Self {
        Self {
            connection_data: ConnectionManager::new(),
            receive_thread: None,
            stop_signal: Arc::new(AtomicBool::new(false)),
            active_subscriptions: Arc::new(RwLock::new(HashMap::new())),
        }
    }

    /// Establishes a connection to a specified server URL using the given configuration options.
    ///
    /// Opens a SUB socket, connects to the server, and spawns a background
    /// receive thread for processing incoming subscription updates.
    ///
    /// # Arguments:
    /// * `url` - The server's address (e.g., "wss://127.0.0.1:5569").
    /// * `connection_options` - The connection settings, including TLS certificate and timeouts.
    ///
    /// # Errors:
    /// Returns `MotorcortexError::Connection` if the socket cannot be opened or the connection fails.
    pub fn connect(&mut self, url: &str, connection_options: ConnectionOptions) -> Result<()> {
        let res = self
            .connection_data
            .connect(url, connection_options, nng_c_sys::nng_sub0_open);
        let stop_signal_clone = Arc::clone(&self.stop_signal);
        let active_subscriptions_clone = Arc::clone(&self.active_subscriptions);

        let sock = self.connection_data.sock.unwrap();
        self.receive_thread = Some(thread::spawn(move || {
            const HEADER_LEN: usize = 4usize;
            loop {
                if stop_signal_clone.load(Ordering::Relaxed) {
                    println!("Receive thread stopping...");
                    break;
                }

                let buffer = match receive_message(&sock) {
                    Ok(buffer) => buffer,
                    Err(_) => continue,
                };

                if buffer.len() > HEADER_LEN {
                    let id =
                        (buffer[0] as u32) | ((buffer[1] as u32) << 8) | ((buffer[2] as u32) << 16);
                    let subscriptions = active_subscriptions_clone.write().unwrap();
                    if let Some(sub) = subscriptions.get(&id) {
                        let mut sub = sub.write().unwrap();
                        sub.update(buffer);
                    }
                }
            }
        }));

        res
    }

    /// Disconnects the current connection and frees the associated resources.
    ///
    /// # Errors:
    /// Returns `MotorcortexError::Connection` if the disconnection fails.
    pub fn disconnect(&mut self) -> Result<()> {
        self.stop_signal.store(true, Ordering::Relaxed);
        let res = self.connection_data.disconnect();
        self.receive_thread.take().unwrap().join().unwrap();
        res
    }

    /// Removes a subscription and its associated server-side group.
    ///
    /// # Arguments:
    /// * `request` - The Request client used to remove the group on the server.
    /// * `id` - The subscription ID to remove.
    ///
    /// # Errors:
    /// Returns `MotorcortexError::Subscription` if the NNG unsubscribe fails.
    pub fn unsubscribe(&mut self, request: &Request, id: u32) -> Result<()> {
        let sub = self.active_subscriptions.write().unwrap().remove(&id);
        if let Some(sub) = sub {
            // Remove subscription hash from nng
            let bytes = id.to_le_bytes();
            let rv = unsafe {
                nng_c_sys::nng_setopt(
                    self.connection_data.sock.unwrap(),
                    nng_c_sys::NNG_OPT_SUB_UNSUBSCRIBE.as_ptr() as *const core::ffi::c_char,
                    bytes.as_ptr() as *const std::ffi::c_void,
                    Subscribe::ID_BYTE_SIZE,
                )
            };
            if rv != 0 {
                return Err(MotorcortexError::Subscription(format!(
                    "Failed to unsubscribe from NNG. Error code: {}",
                    rv
                )));
            }

            // Send remove group message
            let name = sub.read().unwrap().name();
            request.remove_group(&name)?;
        }
        Ok(())
    }

    /// Creates a new subscription for the specified parameters.
    ///
    /// # Arguments:
    /// * `request` - The Request client used to create the group on the server.
    /// * `paths` - The parameter paths to subscribe to.
    /// * `group_name` - An alias for the subscription group.
    /// * `frequency_divider` - Update rate divider relative to the server's base frequency.
    ///
    /// # Errors:
    /// Returns `MotorcortexError::Subscription` if the group creation or NNG subscribe fails.
    pub fn subscribe<I>(
        &mut self,
        request: &Request,
        paths: I,
        group_name: &str,
        frequency_divider: u32,
    ) -> Result<ReadOnlySubscription>
    where
        I: Parameters,
    {
        let sub = request.create_group(paths, group_name, frequency_divider)?;

        if sub.status != StatusCode::Ok as i32 {
            return Err(MotorcortexError::Subscription(format!(
                "Failed to create group, status: {}",
                sub.status
            )));
        }

        let id = sub.id;
        let bytes = id.to_le_bytes();
        let rv = unsafe {
            nng_c_sys::nng_setopt(
                self.connection_data.sock.unwrap(),
                nng_c_sys::NNG_OPT_SUB_SUBSCRIBE.as_ptr() as *const core::ffi::c_char,
                bytes.as_ptr() as *const std::ffi::c_void,
                Subscribe::ID_BYTE_SIZE,
            )
        };
        if rv != 0 {
            return Err(MotorcortexError::Subscription(format!(
                "Failed to subscribe via NNG. Error code: {}",
                rv
            )));
        }

        // Create the subscription and store it in the HashMap
        let subscription = Arc::new(RwLock::new(Subscription::new(sub)));
        let mut subscriptions = self.active_subscriptions.write().unwrap();
        subscriptions.insert(id, subscription.clone());

        Ok(ReadOnlySubscription::new(subscription))
    }
}

impl Drop for Subscribe {
    fn drop(&mut self) {
        self.stop_signal.store(true, Ordering::Relaxed);
        if let Some(handle) = self.receive_thread.take() {
            let _ = handle.join();
        }
        // ConnectionManager::drop() handles socket and TLS cleanup
    }
}