ds1090/

lib.rs

pub mod defaultconfig;
pub mod flight;

use datastreamcorelib::datamessage::PubSubDataMessage;
use datastreamservicelib::heartbeat::heartbeat_task;
use datastreamservicelib::utils::{configval_as_string_vec, wait_for_tasks};
use datastreamservicelib::zmqwrappers::TokioPubSubManager;
use datastreamservicelib::TerminationFlag;
// we need to explicitly import this trait for it to work
use datastreamcorelib::pubsub::PubSubManager;
use failure::Fallible;
use flight::{Flights, Message};
use futures::stream::FuturesUnordered;
use futures::{future, StreamExt};
use log;
use std::net::SocketAddr;
use std::sync::atomic::Ordering;
use std::time::Duration;
use tokio::net::TcpStream;
use tokio::time::delay_for;
use tokio_util;
use tokio_util::codec::{FramedRead, LinesCodec};
use toml;

/// Main entrypoint
pub async fn mainloop(config: toml::Value, term: TerminationFlag) -> Fallible<()> {
    // Form dump1090 address
    let dump1090addr = match validate_dump1090addr(&config) {
        Ok(addr) => addr,
        Err(e) => {
            return Err(failure::err_msg(format!(
                "dump1090 address validation failed: {}",
                e
            )));
        }
    };
    // Whether or not to combine data from multiple packets and only output
    // when full flight info is available, or it changes
    let track: bool = match config["process"]["track"].as_bool() {
        Some(track) => track,
        None => false,
    };
    // These need be available for a Flight
    // FIXME: turn these into an enum
    let required_fields: Vec<String> = match config["process"]["required_fields"].as_array() {
        Some(required_fields) => required_fields
            .to_vec()
            .iter()
            .map(|x| String::from(x.as_str().unwrap()))
            .collect(),
        None => vec![],
    };

    // Instances of things we are going to need
    let psmgr = TokioPubSubManager::instance();

    // Just using ? in the lock will fail with something not being sendable
    match psmgr.lock() {
        Err(_) => {
            term.store(true, Ordering::Relaxed);
            return Err(failure::err_msg("Could not lock psmgr"));
        }
        Ok(mut psmgr_locked) => {
            psmgr_locked.term_flag = term.clone();
            psmgr_locked
                .set_default_pub_uris(&configval_as_string_vec(&config["zmq"]["pub_sockets"])?)?;
        }
    }
    // Give the socket a moment to stabilize
    delay_for(Duration::from_millis(150)).await;

    // Create tasks
    let tasks = FuturesUnordered::new();
    tasks.push(tokio::spawn(heartbeat_task(term.clone())));
    tasks.push(tokio::spawn(reader_publisher_task(
        track,
        required_fields,
        dump1090addr,
        "messages".to_string(),
        term.clone(),
    )));
    wait_for_tasks(tasks).await?;
    Ok(())
}

/// Validates that dump1090 address and port in configuration look sane
pub fn validate_dump1090addr(config: &toml::Value) -> Fallible<SocketAddr> {
    let host: String = match config["dump1090"]["host"].as_str() {
        Some(host) => host.to_string(),
        None => {
            return Err(failure::err_msg("dump1090.host not supplied"));
        }
    };
    let port: u16 = match config["dump1090"]["port"].as_integer() {
        Some(port) => {
            if port < 1 || port > 65535 {
                return Err(failure::err_msg(
                    "dump1090.port does not look like a valid TCP port number",
                ));
            }
            port as u16
        }
        None => 30003,
    };

    let dump1090addr = match format!("{}:{}", host, port).parse::<SocketAddr>() {
        Ok(addr) => addr,
        Err(e) => {
            return Err(failure::err_msg(format!(
                "Could not parse dump1090 address: {}",
                e
            )))
        }
    };

    Ok(dump1090addr)
}

