pub mod conn; pub mod proto; pub mod search; pub mod store; use self::store::DataStore; use crate::ca::conn::ConnCommand; use crate::store::{CommonInsertItemQueue, QueryItem}; use async_channel::Sender; use conn::CaConn; use err::Error; use futures_util::StreamExt; use log::*; use netpod::{Database, ScyllaConfig}; use scylla::batch::Consistency; use serde::{Deserialize, Serialize}; use stats::{CaConnStats, CaConnStatsAgg, CaConnStatsAggDiff}; use std::collections::{BTreeMap, VecDeque}; use std::net::{Ipv4Addr, SocketAddrV4}; use std::path::PathBuf; use std::sync::atomic::{AtomicU64, Ordering}; use std::sync::{Arc, Mutex, Once}; use std::time::Duration; use tokio::fs::OpenOptions; use tokio::io::AsyncReadExt; use tokio_postgres::Client as PgClient; static mut METRICS: Option>> = None; static METRICS_ONCE: Once = Once::new(); pub fn get_metrics() -> &'static mut Option { METRICS_ONCE.call_once(|| unsafe { METRICS = Some(Mutex::new(None)); }); let mut g = unsafe { METRICS.as_mut().unwrap().lock().unwrap() }; let ret: &mut Option = &mut *g; let ret = unsafe { &mut *(ret as *mut _) }; ret } #[derive(Debug, Serialize, Deserialize)] struct ChannelConfig { channels: Vec, search: Vec, addr_bind: Ipv4Addr, addr_conn: Ipv4Addr, whitelist: String, blacklist: String, max_simul: Option, timeout: Option, #[serde(default)] abort_after_search: u32, pgconf: Database, scyconf: ScyllaConfig, array_truncate: Option, insert_worker_count: Option, insert_scylla_sessions: Option, insert_queue_max: Option, insert_item_queue_cap: Option, api_bind: Option, local_epics_hostname: String, } pub struct ListenFromFileOpts { pub config: PathBuf, } pub async fn parse_config(config: PathBuf) -> Result { let mut file = OpenOptions::new().read(true).open(config).await?; let mut buf = vec![]; file.read_to_end(&mut buf).await?; let mut conf: ChannelConfig = serde_yaml::from_slice(&buf).map_err(|e| Error::with_msg_no_trace(format!("{:?}", e)))?; let re_p = regex::Regex::new(&conf.whitelist)?; let re_n = regex::Regex::new(&conf.blacklist)?; conf.channels = conf .channels .into_iter() .filter(|ch| { if let Some(_cs) = re_p.captures(&ch) { true } else if re_n.is_match(&ch) { false } else { true } }) .collect(); Ok(CaConnectOpts { channels: conf.channels, search: conf.search, addr_bind: conf.addr_bind, addr_conn: conf.addr_conn, timeout: conf.timeout.unwrap_or(2000), abort_after_search: conf.abort_after_search, pgconf: conf.pgconf, scyconf: conf.scyconf, array_truncate: conf.array_truncate.unwrap_or(512), insert_worker_count: conf.insert_worker_count.unwrap_or(8), insert_scylla_sessions: conf.insert_scylla_sessions.unwrap_or(1), insert_queue_max: conf.insert_queue_max.unwrap_or(32), insert_item_queue_cap: conf.insert_item_queue_cap.unwrap_or(380000), api_bind: conf.api_bind.unwrap_or_else(|| "0.0.0.0:3011".into()), local_epics_hostname: conf.local_epics_hostname, }) } pub struct CaConnectOpts { pub channels: Vec, pub search: Vec, pub addr_bind: Ipv4Addr, pub addr_conn: Ipv4Addr, pub timeout: u64, pub abort_after_search: u32, pub pgconf: Database, pub scyconf: ScyllaConfig, pub array_truncate: usize, pub insert_worker_count: usize, pub insert_scylla_sessions: usize, pub insert_queue_max: usize, pub insert_item_queue_cap: usize, pub api_bind: String, pub local_epics_hostname: String, } async fn spawn_scylla_insert_workers( scyconf: ScyllaConfig, insert_scylla_sessions: usize, insert_worker_count: usize, insert_item_queue: &CommonInsertItemQueue, insert_frac: Arc, pg_client: Arc, store_stats: Arc, ) -> Result<(), Error> { let mut data_stores = vec![]; for _ in 0..insert_scylla_sessions { let scy = scylla::SessionBuilder::new() .known_nodes(&scyconf.hosts) .default_consistency(Consistency::One) .use_keyspace(&scyconf.keyspace, true) .build() .await .map_err(|e| Error::with_msg_no_trace(format!("{e:?}")))?; let scy = Arc::new(scy); let data_store = Arc::new(DataStore::new(pg_client.clone(), scy.clone()).await?); data_stores.push(data_store); } for i1 in 0..insert_worker_count { let data_store = data_stores[i1 * data_stores.len() / insert_worker_count].clone(); let stats = store_stats.clone(); let recv = insert_item_queue.receiver(); let insert_frac = insert_frac.clone(); let fut = async move { let mut i1 = 0; while let Ok(item) = recv.recv().await { match item { QueryItem::Insert(item) => { stats.store_worker_item_recv_inc(); let insert_frac = insert_frac.load(Ordering::Acquire); if i1 % 1000 < insert_frac { match crate::store::insert_item(item, &data_store, &stats).await { Ok(_) => { stats.store_worker_item_insert_inc(); } Err(e) => { stats.store_worker_item_error_inc(); // TODO introduce more structured error variants. if e.msg().contains("WriteTimeout") { tokio::time::sleep(Duration::from_millis(100)).await; } else { // TODO back off but continue. error!("insert worker sees error: {e:?}"); break; } } } } else { stats.store_worker_item_drop_inc(); } i1 += 1; } QueryItem::Mute(item) => { let values = ( (item.series & 0xff) as i32, item.series as i64, item.ts as i64, item.ema, item.emd, ); let qres = data_store .scy .query( "insert into muted (part, series, ts, ema, emd) values (?, ?, ?, ?, ?)", values, ) .await; match qres { Ok(_) => {} Err(_) => { stats.store_worker_item_error_inc(); } } } QueryItem::Ivl(item) => { let values = ( (item.series & 0xff) as i32, item.series as i64, item.ts as i64, item.ema, item.emd, ); let qres = data_store .scy .query( "insert into item_recv_ivl (part, series, ts, ema, emd) values (?, ?, ?, ?, ?)", values, ) .await; match qres { Ok(_) => {} Err(_) => { stats.store_worker_item_error_inc(); } } } } } }; tokio::spawn(fut); } Ok(()) } pub struct CommandQueueSet { queues: tokio::sync::Mutex>>, } impl CommandQueueSet { pub fn new() -> Self { Self { queues: tokio::sync::Mutex::new(VecDeque::>::new()), } } pub fn queues(&self) -> &tokio::sync::Mutex>> { &self.queues } } pub async fn ca_connect(opts: ListenFromFileOpts) -> Result<(), Error> { let facility = "scylla"; let insert_frac = Arc::new(AtomicU64::new(1000)); let insert_ivl_min = Arc::new(AtomicU64::new(8800)); let opts = parse_config(opts.config).await?; let scyconf = opts.scyconf.clone(); let command_queue_set = Arc::new(CommandQueueSet::new()); tokio::spawn(crate::metrics::start_metrics_service( opts.api_bind.clone(), insert_frac.clone(), insert_ivl_min.clone(), command_queue_set.clone(), )); let d = Database { name: opts.pgconf.name.clone(), host: opts.pgconf.host.clone(), port: opts.pgconf.port.clone(), user: opts.pgconf.user.clone(), pass: opts.pgconf.pass.clone(), }; let (pg_client, pg_conn) = tokio_postgres::connect( &format!("postgresql://{}:{}@{}:{}/{}", d.user, d.pass, d.host, d.port, d.name), tokio_postgres::tls::NoTls, ) .await .unwrap(); // TODO allow clean shutdown on ctrl-c and join the pg_conn in the end: tokio::spawn(pg_conn); let pg_client = Arc::new(pg_client); let scy = scylla::SessionBuilder::new() .known_nodes(&scyconf.hosts) .default_consistency(Consistency::One) .use_keyspace(scyconf.keyspace, true) .build() .await .map_err(|e| Error::with_msg_no_trace(format!("{e:?}")))?; let scy = Arc::new(scy); // TODO use new struct: let local_stats = Arc::new(CaConnStats::new()); // TODO factor the find loop into a separate Stream. let qu_find_addr = pg_client .prepare("with q1 as (select t1.facility, t1.channel, t1.addr from ioc_by_channel t1 where t1.facility = $1 and t1.channel in ($2, $3, $4, $5, $6, $7, $8, $9) and t1.addr != '' order by t1.tsmod desc) select distinct on (q1.facility, q1.channel) q1.facility, q1.channel, q1.addr from q1") .await .map_err(|e| Error::with_msg_no_trace(format!("{e:?}")))?; let mut channels_by_host = BTreeMap::new(); let mut chns_todo = &opts.channels[..]; let mut chstmp = ["__NONE__"; 8]; let mut ix = 0; while chns_todo.len() > 0 { for (s1, s2) in chns_todo.iter().zip(chstmp.iter_mut()) { *s2 = s1; } chns_todo = &chns_todo[chstmp.len().min(chns_todo.len())..]; let rows = pg_client .query( &qu_find_addr, &[ &facility, &chstmp[0], &chstmp[1], &chstmp[2], &chstmp[3], &chstmp[4], &chstmp[5], &chstmp[6], &chstmp[7], ], ) .await .map_err(|e| Error::with_msg_no_trace(format!("pg lookup error: {e:?}")))?; if rows.is_empty() { error!("can not find any addresses of channels {:?}", chstmp); } else { for row in rows { let ch: &str = row.get(1); let addr: &str = row.get(2); if addr == "" { // TODO the address was searched before but could not be found. } else { let addr: SocketAddrV4 = match addr.parse() { Ok(k) => { local_stats.ioc_lookup_inc(); k } Err(e) => { error!("can not parse {addr:?} for channel {ch:?} {e:?}"); continue; } }; ix += 1; if ix % 1000 == 0 { info!("{} of {} {} {:?}", ix, opts.channels.len(), ch, addr); } if !channels_by_host.contains_key(&addr) { channels_by_host.insert(addr, vec![ch.to_string()]); } else { channels_by_host.get_mut(&addr).unwrap().push(ch.to_string()); } } } } } if opts.abort_after_search == 1 { return Ok(()); } let data_store = Arc::new(DataStore::new(pg_client.clone(), scy.clone()).await?); let insert_item_queue = CommonInsertItemQueue::new(opts.insert_item_queue_cap); // TODO use a new stats struct let store_stats = Arc::new(CaConnStats::new()); spawn_scylla_insert_workers( opts.scyconf.clone(), opts.insert_scylla_sessions, opts.insert_worker_count, &insert_item_queue, insert_frac.clone(), pg_client.clone(), store_stats.clone(), ) .await?; let mut conn_jhs = vec![]; let mut conn_stats = vec![]; info!("channels_by_host len {}", channels_by_host.len()); for (host, channels) in channels_by_host { let data_store = data_store.clone(); let addr = SocketAddrV4::new(host.ip().clone(), host.port()); let mut conn = CaConn::new( addr, opts.local_epics_hostname.clone(), data_store.clone(), insert_item_queue.sender(), opts.array_truncate, opts.insert_queue_max, insert_ivl_min.clone(), ); conn_stats.push(conn.stats()); for c in channels { conn.channel_add(c); } let stats2 = conn.stats(); let conn_command_tx = conn.conn_command_tx(); command_queue_set.queues().lock().await.push_back(conn_command_tx); let conn_block = async move { while let Some(item) = conn.next().await { match item { Ok(_) => { stats2.conn_item_count_inc(); } Err(e) => { error!("CaConn gives error: {e:?}"); break; } } } Ok::<_, Error>(()) }; let jh = tokio::spawn(conn_block); conn_jhs.push(jh); } let mut agg_last = CaConnStatsAgg::new(); loop { tokio::time::sleep(Duration::from_millis(671)).await; let agg = CaConnStatsAgg::new(); agg.push(&local_stats); agg.push(&store_stats); for g in &conn_stats { agg.push(&g); } let m = get_metrics(); *m = Some(agg.clone()); if false { let diff = CaConnStatsAggDiff::diff_from(&agg_last, &agg); info!("{}", diff.display()); } for _s1 in &conn_stats {} agg_last = agg; if false { break; } } for jh in conn_jhs { match jh.await { Ok(k) => match k { Ok(_) => {} Err(e) => { error!("{e:?}"); } }, Err(e) => { error!("{e:?}"); } } } Ok(()) }