datastreamservicelib 1.0.0

Rust version of https://gitlab.com/advian-oss/python-datastreamservicelib
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155

/// Simple test message publisher
use bytes::BytesMut;
use clap::{crate_authors, crate_version, App, Arg};
use datastreamcorelib::datamessage::PubSubDataMessage;
use datastreamcorelib::imagemessage::PubSubImageMessage;
use datastreamcorelib::logging::init_logging;
use datastreamservicelib::heartbeat::heartbeat_task;
use datastreamservicelib::utils::{setup_signals_termination, 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 futures::stream::FuturesUnordered;
use log;
use serde_json;
use std::sync::atomic::Ordering;
use std::time::Duration;
use tokio::runtime;
use tokio::time::delay_for;

/// Main entrypoint, handler CLI parsing and runtime setup
fn main() -> Fallible<()> {
    init_logging(log::LevelFilter::Trace)?;
    let app = App::new("testpublisher")
        .about("PUBlish test messages ")
        .version(crate_version!())
        .author(crate_authors!())
        .arg(
            Arg::with_name("socket_uri")
                .short("s")
                .long("socket_uri")
                .takes_value(true)
                .empty_values(false)
                .required(true)
                .help("For example ipc:///tmp/publisher.sock"),
        )
        .arg(
            Arg::with_name("topic")
                .short("t")
                .long("topic")
                .takes_value(true)
                .empty_values(false)
                .required(true)
                .help("The topic to use for sending"),
        )
        .arg(
            Arg::with_name("images")
                .short("i")
                .long("images")
                .help("Send images (one red pixel)"),
        )
        .arg(
            Arg::with_name("count")
                .short("c")
                .long("count")
                .takes_value(true)
                .empty_values(false)
                .help("Number of messages to send"),
        );
    let app_args = app.get_matches();
    let socket_uri = app_args.value_of("socket_uri").unwrap().to_string();
    let topic = app_args.value_of("topic").unwrap().to_string();
    let count_to: i64 = match app_args.value_of("count") {
        None => -1,
        Some(val) => val.parse()?,
    };
    let images = app_args.is_present("images");

    // Use the non-threaded scheduler explicitly
    let mut rt = runtime::Builder::new()
        .basic_scheduler()
        .enable_time()
        .enable_io()
        .build()?;
    rt.block_on(mainloop(socket_uri, topic, count_to, images))?;
    Ok(())
}

/// Mainloop task, runtime will block on this one.
async fn mainloop(socket_uri: String, topic: String, count_to: i64, images: bool) -> Fallible<()> {
    let term = setup_signals_termination()?;
    let psmgr = TokioPubSubManager::instance();
    // Scope limit the lock to this critical part
    {
        let mut psmgr_locked = psmgr.lock();
        psmgr_locked.term_flag = term.clone();
        psmgr_locked.set_default_pub_uris(&vec![socket_uri])?;
    }
    // Give the socket a moment to stabilize
    delay_for(Duration::from_millis(150)).await;

    // Track tasks
    let tasks = FuturesUnordered::new();
    tasks.push(tokio::spawn(heartbeat_task(term.clone())));
    if images {
        tasks.push(tokio::spawn(img_publisher_task(
            topic,
            count_to,
            term.clone(),
        )));
    } else {
        tasks.push(tokio::spawn(publisher_task(topic, count_to, term.clone())));
    }
    wait_for_tasks(tasks).await?;

    Ok(())
}

/// Publisher task, will publish messages with counter in the given topic
async fn publisher_task(topic: String, count_to: i64, term: TerminationFlag) -> Fallible<()> {
    let psmgr = TokioPubSubManager::instance();
    let mut msgno: i64 = 0;
    loop {
        msgno = msgno + 1;
        let mut msg = PubSubDataMessage::new(topic.clone())?;
        msg.data["msgno"] = serde_json::to_value(msgno)?;
        psmgr.lock().publish(&msg)?;
        if count_to > 0 && msgno >= count_to {
            log::info!("Sent {} messages, exiting now", msgno);
            term.store(true, Ordering::Relaxed);
            return Ok(());
        }
        if term.load(Ordering::Relaxed) {
            log::trace!("Got term flag, exiting publisher task");
            return Ok(());
        }
        delay_for(Duration::from_millis(100)).await;
    }
}

async fn img_publisher_task(topic: String, count_to: i64, term: TerminationFlag) -> Fallible<()> {
    let psmgr = TokioPubSubManager::instance();
    let mut msgno: i64 = 0;
    loop {
        msgno = msgno + 1;
        let mut msg = PubSubImageMessage::new(topic.clone())?;
        msg.data["msgno"] = serde_json::to_value(msgno)?;
        msg.imginfo["w"] = serde_json::json!(1);
        msg.imginfo["h"] = serde_json::json!(1);
        let imgdata: [u8; 3] = [0, 0, 255];
        msg.imgdata = BytesMut::from(&imgdata[..]);
        psmgr.lock().publish(&msg)?;
        if count_to > 0 && msgno >= count_to {
            log::info!("Sent {} messages, exiting now", msgno);
            term.store(true, Ordering::Relaxed);
            return Ok(());
        }
        if term.load(Ordering::Relaxed) {
            log::trace!("Got term flag, exiting publisher task");
            return Ok(());
        }
        delay_for(Duration::from_millis(100)).await;
    }
}