romp 0.5.2

//!Destination represents a single queue or topic.

use std::collections::HashMap;
use std::slice::Iter;
use std::sync::{Arc, RwLock};
use std::sync::atomic::{AtomicUsize, Ordering};

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

use crate::errors::ServerError;
use crate::init::SERVER;
use crate::message::stomp_message::{Ownership, StompMessage, StompCommand, Header};
use crate::session::subscription::Subscription;
use std::time::Duration;
use crate::config::config;

/// A Queue / Topic for messages
/// Messages are send in strict order to all subscribers, when min_delivery messages have been sent
/// the message is removed from the front of the queue.  Until then the queue does not progress.
// TODO support for destinations that send any messages they get immediately

pub struct Destination {
    // data
    name: String,
    id: usize,
    q: Box<Vec<Arc<StompMessage>>>,
    subs: HashMap<String, Arc<RwLock<Subscription>>>,

    // config
    /// maximum number of subscribers
    max_connections: usize,
    /// maximum number of message to store on the queue
    max_messages: usize,
    /// number of times a message must be sent to remove it from the queue, set min_delivery==1 for a queue,  set to 2 for a queue with redundancy, set to zero for a topic that discards messages if no one is listening
    min_delivery: usize,
    /// default expiry time
    expiry: u64,
    /// check each message has not expired before sending it
    pedantic_expiry: bool,
    filter: bool,
    filter_self: bool,
    sub_by_user: bool,
    record_stats: bool,

    // security
    /// force all subs be auto ack (otherwise a client can set ack:client and hang the topic by no Acking the head)
    auto_ack: bool,
    /// limit size of uploaded messages
    max_message_size: usize,
    /// prevent reads from the network
    read_block: bool,
    /// prevent writes from the network
    write_block: bool,
    /// prevent reads from t'interwebs
    web_read_block: bool,
    /// prevent writes from t'interwebs
    web_write_block: bool,

    // stats
    message_delta: AtomicUsize,
    message_total: AtomicUsize,

    // state
    task: Option<Task>,
    /// when true accept no more messages
    drain: bool,
    /// when true its time to die
    shutdown: bool,
    /// set to true when the head of the queue has not been delivered enough times (typically min_delivery==1 for a queue)
    queued: bool,
}

/// attributes that are configurable at runtime
pub struct RuntimeConfig {
    pub max_connections: Option<usize>,
    pub max_messages: Option<usize>,
    pub expiry: Option<u64>,
    pub max_message_size: Option<usize>,
    pub write_block: Option<bool>,
    pub read_block: Option<bool>,
    pub web_write_block: Option<bool>,
    pub web_read_block: Option<bool>,
}

impl Default for RuntimeConfig {
    fn default() -> Self {
        RuntimeConfig {
            max_connections: None,
            max_messages: None,
            expiry: None,
            max_message_size: None,
            write_block: None,
            read_block: None,
            web_write_block: None,
            web_read_block: None
        }
    }
}

impl Destination {

    pub(crate) fn new(name: String, id: usize,cfg: &config::Destination) -> Destination {
        Destination {
            name,
            q: Box::new(vec![]),
            id,
            subs: Default::default(),

            // config
            max_connections: cfg.max_connections,
            max_messages: cfg.max_messages,
            min_delivery: cfg.min_delivery,
            expiry: cfg.expiry,
            pedantic_expiry: cfg.pedantic_expiry,
            filter: cfg.filter,
            filter_self: cfg.filter_self,
            sub_by_user: cfg.sub_by_user,
            record_stats: cfg.stats,

            // security
            auto_ack: cfg.auto_ack,
            max_message_size: cfg.max_message_size,
            read_block: cfg.read_block,
            write_block: cfg.write_block,
            web_read_block: cfg.web_read_block,
            web_write_block: cfg.web_write_block,

            // stats
            message_delta: AtomicUsize::new(0),
            message_total: AtomicUsize::new(0),

            // state
            task: None,
            drain: false,
            shutdown: false,
            queued: false,
        }
    }

    // accessors

    pub fn name(&self) -> &String {
        &self.name
    }

    pub fn id(&self) -> usize {
        self.id
    }

    pub fn expiry(&self) -> u64 {
        self.expiry
    }

    pub fn auto_ack(&self) -> bool {
        self.auto_ack
    }

