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;
pub struct DestinationServer {
message_id : AtomicUsize,
session_id: AtomicUsize,
destination_id: AtomicUsize,
destinations: HashMap<String, Arc<RwLock<Destination>>>,
destinations_by_id: HashMap<usize, Arc<RwLock<Destination>>>,
shutdown: AtomicBool,
task_set: AtomicBool,
task: RwLock<Option<Task>>,
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
}
pub(crate) fn new_message_id(&self) -> usize {
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));
}
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()
}
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");
}
pub(crate) fn tick(&self) -> Result<(), tokio::timer::Error> {
debug!("server tick");
if ! self.task_set.load(Ordering::Relaxed) {
self.task_set.store(true, Ordering::Relaxed);
self.task.write().unwrap().replace(futures::task::current());
return Ok(());
}
if self.shutdown.load(Ordering::Relaxed) {
if self.connections_count.load(Ordering::SeqCst) == 0 {
warn!("shutting down");
std::process::exit(0);
} 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(); }
}
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(); }
}
}
}
}
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());
}
}