romp 0.5.2

STOMP server and WebSockets platform
Documentation
//! DestinationServer is the in memory representation of queues and topics.
//! Holds all the servers Destinations, which contain messages to be sent out and subscriber information.
//! Holds all global state outside `CONFIG` and tokio.

use std::collections::hash_map::Iter;
use std::collections::HashMap;
use std::io::Write;
use std::sync::{Arc, RwLock};
use std::sync::atomic::{AtomicI64, AtomicUsize, Ordering, AtomicBool};
use std::ops::Deref;

use chrono::Utc;
use futures::task::Task;

use log::*;
use log::Level::Debug;

use crate::init::NGIN;
use crate::init::CONFIG;
use crate::config;
use crate::config::config::ServerConfig;
use crate::workflow::destination::destination::Destination;
use crate::message::response::{get_response_stats_msg, get_response_cc_msg};
use crate::message::stomp_message::StompMessage;
use crate::errors::ServerError;



// const STATS_TOPIC: &str = "/xtomp/stat";

pub struct DestinationServer {
    // state
    message_id : AtomicUsize,
    session_id: AtomicUsize,
    destination_id: AtomicUsize,
    destinations: HashMap<String, Arc<RwLock<Destination>>>,
    destinations_by_id: HashMap<usize, Arc<RwLock<Destination>>>,

    // life cycle
    shutdown: AtomicBool,
    task_set: AtomicBool,
    task: RwLock<Option<Task>>,

    // stats
    pub uptime: AtomicI64,
    pub connections_count: AtomicUsize,
    pub connections_total: AtomicUsize,
    pub messages_total: AtomicUsize,
}

impl DestinationServer {

    pub(crate) fn new() -> DestinationServer {
        DestinationServer {
            message_id: AtomicUsize::new(DestinationServer::init_id()),
            session_id: AtomicUsize::new(1),
            destination_id: AtomicUsize::new(1),
            destinations: HashMap::new(),
            destinations_by_id: HashMap::new(),
            shutdown: AtomicBool::new(false),
            task_set: AtomicBool::new(false),
            task: RwLock::new(None),
            uptime: AtomicI64::new(Utc::now().timestamp_millis()),
            connections_count: AtomicUsize::new(0),
            connections_total: AtomicUsize::new(0),
            messages_total: AtomicUsize::new(0),
        }
    }

    #[cfg(target_pointer_width = "64")]
    fn init_id() -> usize {
        Utc::now().timestamp_millis() as usize
    }

    #[cfg(target_pointer_width = "32")]
    fn init_id()  -> usize {
        1
    }

    // STATE

    pub(crate) fn new_message_id(&self) -> usize {
        // this wraps on overflow
        self.messages_total.fetch_add(1, Ordering::SeqCst);
        return self.message_id.fetch_add(1, Ordering::SeqCst);
    }
    
    pub(crate) fn new_session(&self) -> usize {
        self.connections_count.fetch_add(1, Ordering::SeqCst);
        self.connections_total.fetch_add(1, Ordering::SeqCst);
        return self.session_id.fetch_add(1, Ordering::SeqCst);
    }

    pub(crate) fn drop_session(&self) -> usize {
        return self.connections_count.fetch_sub(1, Ordering::SeqCst);
    }

    fn create_pid_file(&self) -> std::io::Result<()> {
        let pid = std::process::id().to_string();
        let mut file = std::fs::File::create(&CONFIG.pid_file)?;
        file.write_all(pid.as_bytes())?;
        Ok(())
    }

    fn init_destination(&self, destination: &config::config::Destination) -> Arc<RwLock<Destination>> {
        Arc::new(RwLock::new(Destination::new(
            destination.name.clone(),
            self.destination_id.fetch_sub(1, Ordering::SeqCst),
            destination
        )))
    }

    pub fn find_destination(&self, name: &str) -> Option<&Arc<RwLock<Destination>>> {
        return self.destinations.get(&String::from(name));
    }

    /// publish message without security checks, this is an internal API and should not be called from filters.
    pub fn publish_message_internal(&self, name: &str, message: StompMessage) -> Result<usize, ServerError> {
        if let Some(destination) = self.destinations.get(&String::from(name)) {
            return destination.write().unwrap().push_owned(message);
        }
        Err(ServerError::DestinationUnknown)
    }

    pub fn find_destination_by_id(&self, id: &usize) -> Option<&Arc<RwLock<Destination>>> {
        return self.destinations_by_id.get(id);
    }

    pub fn iter(&self) -> Iter<String, Arc<RwLock<Destination>>> {
        self.destinations.iter()
    }

    // LIFECYCLE

    /// called once on boot
    pub(crate) fn init(&mut self, config: &ServerConfig) {
        if let Err(_) = self.create_pid_file() {
            warn!("error creating pid file: {}", std::process::id());
        }
        for (k,v) in &config.destinations {
            let destination = self.init_destination(v);
            let destination_id= destination.read().unwrap().id();
            self.destinations.insert(k.clone(), destination.clone());
            self.destinations_by_id.insert(destination_id, destination.clone());
            info!("created destination {}", k);
        }
        if config.enable_console {
            NGIN.init();
        }
        debug!("init finished");
    }