    pub fn write_block(&self) -> bool {
        self.write_block
    }

    pub fn read_block(&self) -> bool {
        self.read_block
    }

    pub fn web_write_block(&self) -> bool {
        self.web_write_block
    }

    pub fn web_read_block(&self) -> bool {
        self.web_read_block
    }

    pub fn record_stats(&self) -> bool {
        self.record_stats
    }

    pub fn filter_self(&self) -> bool {
        self.filter_self
    }

    pub fn max_connections(&self) -> usize {
        self.max_connections
    }

    pub fn message_total(&self) -> usize {
        self.message_total.load(Ordering::Relaxed)
    }

    pub fn message_delta(&self) -> usize {
        self.message_delta.swap(0, Ordering::Relaxed)
    }

    pub fn len(&self) -> usize {
        self.q.len()
    }

    pub fn iter(&mut self) -> Iter<Arc<StompMessage>> {
        self.q.iter()
    }

    // mutators

    pub(crate) fn set_task(&mut self, task: Task) {
        debug!("captured task");
        self.task = Some(task);
    }

    pub fn reconfigure(&mut self, cfg: RuntimeConfig) {
        if cfg.max_connections.is_some() {
            self.max_connections = cfg.max_connections.unwrap();
        }
        if cfg.max_messages.is_some() {
            self.max_messages = cfg.max_messages.unwrap();
        }
        if cfg.expiry.is_some() {
            self.expiry = cfg.expiry.unwrap();
        }
        if cfg.max_message_size.is_some() {
            self.max_message_size = cfg.max_message_size.unwrap();
        }
        if cfg.read_block.is_some() {
            self.read_block = cfg.read_block.unwrap();
        }
        if cfg.write_block.is_some() {
            self.write_block = cfg.write_block.unwrap();
        }
        if cfg.web_read_block.is_some() {
            self.web_read_block = cfg.web_read_block.unwrap();
        }
        if cfg.web_write_block.is_some() {
            self.web_write_block = cfg.web_write_block.unwrap();
        }
    }

    /// Put a copy of the message on the q
    // TODO return future for notification of durability
    pub fn push(&mut self, message: &StompMessage) -> Result<usize, ServerError> {
        return match self.push_security(&message) {
            Err(se) => Err(se),
            Ok(_) => {
                let id = SERVER.new_message_id();
                self.message_delta.fetch_add(1, Ordering::SeqCst);
                self.message_total.fetch_add(1, Ordering::SeqCst);

                let mut copy = message.clone_to_message(Ownership::Destination, id);
                copy.expiry = Duration::from_millis(self.expiry);

                self.q.push(Arc::new(copy));

                self.notify();

                Ok(id)
            }
        }
    }

    /// put an owned message on the q
    pub fn push_owned(&mut self, mut message: StompMessage) -> Result<usize, ServerError> {
        return match self.push_security(&message) {
            Err(se) => Err(se),
            Ok(_) => {
                let id = message.id;
                if id == 0 {
                    let id = SERVER.new_message_id();
                    self.message_total.fetch_add(1, Ordering::SeqCst);
                    message.id = id;
                }

                message.expiry = Duration::from_millis(self.expiry);

                self.q.push(Arc::new(message));

                self.notify();

                Ok(id)
            }
        }

    }

    /// validation that its safe to push this message to the destination
    fn push_security(&mut self, message: &StompMessage) -> Result<(), ServerError>  {
        if self.drain {
            return Err(ServerError::ShuttingDown);
        }
        if self.q.len() == self.max_messages {
            return Err(ServerError::DestinationFlup);
        }
        if message.body_len() > self.max_message_size {
            return Err(ServerError::MessageFlup);
        }
        Ok(())
    }

    /// push a message to a specific logged in user, destination must be hashed by username
    pub fn push_direct(&mut self, message: StompMessage) -> bool {
        if self.sub_by_user {
            // how to borrow an Option<String>
            let to = &message.to.as_ref();

            if let Some(sub) = self.subs.get(to.unwrap()) {
                return sub.write().unwrap().publish(self.id, Arc::new(message), false);
            }
        }
        false
    }