/// Reader task reads dump1090 messages over TCP, and parses and PUBlishes them
async fn reader_publisher_task(
    track: bool,
    required_fields: Vec<String>,
    addr: SocketAddr,
    topic: String,
    term: TerminationFlag,
) -> Fallible<()> {
    let psmgr = TokioPubSubManager::instance();
    let mut flights = Flights::new(required_fields);
    loop {
        // We can't await inside the lock match so yield here instead to the arm
        // where we fail to acquire the lock
        // https://docs.rs/tokio/0.2.16/tokio/fn.spawn.html#using-send-values-from-a-task
        tokio::task::yield_now().await;
        if term.load(Ordering::Relaxed) {
            log::trace!("Got term flag, exiting reader task");
            return Ok(());
        }

        // Connect to dump1090 port
        let mut stream = match TcpStream::connect(&addr).await {
            Ok(stream) => {
                log::debug!("Connected: {:?}", stream);
                stream
            }
            Err(e) => {
                log::debug!("Connection failed to: {:?}: {}, retrying", addr, e);
                delay_for(Duration::from_millis(1000)).await;
                continue;
            }
        };
        let (r, _w) = stream.split();
        let mut stream =
            FramedRead::new(r, LinesCodec::new_with_max_length(1024)).filter_map(|i| match i {
                Ok(i) => future::ready(Some(i)),
                Err(e) => {
                    log::error!("Could not read from socket: {}", e);
                    future::ready(None)
                }
            });

        // Consume input lines
        while let Some(line) = stream.next().await {
            if term.load(Ordering::Relaxed) {
                log::trace!("Got term flag, exiting reader task");
                return Ok(());
            }

            // Attempt to parse a line into a message
            let mut adsbmessage = Message::parse(&line);
            let mut pubsubmessage = PubSubDataMessage::new(topic.clone())?;

            // Track flights if configured
            if track && adsbmessage.valid() {
                if let Some(_flight) = flights.update(&mut adsbmessage) {
                    pubsubmessage.data = _flight.values.json();
                } else {
                    continue;
                }
            } else if !track && adsbmessage.valid() {
                pubsubmessage.data = adsbmessage.json();
            } else {
                continue;
            }

            // Publish message
            match psmgr.lock() {
                Err(e) => {
                    // Wait for the socket lock
                    log::debug!("Could not acquire psmgr lock: {:?}", e);
                    continue;
                }
                Ok(psmgr_locked) => {
                    psmgr_locked.publish(&pubsubmessage)?;
                }
            }
        }

        // Pause for a while before reconnect
        delay_for(Duration::from_millis(1000)).await;
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::defaultconfig::default_config_str;
    use datastreamcorelib::logging::init_logging;
    use std::env::temp_dir;
    use std::sync::atomic::AtomicBool;
    use std::sync::Arc;
    use toml;

    #[tokio::test]
    async fn mainloop_heartbeat() {
        init_logging(log::LevelFilter::Trace).unwrap();
        let term: TerminationFlag = Arc::new(AtomicBool::new(false));
        let mut config: toml::Value = toml::from_str(default_config_str().as_str()).unwrap();
        let mut tmppath1 = temp_dir();
        tmppath1.push("3f7639c2-df18-49fa-a788-11a0486e7ac1_pub.sock");
        let sockpath1 = "ipc://".to_string() + &tmppath1.to_string_lossy();
        let new_toml_sockets = toml::Value::try_from(vec![sockpath1.clone()]).unwrap();
        config["zmq"]["pub_sockets"] = new_toml_sockets;
        assert_eq!(
            configval_as_string_vec(&config["zmq"]["pub_sockets"]).unwrap(),
            vec![sockpath1.clone()]
        );
        let maintask = tokio::spawn(mainloop(config, term.clone()));

        // Be fancy and get rid of the delay
        delay_for(Duration::from_millis(1000)).await;

        // Tell things to shut down
        term.store(true, Ordering::Relaxed);
        maintask.await.expect("Main task failed").unwrap();
    }
}