    /// fired every 60 seconds
    pub(crate) fn tick(&self) -> Result<(), tokio::timer::Error> {
        debug!("server tick");
        // bit crap, no real need for these atomic locks except we cannot get mut ref on a static
        // TODO we can but its unsafe, need to check this is really single threaded
        if ! self.task_set.load(Ordering::Relaxed) {
            self.task_set.store(true, Ordering::Relaxed);
            self.task.write().unwrap().replace(futures::task::current());
            // return OK here so when we boot we don't print stats
            return Ok(());
        }

        if self.shutdown.load(Ordering::Relaxed) {
            if self.connections_count.load(Ordering::SeqCst) == 0 {
                warn!("shutting down");
                std::process::exit(0);
                // TODO tokio Runtime that shutdown gracefully
                // return Err(tokio::timer::Error::shutdown());
            } else {
                if log_enabled!(Debug) {
                    debug!("connections still open {}", self.connections_count.load(Ordering::Relaxed));
                }
            }
        }

        if let Some(stats_dest) = self.destinations.get(&String::from("/xtomp/stat")) {
            {
                if let Ok(mut stats_dest) = stats_dest.write() {
                    debug!("pushing cc message");
                    stats_dest.push(&get_response_cc_msg(self.connections_count.load(Ordering::Relaxed),
                                                         self.uptime.load(Ordering::Relaxed),
                                                         self.connections_total.load(Ordering::Relaxed),
                                                         self.messages_total.load(Ordering::Relaxed)
                    )).ok(); // swallow errors
                }
            }
            for (name, d) in &self.destinations {
                if name.as_str() == "/xtomp/stat" {
                    continue;
                }
                let destination = d.read().unwrap();
                if destination.record_stats() {
                    debug!("pushing stats message for {}", destination.name());
                    {
                        if let Ok(mut stats_dest) = stats_dest.write() {
                            stats_dest.push(&get_response_stats_msg(destination.name(),
                                                                    destination.max_connections(),
                                                                    destination.len(),
                                                                    destination.message_delta(),
                                                                    destination.message_total()
                            )).ok(); // swallow errors
                        }
                    }
                }
            }
        }
        if CONFIG.enable_console {
            NGIN.set_stats_cc();
            for (i, destination) in self.destinations.values().enumerate() {
                if let Ok(destination) = destination.read() {
                    NGIN.set_stats_destination(i,
                                               destination.max_connections(),
                                               destination.len(),
                                               destination.message_delta(),
                                               destination.message_total());
                }
            }
        }
        info!("stat server sz={}, up={}, Σ={}, Σµ={},",
            self.connections_count.load(Ordering::Relaxed),
            self.uptime.load(Ordering::Relaxed),
            self.connections_total.load(Ordering::Relaxed),
            self.messages_total.load(Ordering::Relaxed)
        );
        for (name, d) in &self.destinations {
            if name.as_str() == "/xtomp/stat" {
                continue;
            }
            let destination = d.read().unwrap();
            if destination.record_stats() {
                info!("stat d={}, s={}, q={}, Δ={}, Σ={}", destination.name(),
                                                        destination.max_connections(),
                                                        destination.len(),
                                                        destination.message_delta(),
                                                        destination.message_total())
            }
        }

        Ok(())
    }

    pub fn is_shutdown(&self) -> bool {
        self.shutdown.load(Ordering::Relaxed)
    }

    pub fn shutdown(&self) {
        self.shutdown.store(true, Ordering::Relaxed);
        for (_k, d) in &self.destinations {
            d.write().unwrap().shutdown();
        }
    }

    pub fn notify(&self) {
        let task = self.task.write().unwrap();
        if let Some(task) = task.deref() {
            task.notify();
        }
    }
}


#[cfg(test)]
mod tests {
    use crate::config;

    use super::*;

    #[test]
    fn test_init() {
        let mut server = DestinationServer::new();

        let mut config = ServerConfig { ..Default::default() };

        let d1 = config::config::Destination {
            name: "memtop-a".to_string(), ..Default::default()
        };
        let d2 = config::config::Destination {
            name: "memtop-b".to_string(), ..Default::default()
        };
        let destinations = &mut config.destinations;
        destinations.insert(String::from("memtop-a"), d1);
        destinations.insert(String::from("memtop-b"), d2);

        server.init(&config);

        match server.find_destination("memtop-a") {
            Some(_d) => {}
            _ => panic!("destination not found")
        }
        match server.find_destination("memtop-b") {
            Some(_d) => {}
            _ => panic!("destination not found")
        }
        match server.find_destination("memtop-c") {
            Some(_d) => panic!("wrong destination found"),
            _ => {}
        }
    }


    #[test]
    fn test_new_id() {
        let server = DestinationServer::new();
        assert_ne!(server.new_message_id(), server.new_message_id());
    }
}