    /// forward a downstream MESSAGE to a SEND on this destination
    pub fn push_resend(&mut self, message: &StompMessage) -> Result<usize, ServerError> {
        return match self.push_security(&message) {
            Err(se) => Err(se),
            Ok(_) => {
                // for now we increment the id because we dont know how upstream allocates ids.
                let id = SERVER.new_message_id();
                let original_id = message.id;
                self.message_delta.fetch_add(1, Ordering::SeqCst);
                self.message_total.fetch_add(1, Ordering::SeqCst);

                let mut copy = message.clone_to_message(Ownership::Destination, id);
                copy.expiry = Duration::from_millis(self.expiry);

                if copy.command == StompCommand::Message {
                    copy.command = StompCommand::Send;
                }
                copy.push_header(Header {
                    name: "orig-id".to_string(),
                    value: original_id.to_string()
                });

                self.q.push(Arc::new(copy));

                self.notify();

                Ok(id)
            }
        }
    }

    pub fn find_sub(&self, to: &String) -> Option<&Arc<RwLock<Subscription>>> {
        return if let Some(sub) = self.subs.get(to) {
            Some(sub)
        } else {
            None
        }
    }

    /// called periodically for message expiry
    pub fn timeout_messages(&mut self) -> Result<(), tokio::timer::Error> {
        self.expire();
        if self.drain || self.shutdown {
            return Err(tokio::timer::Error::shutdown());
        }
        Ok(())
    }

    // TODO pushes HashMap responsibility to client, alternative it get a read lock on Session
    pub(crate) fn subscribe(&mut self, sub: Arc<RwLock<Subscription>>) -> Result<(), ServerError> {
        if self.drain || self.shutdown {
            return Err(ServerError::ShuttingDown);
        }

        if self.subs.len() == self.max_connections {
            return Err(ServerError::SubscriptionFlup);
        }

        // clone the string because borrow checker cant handle realities of subscription lifetime
        let hash_id;
        {
            if self.sub_by_user {
                hash_id = sub.read().unwrap().user_id();
            } else {
                hash_id = sub.read().unwrap().hash_id().clone();
            }
        }
        self.subs.insert(hash_id, sub);

        if self.queued {
            debug!("subscribe while queued, notify");
            self.notify();
        } else {
            debug!("subscribed to {}", self.name);
        }

        Ok(())
    }

    // Unsubscribe return true on success
    pub(crate) fn unsubscribe(&mut self, id: &String) -> bool {
        self.subs.remove(id).is_some()
    }

    pub fn shutdown(&mut self) {
        self.shutdown = true;
        self.q.clear();
        match &self.task {
            Some(task) => task.notify(),
            _ => {}
        }
    }

    pub fn drain(&mut self) {
        self.drain = true;
        self.expire();
        match &self.task {
            Some(task) => task.notify(),
            _ => {}
        }
    }

    pub fn clean(&mut self) -> usize {
        let len = self.q.len();
        self.q.clear();
        return len;
    }

    fn notify(&self) {
        match &self.task {
            None => warn!("publish before Destination is started, message queued"),
            Some(task) => {
                debug!("destination notified");
                task.notify()
            },
        }
    }

    /// expire messages that are past their sell by date
    fn expire(&mut self) {
        let now = Utc::now().timestamp() as u64;
        let len = self.q.len();
        self.q.retain(|message| {
            message.timestamp.timestamp() as u64 + message.expiry.as_secs() as u64 > now
        });
        let expired = len - self.q.len();
        if expired > 0 {
            info!("expired {} messages, remain={}", expired, self.q.len());
        }
    }

    /// Acknowledge a message as consumed
    /// This only makes sense for queues, since we don't track Acks per subscription
    pub fn ack(&self, id: usize) {
        debug!("ACKing {}", id);
        for message in self.q.iter() {
            if message.id == id {
                if message.increment_delivered(1) + 1 >= self.min_delivery {
                    self.notify();
                }
                break;
            }
        }
    }

    /// Acknowledge a message as unable to be consumed
    // TODO DLQ
    pub fn nack(&self, id: usize) {
        self.ack(id);
    }

