diff --git a/Cargo.lock b/Cargo.lock index 102295d..95375d5 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -441,9 +441,9 @@ checksum = "a6cb138bb79a146c1bd460005623e142ef0181e3d0219cb493e02f7d08a35695" [[package]] name = "isis_streaming_data_types" -version = "0.1.2" +version = "0.1.3" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "4a40f55610cfa45cc165129d63849bd6a805a918c1e1f64c7de58cba09d181f9" +checksum = "7189171301072e5e8c6541db5f7f086a45045a20c190cf946a208f8cc4b8f142" dependencies = [ "flatbuffers", ] diff --git a/Cargo.toml b/Cargo.toml index 7c5cdfc..3ddd67b 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -10,7 +10,7 @@ log = "0.4.29" uuid = { version = "1.22.0", features = ["v4"] } env_logger = "0.11.9" clap-verbosity-flag = "3.0.4" -isis_streaming_data_types = "0.1.1" +isis_streaming_data_types = "0.1.3" rand = { version = "0.10.1", features = ["thread_rng"]} rand_distr = "0.6.0" flatbuffers = "25.12.19" diff --git a/src/cli_utils.rs b/src/cli_utils.rs index b19b553..4e5e986 100644 --- a/src/cli_utils.rs +++ b/src/cli_utils.rs @@ -1,4 +1,5 @@ use anyhow::{Context, Result, bail}; +use rdkafka::ClientConfig; use std::str::FromStr; pub(crate) fn parse_broker_spec(s: &str) -> Result { @@ -68,6 +69,21 @@ pub struct KafkaOption { pub value: String, } +pub fn set_kafka_options( + client_config: &mut ClientConfig, + kafka_config: &Option>, +) { + if let Some(kafka_options) = kafka_config { + for option in kafka_options { + println!( + "Setting Kafka config option {}={}", + option.key, option.value + ); + client_config.set(&option.key, &option.value); + } + } +} + impl FromStr for KafkaOption { type Err = String; diff --git a/src/consume.rs b/src/consume.rs index 3cf9993..84ea6ea 100644 --- a/src/consume.rs +++ b/src/consume.rs @@ -1,5 +1,5 @@ use crate::KafkaOption; -use crate::cli_utils::BrokerAndTopic; +use crate::cli_utils::{BrokerAndTopic, set_kafka_options}; use isis_streaming_data_types::{deserialize_message, get_schema_id}; use log::{debug, error, info}; @@ -35,16 +35,7 @@ pub fn consume(config: &ConsumeConfig) { let mut client_config = ClientConfig::new(); client_config.set("group.id", Uuid::new_v4().to_string()); client_config.set("bootstrap.servers", config.topic.broker()); - - if let Some(kafka_options) = &config.kafka_config { - for option in kafka_options { - println!( - "Setting Kafka config option {}={}", - option.key, option.value - ); - client_config.set(&option.key, &option.value); - } - } + set_kafka_options(&mut client_config, &config.kafka_config); let consumer: BaseConsumer = client_config.create().expect("Base creation failed"); diff --git a/src/count.rs b/src/count.rs index d707bea..7941bd3 100644 --- a/src/count.rs +++ b/src/count.rs @@ -1,5 +1,5 @@ use crate::KafkaOption; -use crate::cli_utils::BrokerAndTopic; +use crate::cli_utils::{BrokerAndTopic, set_kafka_options}; use futures::stream::StreamExt; use log::error; use rdkafka::consumer::{Consumer, DefaultConsumerContext, StreamConsumer}; @@ -15,16 +15,7 @@ pub async fn count( let mut config = ClientConfig::new(); config.set("group.id", Uuid::new_v4().to_string()); config.set("bootstrap.servers", topic.broker()); - - if let Some(kafka_options) = kafka_config { - for option in kafka_options { - println!( - "Setting Kafka config option {}={}", - option.key, option.value - ); - config.set(&option.key, &option.value); - } - } + set_kafka_options(&mut config, &kafka_config); let consumer: StreamConsumer = config.create().expect("Consumer creation failed"); diff --git a/src/howl.rs b/src/howl.rs index bfdb1df..40da051 100644 --- a/src/howl.rs +++ b/src/howl.rs @@ -2,13 +2,18 @@ use crate::KafkaOption; use std::thread; use std::time::{Duration, SystemTime}; -use flatbuffers::FlatBufferBuilder; +use flatbuffers::{FlatBufferBuilder, WIPOffset}; use isis_streaming_data_types::flatbuffers_generated::events_ev44::{ Event44Message, Event44MessageArgs, finish_event_44_message_buffer, }; use isis_streaming_data_types::flatbuffers_generated::pulse_metadata_pu00::{ Pu00Message, Pu00MessageArgs, finish_pu_00_message_buffer, }; +use isis_streaming_data_types::flatbuffers_generated::veto_configuration_vc00::{ + Vetoes, VetoesArgs, finish_vetoes_buffer, +}; + +use crate::cli_utils::set_kafka_options; use isis_streaming_data_types::flatbuffers_generated::run_start_pl72::{ RunStart, RunStartArgs, SpectraDetectorMapping, SpectraDetectorMappingArgs, finish_run_start_buffer, @@ -25,6 +30,8 @@ use rdkafka::producer::{BaseRecord, DefaultProducerContext, ThreadedProducer}; use serde_json::json; use uuid::Uuid; +pub const VETO_COUNT: usize = 32; + fn generate_run_start<'a>( fbb: &'a mut FlatBufferBuilder<'_>, det_max: i32, @@ -114,12 +121,56 @@ fn generate_run_stop<'a>(fbb: &'a mut FlatBufferBuilder<'_>, job_id: &str) -> &' fbb.finished_data() } +fn get_veto_names_fbb<'a>( + veto_names: &[String], + fbb: &mut FlatBufferBuilder<'a>, + buf: &mut Vec>, +) { + buf.clear(); + + for i in 0..VETO_COUNT { + let name = veto_names + .get(i) + .cloned() + .unwrap_or(format!("saluki_veto_{i}")); + + buf.push(fbb.create_string(&name.to_string())); + } +} + +fn get_enabled_vetoes(conf: &HowlConfig, rng: &mut ThreadRng) -> u32 { + let mut vetoes = 0; + let mut prob; + + for i in 0..VETO_COUNT { + prob = conf.veto_probability.get(i).cloned().unwrap_or(0.0); + vetoes = (vetoes << 1) | rng.random_bool(prob) as u32; + } + + vetoes +} + +fn get_active_vetoes(conf: &HowlConfig) -> u32 { + let mut vetoes = 0; + let mut active; + + for i in 0..VETO_COUNT { + active = conf.enabled_vetoes.get(i).cloned().unwrap_or(false); + vetoes = (vetoes << 1) | active as u32; + } + + vetoes +} + +#[allow(clippy::too_many_arguments)] fn produce_messages( producer: &ThreadedProducer, fbb: &mut FlatBufferBuilder, rng: &mut ThreadRng, frame: u32, conf: &HowlConfig, + active_vetoes: &u32, + enabled_vetoes: &u32, current_job_id: &mut String, ) { // get current time @@ -130,77 +181,23 @@ fn produce_messages( .try_into() .expect("This will fail after April 11th, 2262"); - match producer.send( - BaseRecord::to(conf.event_topic) - .key("") - .payload(generate_fake_metadata( - rng, - fbb, - now_nanos, - conf.veto_probability, - )) - .timestamp(now_nanos / 1_000_000), - ) { - Ok(_) => {} - Err(err) => { - error!("Failed to send messages: {}", err.0); - } - } - - let ev44 = generate_fake_events(fbb, rng, frame, conf.event_message_config, now_nanos).to_vec(); - - for _ in 0..conf.messages_per_frame { - match producer.send( - BaseRecord::to(conf.event_topic) - .key("") - .payload(if conf.fast { - ev44.as_slice() - } else { - generate_fake_events(fbb, rng, frame, conf.event_message_config, now_nanos) - }) - .timestamp(now_nanos / 1_000_000), - ) { - Ok(_) => {} - Err(err) => { - error!("Failed to send messages: {}", err.0); - } - } - } - if conf.frames_per_run > 0 && frame.is_multiple_of(conf.frames_per_run) { info!( "Starting new run after {} simulated frames", conf.frames_per_run ); - match producer.send( - BaseRecord::to(conf.run_info_topic) - .key("") - .payload(generate_run_stop(fbb, current_job_id)) - .timestamp(now_nanos / 1_000_000), - ) { - Ok(_) => {} - Err(err) => { - error!("Failed to send run stop: {}", err.0); - } - } - *current_job_id = Uuid::new_v4().to_string(); - match producer.send( - BaseRecord::to(conf.run_info_topic) - .key("") - .payload(generate_run_start( - fbb, - conf.event_message_config.det_max, - conf.event_topic, - current_job_id, - )) - .timestamp(now_nanos / 1_000_000), - ) { - Ok(_) => {} - Err(err) => { - error!("Failed to send run start: {}", err.0); - } + + if frame != 0 { + send_run_stop(producer, fbb, conf, current_job_id, now_nanos); + *current_job_id = Uuid::new_v4().to_string(); } + + send_run_start(producer, fbb, conf, current_job_id, now_nanos); + send_veto_config(producer, fbb, conf, enabled_vetoes, now_nanos); } + + send_frame_metadata(producer, fbb, conf, active_vetoes, now_nanos); + send_frame_data(producer, fbb, conf, rng, frame, now_nanos); } pub struct EventMessageConfig { @@ -244,61 +241,61 @@ fn generate_fake_events<'a>( } fn generate_fake_metadata<'a>( - rng: &mut ThreadRng, + vetoes_mask: &u32, fbb: &'a mut FlatBufferBuilder<'_>, timestamp_ns: i64, - veto_probability: f64, ) -> &'a [u8] { fbb.reset(); - let is_vetoed = rng.random_range(0.0..1.0) < veto_probability; + let args = Pu00MessageArgs { reference_time: timestamp_ns, message_id: 0, source_name: Some(fbb.create_string("saluki")), period_number: Some(0), - vetos: Some(if is_vetoed { 1 } else { 0 }), + vetos: Some(*vetoes_mask), // active proton_charge: Some(0.1), }; let pu00 = Pu00Message::create(fbb, &args); finish_pu_00_message_buffer(fbb, pu00); + fbb.finished_data() } -pub struct HowlConfig<'a> { - pub broker: &'a str, - pub event_topic: &'a str, - pub run_info_topic: &'a str, - pub messages_per_frame: u32, - pub frames_per_second: u32, - pub frames_per_run: u32, - pub veto_probability: f64, // 1 = always vetoed, 0 = never vetoed - pub event_message_config: &'a EventMessageConfig, - pub fast: bool, - pub kafka_config: Option>, -} +fn generate_veto_config<'a>( + veto_names: &[String], + fbb: &'a mut FlatBufferBuilder<'_>, + timestamp_ns: i64, + vetoes_mask: &u32, +) -> &'a [u8] { + fbb.reset(); + let mut veto_names_fbb: Vec> = Vec::new(); + get_veto_names_fbb(veto_names, fbb, &mut veto_names_fbb); -pub fn howl(conf: &HowlConfig) { - // create producer - let mut fbb = FlatBufferBuilder::new(); - let mut rng = rand::rng(); + let args = VetoesArgs { + timestamp: timestamp_ns, + vetoes: *vetoes_mask, // enabled + veto_names: Some(fbb.create_vector(&veto_names_fbb)), + }; + let vc00 = Vetoes::create(fbb, &args); + finish_vetoes_buffer(fbb, vc00); - let now_nanos = SystemTime::now() - .duration_since(SystemTime::UNIX_EPOCH) - .expect("Failed to get system time") - .as_nanos() - .try_into() - .expect("This will fail after April 11th, 2262"); + fbb.finished_data() +} +fn calculate_data_rate( + fbb: &mut FlatBufferBuilder<'_>, + rng: &mut ThreadRng, + conf: &HowlConfig, + timestamp_ns: i64, + vetoes_mask: &u32, +) { let ev44_size = - generate_fake_events(&mut fbb, &mut rng, 0, conf.event_message_config, now_nanos).len() - as u32; + generate_fake_events(fbb, rng, 0, conf.event_message_config, timestamp_ns).len() as u32; debug!("ev44 size is {ev44_size} bytes"); - let pu00_size = - generate_fake_metadata(&mut rng, &mut fbb, now_nanos, conf.veto_probability).len() as u32; + let pu00_size = generate_fake_metadata(vetoes_mask, fbb, timestamp_ns).len() as u32; debug!("pu00 size is {pu00_size} bytes"); - // calculate overall rate (with both ev44 and pu00) let rate_bytes_per_sec = ev44_size * conf.messages_per_frame * conf.frames_per_second + pu00_size * conf.frames_per_second; debug!("bytes per second: {rate_bytes_per_sec}"); @@ -311,39 +308,128 @@ pub fn howl(conf: &HowlConfig) { ); println!("Each pu00 is {pu00_size} bytes"); println!("Each ev44 is {ev44_size} bytes"); +} - let mut config: ClientConfig = ClientConfig::new(); - config.set("bootstrap.servers", conf.broker); - - if let Some(kafka_options) = &conf.kafka_config { - for option in kafka_options { - println!( - "Setting Kafka config option {}={}", - option.key, option.value - ); - config.set(&option.key, &option.value); +fn send_frame_metadata( + producer: &ThreadedProducer, + fbb: &mut FlatBufferBuilder<'_>, + conf: &HowlConfig, + vetoes_mask: &u32, + now_nanos: i64, +) { + match producer.send( + BaseRecord::to(conf.event_topic) + .key("") + .payload(generate_fake_metadata(vetoes_mask, fbb, now_nanos)) + .timestamp(now_nanos / 1_000_000), + ) { + Ok(_) => {} + Err(err) => { + error!("Failed to send messages: {}", err.0); } } +} - let producer: ThreadedProducer = - config.create().expect("Producer creation error"); - - let mut current_job_id = Uuid::new_v4().to_string(); +fn send_frame_data( + producer: &ThreadedProducer, + fbb: &mut FlatBufferBuilder<'_>, + conf: &HowlConfig, + rng: &mut ThreadRng, + frame: u32, + now_nanos: i64, +) { + let ev44 = generate_fake_events(fbb, rng, frame, conf.event_message_config, now_nanos).to_vec(); + for _ in 0..conf.messages_per_frame { + match producer.send( + BaseRecord::to(conf.event_topic) + .key("") + .payload(if conf.fast { + ev44.as_slice() + } else { + generate_fake_events(fbb, rng, frame, conf.event_message_config, now_nanos) + }) + .timestamp(now_nanos / 1_000_000), + ) { + Ok(_) => {} + Err(err) => { + error!("Failed to send messages: {}", err.0); + } + } + } +} +fn send_run_start( + producer: &ThreadedProducer, + fbb: &mut FlatBufferBuilder<'_>, + conf: &HowlConfig, + current_job_id: &str, + now_nanos: i64, +) { producer .send( BaseRecord::to(conf.run_info_topic) .key("") .payload(generate_run_start( - &mut fbb, + fbb, conf.event_message_config.det_max, conf.event_topic, - ¤t_job_id, + current_job_id, )) .timestamp(now_nanos / 1_000_000), ) .expect("Failed to enqueue run start message"); +} + +fn send_run_stop( + producer: &ThreadedProducer, + fbb: &mut FlatBufferBuilder<'_>, + conf: &HowlConfig, + current_job_id: &str, + now_nanos: i64, +) { + match producer.send( + BaseRecord::to(conf.run_info_topic) + .key("") + .payload(generate_run_stop(fbb, current_job_id)) + .timestamp(now_nanos / 1_000_000), + ) { + Ok(_) => {} + Err(err) => { + error!("Failed to send run stop: {}", err.0); + } + } +} +fn send_veto_config( + producer: &ThreadedProducer, + fbb: &mut FlatBufferBuilder<'_>, + conf: &HowlConfig, + vetoes_mask: &u32, + now_nanos: i64, +) { + producer + .send( + BaseRecord::to(conf.veto_config_topic) + .key("") + .payload(generate_veto_config( + &conf.veto_names, + fbb, + now_nanos, + vetoes_mask, + )) + .timestamp(now_nanos / 1_000_000), + ) + .expect("Failed to enqueue run veto configuration message"); +} + +fn howl_begin( + producer: &mut ThreadedProducer, + fbb: &mut FlatBufferBuilder<'_>, + rng: &mut ThreadRng, + conf: &HowlConfig, + active_vetoes: &u32, + enabled_vetoes: &u32, +) { let target_frame_time = Duration::from_secs_f64(1.0 / conf.frames_per_second as f64); debug!("Target frame time: {target_frame_time:?}"); @@ -354,17 +440,20 @@ pub fn howl(conf: &HowlConfig) { .expect("Failed to get system time"); debug!("Target time: {target_time:?}"); + let mut current_job_id = Uuid::new_v4().to_string(); + loop { target_time += target_frame_time; debug!("New target: {target_time:?}"); frames += 1; - debug!("current job id: {current_job_id}"); produce_messages( - &producer, - &mut fbb, - &mut rng, + producer, + fbb, + rng, frames, conf, + active_vetoes, + enabled_vetoes, &mut current_job_id, ); let now = SystemTime::now() @@ -386,3 +475,53 @@ pub fn howl(conf: &HowlConfig) { } } } + +pub struct HowlConfig<'a> { + pub broker: &'a str, + pub event_topic: &'a str, + pub run_info_topic: &'a str, + pub veto_config_topic: &'a str, + pub messages_per_frame: u32, + pub frames_per_second: u32, + pub frames_per_run: u32, + pub veto_probability: Vec, + pub enabled_vetoes: Vec, + pub veto_names: Vec, + pub event_message_config: &'a EventMessageConfig, + pub fast: bool, + pub kafka_config: Option>, +} + +pub fn howl(conf: &HowlConfig) { + let mut fbb = FlatBufferBuilder::new(); + let mut rng = rand::rng(); + + let active_vetoes = get_active_vetoes(conf); + let enabled_vetoes = get_enabled_vetoes(conf, &mut rng); + + let now_nanos = SystemTime::now() + .duration_since(SystemTime::UNIX_EPOCH) + .expect("Failed to get system time") + .as_nanos() + .try_into() + .expect("This will fail after April 11th, 2262"); + + calculate_data_rate(&mut fbb, &mut rng, conf, now_nanos, &active_vetoes); + + let mut client_config: ClientConfig = ClientConfig::new(); + client_config.set("bootstrap.servers", conf.broker); + set_kafka_options(&mut client_config, &conf.kafka_config); + + // create producer + let mut producer: ThreadedProducer = + client_config.create().expect("Producer creation error"); + + howl_begin( + &mut producer, + &mut fbb, + &mut rng, + conf, + &active_vetoes, + &enabled_vetoes, + ); +} diff --git a/src/main.rs b/src/main.rs index 0ac9c45..80a2db6 100644 --- a/src/main.rs +++ b/src/main.rs @@ -8,7 +8,7 @@ use crate::cli_utils::BrokerAndOptionalTopic; use crate::cli_utils::KafkaOption; use crate::consume::ConsumeConfig; use crate::count::count; -use crate::howl::{EventMessageConfig, HowlConfig, howl}; +use crate::howl::{EventMessageConfig, HowlConfig, VETO_COUNT, howl}; use crate::sniff::sniff; use clap::{Parser, Subcommand}; use cli_utils::{BrokerAndTopic, parse_broker_spec, parse_broker_spec_optional_topic}; @@ -95,9 +95,15 @@ enum Commands { /// Maximum detector ID #[arg(long, default_value = "1000")] det_max: i32, - /// Veto probability (0 = never vetoed; 1 = always vetoed) - #[arg(long, default_value = "0.0")] - veto_probability: f64, + /// Veto probabilities + #[arg(long, num_args = 0..VETO_COUNT+1, value_delimiter = ' ')] + veto_probability: Vec, + /// Enabled vetoes + #[arg(long, num_args = 0..VETO_COUNT+1, value_delimiter = ' ')] + enabled_vetoes: Vec, + /// Veto names + #[arg(long, num_args = 0..VETO_COUNT+1, value_delimiter = ' ')] + veto_names: Vec, /// Enable howl fast mode (Disables randomised ev44 blob generation) #[arg(long, action=clap::ArgAction::SetTrue)] fast: bool, @@ -167,6 +173,8 @@ async fn main() { det_min, det_max, veto_probability, + enabled_vetoes, + veto_names, fast, kafka_config, } => howl(&HowlConfig { @@ -174,6 +182,7 @@ async fn main() { broker: &broker, event_topic: &format!("{topic_prefix}_rawEvents"), run_info_topic: &format!("{topic_prefix}_runInfo"), + veto_config_topic: &format!("{topic_prefix}_vetoConfig"), messages_per_frame, frames_per_second, frames_per_run, @@ -185,6 +194,8 @@ async fn main() { det_max, }, veto_probability, + enabled_vetoes, + veto_names, fast, }), Commands::Count { diff --git a/src/sniff.rs b/src/sniff.rs index 2a42914..95a80ab 100644 --- a/src/sniff.rs +++ b/src/sniff.rs @@ -1,5 +1,5 @@ use crate::KafkaOption; -use crate::cli_utils::BrokerAndOptionalTopic; +use crate::cli_utils::{BrokerAndOptionalTopic, set_kafka_options}; use rdkafka::ClientConfig; use rdkafka::consumer::{BaseConsumer, Consumer}; use std::time::Duration; @@ -7,16 +7,7 @@ use std::time::Duration; pub fn sniff(broker: &BrokerAndOptionalTopic, kafka_config: Option>) { let mut config = ClientConfig::new(); config.set("bootstrap.servers", broker.broker()); - - if let Some(kafka_options) = kafka_config { - for option in kafka_options { - println!( - "Setting Kafka config option {}={}", - option.key, option.value - ); - config.set(&option.key, &option.value); - } - } + set_kafka_options(&mut config, &kafka_config); let consumer: BaseConsumer = config.create().expect("Consumer creation failed");