From a7b76c986822607534d920ec9bb4f75a154ce8c3 Mon Sep 17 00:00:00 2001 From: Dominik Werder Date: Mon, 30 Oct 2023 10:40:45 +0100 Subject: [PATCH] WIP adapt to async-channel 2.0.0 --- dbpg/src/seriesbychannel.rs | 3 +- ingest-bsread/src/bsreadclient.rs | 3 +- ingest-bsread/src/zmtp.rs | 3 +- ingest-bsread/src/zmtp/dumper.rs | 3 +- ingest-bsread/src/zmtp/zmtpproto.rs | 3 +- netfetch/src/ca/conn.rs | 11 ++++-- netfetch/src/ca/connset.rs | 55 +++++++++++++------------- netfetch/src/ca/connset_input_merge.rs | 8 ++-- netfetch/src/ca/finder.rs | 8 ++-- netfetch/src/ca/findioc.rs | 8 ++-- netfetch/src/ca/search.rs | 3 +- netfetch/src/senderpolling.rs | 17 ++++---- scywr/src/insertworker.rs | 2 + 13 files changed, 71 insertions(+), 56 deletions(-) diff --git a/dbpg/src/seriesbychannel.rs b/dbpg/src/seriesbychannel.rs index b5a9477..47855bc 100644 --- a/dbpg/src/seriesbychannel.rs +++ b/dbpg/src/seriesbychannel.rs @@ -289,7 +289,8 @@ impl Worker { } async fn work(&mut self) -> Result<(), Error> { - while let Some(batch) = self.batch_rx.next().await { + let batch_rx = &self.batch_rx; + while let Ok(batch) = batch_rx.recv().await { self.stats.recv_batch().inc(); self.stats.recv_items().add(batch.len() as _); for x in &batch { diff --git a/ingest-bsread/src/bsreadclient.rs b/ingest-bsread/src/bsreadclient.rs index 33638e5..0d41058 100644 --- a/ingest-bsread/src/bsreadclient.rs +++ b/ingest-bsread/src/bsreadclient.rs @@ -214,7 +214,8 @@ impl BsreadClient { if let Some(v) = self.rcvbuf { ingest_linux::net::set_rcv_sock_opts(&mut conn, v as u32)?; } - let mut zmtp = Zmtp::new(conn, SocketType::PULL); + let zmtp = Zmtp::new(conn, SocketType::PULL); + let mut zmtp = Box::pin(zmtp); let mut i1 = 0u64; let mut msgc = 0u64; let mut dh_md5_last = String::new(); diff --git a/ingest-bsread/src/zmtp.rs b/ingest-bsread/src/zmtp.rs index 3b2282d..5491458 100644 --- a/ingest-bsread/src/zmtp.rs +++ b/ingest-bsread/src/zmtp.rs @@ -55,7 +55,8 @@ fn test_service() -> Result<(), Error> { info!("accepting..."); let (conn, remote) = sock.accept().await?; info!("new connection from {:?}", remote); - let mut zmtp = Zmtp::new(conn, SocketType::PUSH); + let zmtp = Zmtp::new(conn, SocketType::PUSH); + let mut zmtp = Box::pin(zmtp); let fut = async move { while let Some(item) = zmtp.next().await { info!("item from {:?} {:?}", remote, item); diff --git a/ingest-bsread/src/zmtp/dumper.rs b/ingest-bsread/src/zmtp/dumper.rs index 50cdf74..4d3341f 100644 --- a/ingest-bsread/src/zmtp/dumper.rs +++ b/ingest-bsread/src/zmtp/dumper.rs @@ -48,7 +48,8 @@ impl BsreadDumper { self.source_addr.clone() }; let conn = tokio::net::TcpStream::connect(&src).await?; - let mut zmtp = Zmtp::new(conn, SocketType::PULL); + let zmtp = Zmtp::new(conn, SocketType::PULL); + let mut zmtp = Box::pin(zmtp); let mut i1 = 0u64; let mut msgc = 0u64; let mut dh_md5_last = String::new(); diff --git a/ingest-bsread/src/zmtp/zmtpproto.rs b/ingest-bsread/src/zmtp/zmtpproto.rs index 43b69c5..7b7f329 100644 --- a/ingest-bsread/src/zmtp/zmtpproto.rs +++ b/ingest-bsread/src/zmtp/zmtpproto.rs @@ -214,7 +214,8 @@ impl Zmtp { let mut item_count = 0; // TODO should I better keep one serialized item in Self so that I know how much space it needs? let serialized: Int> = if self.out_enable && self.outbuf.wcap() >= self.outbuf.cap() / 2 { - match self.data_rx.poll_next_unpin(cx) { + let data_rx = std::pin::pin!(self.data_rx); + match data_rx.poll_next(cx) { Ready(Some(_item)) => { // TODO item should be something that we can convert into a zmtp message. Int::Empty diff --git a/netfetch/src/ca/conn.rs b/netfetch/src/ca/conn.rs index a058a83..8704646 100644 --- a/netfetch/src/ca/conn.rs +++ b/netfetch/src/ca/conn.rs @@ -3,6 +3,7 @@ use super::ExtraInsertsConf; use crate::senderpolling::SenderPolling; use crate::throttletrace::ThrottleTrace; use crate::timebin::ConnTimeBin; +use async_channel::Receiver; use async_channel::Sender; use core::fmt; use dbpg::seriesbychannel::CanSendChannelInfoResult; @@ -47,6 +48,7 @@ use std::collections::BTreeMap; use std::collections::VecDeque; use std::net::SocketAddrV4; use std::ops::ControlFlow; +use std::pin::pin; use std::pin::Pin; use std::sync::atomic; use std::sync::atomic::AtomicUsize; @@ -507,8 +509,8 @@ pub struct CaConn { local_epics_hostname: String, stats: Arc, insert_ivl_min_mus: u64, - conn_command_tx: async_channel::Sender, - conn_command_rx: async_channel::Receiver, + conn_command_tx: Sender, + conn_command_rx: Receiver, conn_backoff: f32, conn_backoff_beg: f32, inserts_counter: u64, @@ -792,7 +794,7 @@ impl CaConn { if self.is_shutdown() { Ok(Ready(None)) } else { - match self.conn_command_rx.poll_next_unpin(cx) { + match pin!(self.conn_command_rx).poll_next(cx) { Ready(Some(a)) => { trace3!("handle_conn_command received a command {}", self.remote_addr_dbg); match a.kind { @@ -1886,7 +1888,8 @@ impl CaConn { fn handle_own_ticker_tick(self: Pin<&mut Self>, _cx: &mut Context) -> Result<(), Error> { // debug!("tick CaConn {}", self.remote_addr_dbg); let tsnow = Instant::now(); - let this = self.get_mut(); + // TODO use safe version + let this = unsafe { self.get_unchecked_mut() }; match &this.state { CaConnState::Unconnected(since) => {} CaConnState::Connecting(since, _addr, _) => { diff --git a/netfetch/src/ca/connset.rs b/netfetch/src/ca/connset.rs index ea0f266..aced6de 100644 --- a/netfetch/src/ca/connset.rs +++ b/netfetch/src/ca/connset.rs @@ -54,6 +54,7 @@ use std::collections::BTreeMap; use std::collections::VecDeque; use std::net::SocketAddr; use std::net::SocketAddrV4; +use std::pin::pin; use std::pin::Pin; use std::sync::atomic; use std::sync::Arc; @@ -330,27 +331,28 @@ impl CanSendChannelInfoResult for SeriesLookupSender { } } +#[pin_project::pin_project] pub struct CaConnSet { backend: String, local_epics_hostname: String, ca_conn_ress: BTreeMap, channel_states: ChannelStateMap, - connset_inp_rx: Receiver, + connset_inp_rx: Pin>>, channel_info_query_queue: VecDeque, channel_info_query_sender: SenderPolling, channel_info_query_tx: Option>, - channel_info_res_tx: Sender>, - channel_info_res_rx: Receiver>, + channel_info_res_tx: Pin>>>, + channel_info_res_rx: Pin>>>, find_ioc_query_queue: VecDeque, find_ioc_query_sender: SenderPolling, - find_ioc_res_rx: Receiver>, - storage_insert_tx: Sender, + find_ioc_res_rx: Pin>>>, + storage_insert_tx: Pin>>, storage_insert_queue: VecDeque, storage_insert_sender: SenderPolling, - ca_conn_res_tx: Sender<(SocketAddr, CaConnEvent)>, - ca_conn_res_rx: Receiver<(SocketAddr, CaConnEvent)>, + ca_conn_res_tx: Pin>>, + ca_conn_res_rx: Pin>>, connset_out_queue: VecDeque, - connset_out_tx: Sender, + connset_out_tx: Pin>>, shutdown_stopping: bool, shutdown_done: bool, chan_check_next: Option, @@ -396,26 +398,26 @@ impl CaConnSet { local_epics_hostname, ca_conn_ress: BTreeMap::new(), channel_states: ChannelStateMap::new(), - connset_inp_rx, + connset_inp_rx: Box::pin(connset_inp_rx), channel_info_query_queue: VecDeque::new(), channel_info_query_sender: SenderPolling::new(channel_info_query_tx.clone()), channel_info_query_tx: Some(channel_info_query_tx), - channel_info_res_tx, - channel_info_res_rx, + channel_info_res_tx: Box::pin(channel_info_res_tx), + channel_info_res_rx: Box::pin(channel_info_res_rx), find_ioc_query_queue: VecDeque::new(), find_ioc_query_sender: SenderPolling::new(find_ioc_query_tx), - find_ioc_res_rx, - storage_insert_tx: storage_insert_tx.clone(), + find_ioc_res_rx: Box::pin(find_ioc_res_rx), + storage_insert_tx: Box::pin(storage_insert_tx.clone()), storage_insert_queue: VecDeque::new(), storage_insert_sender: SenderPolling::new(storage_insert_tx), - ca_conn_res_tx, - ca_conn_res_rx, + ca_conn_res_tx: Box::pin(ca_conn_res_tx), + ca_conn_res_rx: Box::pin(ca_conn_res_rx), shutdown_stopping: false, shutdown_done: false, chan_check_next: None, stats: stats.clone(), ca_conn_stats: ca_conn_stats.clone(), - connset_out_tx, + connset_out_tx: Box::pin(connset_out_tx), connset_out_queue: VecDeque::new(), // connset_out_sender: SenderPolling::new(connset_out_tx), ioc_finder_jh, @@ -542,14 +544,13 @@ impl CaConnSet { running_cmd_id: None, health_timeout_count: 0, }); + let tx = self.channel_info_res_tx.as_ref().get_ref().clone(); let item = ChannelInfoQuery { backend: cmd.backend, channel: cmd.name, scalar_type: CHANNEL_STATUS_DUMMY_SCALAR_TYPE, shape_dims: Vec::new(), - tx: Box::pin(SeriesLookupSender { - tx: self.channel_info_res_tx.clone(), - }), + tx: Box::pin(SeriesLookupSender { tx }), }; self.channel_info_query_queue.push_back(item); Ok(()) @@ -965,7 +966,7 @@ impl CaConnSet { add.backend.clone(), addr_v4, add.local_epics_hostname, - self.storage_insert_tx.clone(), + self.storage_insert_tx.as_ref().get_ref().clone(), self.channel_info_query_tx .clone() .ok_or_else(|| Error::with_msg_no_trace("no more channel_info_query_tx available"))?, @@ -974,8 +975,8 @@ impl CaConnSet { ); let conn_tx = conn.conn_command_tx(); let conn_stats = conn.stats(); - let tx1 = self.ca_conn_res_tx.clone(); - let tx2 = self.storage_insert_tx.clone(); + let tx1 = self.ca_conn_res_tx.as_ref().get_ref().clone(); + let tx2 = self.storage_insert_tx.as_ref().get_ref().clone(); let jh = tokio::spawn(Self::ca_conn_item_merge(conn, tx1, tx2, addr, self.stats.clone())); let ca_conn_res = CaConnRes { state: CaConnState::new(CaConnStateValue::Fresh), @@ -997,7 +998,7 @@ impl CaConnSet { stats.ca_conn_task_begin().inc(); trace2!("ca_conn_consumer begin {}", addr); let connstats = conn.stats(); - let mut conn = conn; + let mut conn = Box::pin(conn); let mut ret = Ok(()); while let Some(item) = conn.next().await { match item { @@ -1579,7 +1580,7 @@ impl Stream for CaConnSet { } } - match self.find_ioc_res_rx.poll_next_unpin(cx) { + match pin!(self.find_ioc_res_rx).poll_next(cx) { Ready(Some(x)) => match self.handle_ioc_query_result(x) { Ok(()) => { have_progress = true; @@ -1592,7 +1593,7 @@ impl Stream for CaConnSet { } } - match self.ca_conn_res_rx.poll_next_unpin(cx) { + match pin!(self.ca_conn_res_rx).poll_next(cx) { Ready(Some((addr, ev))) => match self.handle_ca_conn_event(addr, ev) { Ok(()) => { have_progress = true; @@ -1605,7 +1606,7 @@ impl Stream for CaConnSet { } } - match self.channel_info_res_rx.poll_next_unpin(cx) { + match pin!(self.channel_info_res_rx).poll_next(cx) { Ready(Some(x)) => match self.handle_series_lookup_result(x) { Ok(()) => { have_progress = true; @@ -1618,7 +1619,7 @@ impl Stream for CaConnSet { } } - match self.connset_inp_rx.poll_next_unpin(cx) { + match pin!(self.connset_inp_rx).poll_next(cx) { Ready(Some(x)) => match self.handle_event(x) { Ok(()) => { have_progress = true; diff --git a/netfetch/src/ca/connset_input_merge.rs b/netfetch/src/ca/connset_input_merge.rs index 3662e9e..b040d88 100644 --- a/netfetch/src/ca/connset_input_merge.rs +++ b/netfetch/src/ca/connset_input_merge.rs @@ -4,8 +4,8 @@ use crate::ca::connset::ConnSetCmd; use async_channel::Receiver; use dbpg::seriesbychannel::ChannelInfoResult; use err::Error; -use futures_util::StreamExt; use std::collections::VecDeque; +use std::pin::pin; use std::pin::Pin; use std::task::Context; use std::task::Poll; @@ -43,7 +43,7 @@ impl futures_util::Stream for InputMerge { use Poll::*; let ret = { if let Some(inp) = &mut self.inp3 { - match inp.poll_next_unpin(cx) { + match pin!(*inp).poll_next(cx) { Ready(Some(x)) => Some(CaConnSetEvent::ConnSetCmd(todo!())), Ready(None) => { self.inp2 = None; @@ -59,7 +59,7 @@ impl futures_util::Stream for InputMerge { Some(x) } else { if let Some(inp) = &mut self.inp2 { - match inp.poll_next_unpin(cx) { + match pin!(*inp).poll_next(cx) { Ready(Some(x)) => Some(CaConnSetEvent::ConnSetCmd(todo!())), Ready(None) => { self.inp2 = None; @@ -75,7 +75,7 @@ impl futures_util::Stream for InputMerge { Ready(Some(x)) } else { if let Some(inp) = &mut self.inp1 { - match inp.poll_next_unpin(cx) { + match pin!(*inp).poll_next(cx) { Ready(Some(x)) => Ready(Some(x)), Ready(None) => { self.inp1 = None; diff --git a/netfetch/src/ca/finder.rs b/netfetch/src/ca/finder.rs index 5651faf..95b8730 100644 --- a/netfetch/src/ca/finder.rs +++ b/netfetch/src/ca/finder.rs @@ -251,14 +251,14 @@ async fn finder_worker_single( } async fn finder_network_if_not_found( - mut rx: Receiver>, + rx: Receiver>, tx: Sender>, opts: CaIngestOpts, stats: Arc, ) -> Result<(), Error> { let (net_tx, net_rx, jh, jhs) = ca_search_workers_start(&opts, stats.clone()).await.unwrap(); let jh2 = taskrun::spawn(process_net_result(net_rx, tx.clone(), opts.clone())); - 'outer: while let Some(item) = rx.next().await { + 'outer: while let Ok(item) = rx.recv().await { let mut res = VecDeque::new(); let mut net = VecDeque::new(); for e in item { @@ -287,7 +287,7 @@ async fn finder_network_if_not_found( } async fn process_net_result( - mut net_rx: Receiver, Error>>, + net_rx: Receiver, Error>>, tx: Sender>, opts: CaIngestOpts, ) -> Result<(), Error> { @@ -304,7 +304,7 @@ async fn process_net_result( ioc_search_index_worker_jhs.push(jh); } drop(dbrx); - while let Some(item) = net_rx.next().await { + while let Ok(item) = net_rx.recv().await { match item { Ok(item) => { for e in item.iter() { diff --git a/netfetch/src/ca/findioc.rs b/netfetch/src/ca/findioc.rs index f91512d..dbabc17 100644 --- a/netfetch/src/ca/findioc.rs +++ b/netfetch/src/ca/findioc.rs @@ -7,7 +7,6 @@ use err::Error; use futures_util::Future; use futures_util::FutureExt; use futures_util::Stream; -use futures_util::StreamExt; use libc::c_int; use log::*; use stats::IocFinderStats; @@ -88,7 +87,7 @@ pub struct FindIocRes { pub struct FindIocStream { tgts: Vec, - channels_input: Receiver, + channels_input: Pin>>, in_flight: BTreeMap, in_flight_max: usize, bid_by_sid: BTreeMap, @@ -129,7 +128,7 @@ impl FindIocStream { let afd = AsyncFd::new(sock.0).unwrap(); Self { tgts, - channels_input, + channels_input: Box::pin(channels_input), in_flight: BTreeMap::new(), bid_by_sid: BTreeMap::new(), batch_send_queue: VecDeque::new(), @@ -548,7 +547,8 @@ impl FindIocStream { use Poll::*; let mut ret = Vec::new(); loop { - break match self.channels_input.poll_next_unpin(cx) { + let rx = self.channels_input.as_mut(); + break match rx.poll_next(cx) { Ready(Some(item)) => { ret.push(item); if ret.len() < self.channels_per_batch { diff --git a/netfetch/src/ca/search.rs b/netfetch/src/ca/search.rs index d74537a..af1782a 100644 --- a/netfetch/src/ca/search.rs +++ b/netfetch/src/ca/search.rs @@ -238,7 +238,8 @@ async fn search_tgts_from_opts(opts: &CaIngestOpts) -> Result<(Vec Ok((addrs, blacklist)) } -async fn finder_run(mut finder: FindIocStream, tx: Sender, Error>>) -> Result<(), Error> { +async fn finder_run(finder: FindIocStream, tx: Sender, Error>>) -> Result<(), Error> { + let mut finder = Box::pin(finder); while let Some(item) = finder.next().await { if let Err(_) = tx.send(item).await { break; diff --git a/netfetch/src/senderpolling.rs b/netfetch/src/senderpolling.rs index c6fe7d3..4c6079d 100644 --- a/netfetch/src/senderpolling.rs +++ b/netfetch/src/senderpolling.rs @@ -2,7 +2,6 @@ use async_channel::Send; use async_channel::Sender; use err::thiserror; use futures_util::Future; -use futures_util::FutureExt; use pin_project::pin_project; use std::marker::PhantomPinned; use std::pin::Pin; @@ -24,6 +23,7 @@ where { sender: Option>>, sender_ptr: NonNull>, + #[pin] fut: Option>, _pin: PhantomPinned, } @@ -81,20 +81,23 @@ impl SenderPolling { } } -impl Future for SenderPolling { +impl Future for SenderPolling +where + T: Unpin, +{ type Output = Result<(), Error>; - fn poll(self: Pin<&mut Self>, cx: &mut Context) -> Poll { + fn poll(mut self: Pin<&mut Self>, cx: &mut Context) -> Poll { use Poll::*; let this = self.project(); - match this.fut { - Some(fut) => match fut.poll_unpin(cx) { + match this.fut.as_pin_mut() { + Some(fut) => match fut.poll(cx) { Ready(Ok(())) => { - *this.fut = None; + self.fut = None; Ready(Ok(())) } Ready(Err(e)) => { - *this.fut = None; + self.fut = None; Ready(Err(Error::Closed(e.0))) } Pending => Pending, diff --git a/scywr/src/insertworker.rs b/scywr/src/insertworker.rs index 42f5b3a..08d6590 100644 --- a/scywr/src/insertworker.rs +++ b/scywr/src/insertworker.rs @@ -334,6 +334,8 @@ async fn worker_streamed( insert_worker_opts .insert_workers_running .fetch_add(1, atomic::Ordering::AcqRel); + // TODO possible without box? + let item_inp = Box::pin(item_inp); let mut stream = item_inp .map(|item| { stats.item_recv.inc();