    /// called once each time a message is added to the q, or acked, or a sub
    /// return true if we are still alive
    pub(crate) fn poll(&mut self) -> bool {
        debug!("destination polled {} messages on q", self.q.len());
        if self.drain && self.q.len() == 0 {
            warn!("destination {} dead", self.name);
            return false;
        }

        if self.shutdown {
            warn!("destination {} dead", self.name);
            return false;
        }

        if self.q.len() == 0 {
            self.queued = false;
            return true;
        }

        let now = Utc::now().timestamp() as u64;
        while let Some(message) = self.q.get(0) {
            if self.pedantic_expiry {
                let mut do_drop = false;
                {
                    if message.timestamp.timestamp() as u64 + message.expiry.as_secs() as u64 > now {
                        info!("expired 1 message");
                        do_drop = true;
                    }
                }
                if do_drop {
                    self.q.pop();
                    continue;
                }
            }

            let mut delivered: usize = 0;
            for (key, subscription) in self.subs.iter() {
                // publish clone of Arc, only one copy of message exists
                {
                    debug!("publishing instance of message to {}", key);
                    if subscription.write().unwrap().publish(self.id, message.clone(), self.filter) {
                        delivered += 1;
                    }
                }
            }

            let total_delivered;
            {
                total_delivered = message.increment_delivered(delivered) + delivered;
            }

            if self.min_delivery == 0 {
                // not tracking delivery
                self.q.remove(0);
            } else if self.min_delivery > total_delivered {
                // queue does not progress if don't have enough subscribers
                self.queued = true;
                debug!("queued, not enough subs");
                return true;
            } else {
                // delivered enough times
                self.q.remove(0);
            }
        }

        self.queued = false;

        return true;
    }
}

#[cfg(test)]
mod tests {
    use std::time;
    use std::time::Duration;

    use crate::config::config;

    use super::*;

    #[test]
    fn test() {
        let dest_config = config::Destination{min_delivery: 1, ..Default::default()};
        let mut destination = Destination::new(String::from(""), 1, &dest_config);

        let id_1 = destination.push(&StompMessage::new_send(b"{\"test\": true}", 1)).unwrap_or(99);
        let id_2 = destination.push(&StompMessage::new_send(b"{\"test\": false}", 2)).unwrap_or(99);

        // with no one listening poll should keep messages on the q
        destination.poll();
        assert!(destination.queued);
        assert_eq!(destination.q.len(), 2);

        assert_eq!(destination.q.remove(0).id, id_1);
        assert_eq!(destination.q.remove(0).id, id_2);
        // poll empty dest should set queued flag to false
        destination.poll();
        assert!(!destination.queued);

        let id_3 = destination.push(&StompMessage::new_send(b"{\"test\": false}", 3)).unwrap_or(99);

        assert_eq!(destination.q.remove(0).id, id_3);

        let mut message = StompMessage::new_send(b"{\"timeout\": true}", 4);
        message.expiry = Duration::new(0, 0);

        std::thread::sleep(time::Duration::from_millis(10));

        destination.expire();
        assert_eq!(destination.q.len(), 0);

    }


    #[test]
    fn test_no_min_delivery() {
        let mut destination = Destination::new(String::from(""), 1, &config::Destination{min_delivery: 0, ..Default::default()});

        match destination.push(&StompMessage::new_send(b"{\"test\": true}", 1)) {
            Err(_) => panic!("push failed"),
            _ => {}
        }
        match destination.push(&StompMessage::new_send(b"{\"test\": false}", 2)) {
            Err(_) => panic!("push failed"),
            _ => {}
        }

        // with min_delivery set to zero poll should empty the q
        destination.poll();
        assert!(!destination.queued);
        assert_eq!(destination.q.len(), 0);

    }



    #[test]
    fn test_ack() {
        let mut destination = Destination::new(String::from("ackme"), 1, &config::Destination{min_delivery: 1, ..Default::default()});

        assert_eq!(destination.auto_ack(), false);

        match destination.push(&StompMessage::new_send(b"{\"test\": true}", 2)) {
            Err(_) => panic!("push failed"),
            _ => {}
        }
        match destination.push(&StompMessage::new_send(b"{\"test\": false}", 4)) {
            Err(_) => panic!("push failed"),
            _ => {}
        }

        assert_eq!(destination.q.len(), 2);
        destination.poll();
        // nothing acked we should still have the messages
        assert_eq!(destination.queued, true);
        assert_eq!(destination.q.len(), 2);

        destination.ack(destination.q[0].id);
        destination.poll();
        assert_eq!(destination.q.len(), 1);

        destination.ack(destination.q[0].id);
        destination.poll();
        assert_eq!(destination.q.len(), 0);
    }
}