use crate::create_connection; use async_channel::Receiver; use async_channel::RecvError; use async_channel::Sender; use err::thiserror; use err::ThisError; use netpod::log::*; use netpod::range::evrange::NanoRange; use netpod::ChConf; use netpod::ChannelSearchQuery; use netpod::ChannelSearchResult; use netpod::Database; use netpod::SfDbChannel; use taskrun::tokio; use tokio::task::JoinHandle; use tokio_postgres::Client; #[derive(Debug, ThisError)] #[cstm(name = "PgWorker")] pub enum Error { Error(#[from] err::Error), ChannelSend, ChannelRecv, Join, } impl From for Error { fn from(_value: RecvError) -> Self { Self::ChannelRecv } } impl err::ToErr for Error { fn to_err(self) -> err::Error { err::Error::from_string(self) } } #[derive(Debug)] enum Job { ChConfBestMatchingNameRange(String, String, NanoRange, Sender>), ChConfForSeries(String, u64, Sender>), InfoForSeriesIds( Vec, Sender>, crate::channelinfo::Error>>, ), SearchChannel(ChannelSearchQuery, Sender>), SfChannelBySeries( netpod::SfDbChannel, Sender>, ), } #[derive(Debug, Clone)] pub struct PgQueue { tx: Sender, } impl PgQueue { pub async fn chconf_for_series( &self, backend: &str, series: u64, ) -> Result>, Error> { let (tx, rx) = async_channel::bounded(1); let job = Job::ChConfForSeries(backend.into(), series, tx); self.tx.send(job).await.map_err(|_| Error::ChannelSend)?; Ok(rx) } pub async fn chconf_best_matching_name_range( &self, backend: &str, name: &str, range: NanoRange, ) -> Result>, Error> { let (tx, rx) = async_channel::bounded(1); let job = Job::ChConfBestMatchingNameRange(backend.into(), name.into(), range, tx); self.tx.send(job).await.map_err(|_| Error::ChannelSend)?; Ok(rx) } pub async fn info_for_series_ids( &self, series_ids: Vec, ) -> Result>, crate::channelinfo::Error>>, Error> { let (tx, rx) = async_channel::bounded(1); let job = Job::InfoForSeriesIds(series_ids, tx); self.tx.send(job).await.map_err(|_| Error::ChannelSend)?; Ok(rx) } pub async fn search_channel_scylla( &self, query: ChannelSearchQuery, ) -> Result, Error> { let (tx, rx) = async_channel::bounded(1); let job = Job::SearchChannel(query, tx); self.tx.send(job).await.map_err(|_| Error::ChannelSend)?; let ret = rx.recv().await?; Ok(ret) } pub async fn find_sf_channel_by_series( &self, query: netpod::SfDbChannel, ) -> Result, Error> { let (tx, rx) = async_channel::bounded(1); let job = Job::SfChannelBySeries(query, tx); self.tx.send(job).await.map_err(|_| Error::ChannelSend)?; let ret = rx.recv().await?; Ok(ret) } } #[derive(Debug)] pub struct PgWorker { rx: Receiver, pg: Client, pgjh: Option>>, } impl PgWorker { pub async fn new(pgconf: &Database) -> Result<(PgQueue, Self), Error> { let (tx, rx) = async_channel::bounded(64); let (pg, pgjh) = create_connection(pgconf).await?; let queue = PgQueue { tx }; let worker = Self { rx, pg, pgjh: Some(pgjh), }; Ok((queue, worker)) } pub async fn work(self) -> Result<(), Error> { loop { let x = self.rx.recv().await; let job = match x { Ok(x) => x, Err(_) => { error!("PgWorker can not receive from channel"); return Err(Error::ChannelRecv); } }; match job { Job::ChConfBestMatchingNameRange(backend, name, range, tx) => { let res = crate::channelconfig::chconf_best_matching_for_name_and_range(&backend, &name, range, &self.pg) .await; if tx.send(res.map_err(Into::into)).await.is_err() { // TODO count for stats } } Job::ChConfForSeries(backend, series, tx) => { let res = crate::channelconfig::chconf_for_series(&backend, series, &self.pg).await; if tx.send(res.map_err(Into::into)).await.is_err() { // TODO count for stats } } Job::InfoForSeriesIds(ids, tx) => { let res = crate::channelinfo::info_for_series_ids(&ids, &self.pg).await; if tx.send(res.map_err(Into::into)).await.is_err() { // TODO count for stats } } Job::SearchChannel(query, tx) => { let res = crate::search::search_channel_scylla(query, &self.pg).await; if tx.send(res.map_err(Into::into)).await.is_err() { // TODO count for stats } } Job::SfChannelBySeries(query, tx) => { let res = find_sf_channel_by_series(query, &self.pg).await; if tx.send(res.map_err(Into::into)).await.is_err() { // TODO count for stats } } } } } pub async fn join(&mut self) -> Result<(), Error> { if let Some(jh) = self.pgjh.take() { jh.await.map_err(|_| Error::Join)?.map_err(Error::from)?; Ok(()) } else { Ok(()) } } pub fn close(&self) { self.rx.close(); } } // On sf-databuffer, the channel name identifies the series. But we can also have a series id. // This function is used if the request provides only the series-id, but no name. async fn find_sf_channel_by_series( channel: SfDbChannel, pgclient: &Client, ) -> Result { use crate::FindChannelError; debug!("find_sf_channel_by_series {:?}", channel); let series = channel.series().ok_or_else(|| FindChannelError::BadSeriesId)?; let sql = "select rowid from facilities where name = $1"; let rows = pgclient .query(sql, &[&channel.backend()]) .await .map_err(|e| FindChannelError::Database(e.to_string()))?; let row = rows .into_iter() .next() .ok_or_else(|| FindChannelError::UnknownBackend)?; let backend_id: i64 = row.get(0); let sql = "select name from channels where facility = $1 and rowid = $2"; let rows = pgclient .query(sql, &[&backend_id, &(series as i64)]) .await .map_err(|e| FindChannelError::Database(e.to_string()))?; if rows.len() > 1 { return Err(FindChannelError::MultipleFound); } if let Some(row) = rows.into_iter().next() { let name = row.get::<_, String>(0); let channel = SfDbChannel::from_full(channel.backend(), channel.series(), name); Ok(channel) } else { return Err(FindChannelError::NoFound); } }