|
|
|
@ -15,6 +15,8 @@ use tokio::{
|
|
|
|
|
self,
|
|
|
|
|
error::SendError,
|
|
|
|
|
},
|
|
|
|
|
watch,
|
|
|
|
|
Notify,
|
|
|
|
|
},
|
|
|
|
|
task::JoinHandle,
|
|
|
|
|
time::{
|
|
|
|
@ -24,11 +26,14 @@ use tokio::{
|
|
|
|
|
};
|
|
|
|
|
use futures::StreamExt;
|
|
|
|
|
|
|
|
|
|
pub const DEFAULT_TIMEOUT: Duration= Duration::from_secs(5);
|
|
|
|
|
|
|
|
|
|
/// Settings for chain handler
|
|
|
|
|
#[derive(Debug, Clone, PartialEq)]
|
|
|
|
|
pub struct Settings
|
|
|
|
|
{
|
|
|
|
|
pub backlog: usize,
|
|
|
|
|
pub internal_backlog: usize,
|
|
|
|
|
pub capacity: usize,
|
|
|
|
|
pub timeout: Duration,
|
|
|
|
|
pub throttle: Option<Duration>,
|
|
|
|
@ -38,7 +43,7 @@ pub struct Settings
|
|
|
|
|
impl Settings
|
|
|
|
|
{
|
|
|
|
|
/// Should we keep this string.
|
|
|
|
|
#[inline] fn matches(&self, s: &str) -> bool
|
|
|
|
|
#[inline] fn matches(&self, _s: &str) -> bool
|
|
|
|
|
{
|
|
|
|
|
true
|
|
|
|
|
}
|
|
|
|
@ -51,6 +56,7 @@ impl Default for Settings
|
|
|
|
|
{
|
|
|
|
|
Self {
|
|
|
|
|
backlog: 32,
|
|
|
|
|
internal_backlog: 8,
|
|
|
|
|
capacity: 4,
|
|
|
|
|
timeout: Duration::from_secs(5),
|
|
|
|
|
throttle: Some(Duration::from_millis(200)),
|
|
|
|
@ -64,6 +70,7 @@ impl Default for Settings
|
|
|
|
|
struct HostInner<T>
|
|
|
|
|
{
|
|
|
|
|
input: mpsc::Receiver<Vec<T>>,
|
|
|
|
|
shutdown: watch::Receiver<bool>,
|
|
|
|
|
}
|
|
|
|
|
|
|
|
|
|
#[derive(Debug)]
|
|
|
|
@ -72,7 +79,10 @@ struct Handle<T: Send+ chain::Chainable>
|
|
|
|
|
chain: RwLock<chain::Chain<T>>,
|
|
|
|
|
input: mpsc::Sender<Vec<T>>,
|
|
|
|
|
opt: Settings,
|
|
|
|
|
|
|
|
|
|
notify_write: Arc<Notify>,
|
|
|
|
|
push_now: Arc<Notify>,
|
|
|
|
|
shutdown: watch::Sender<bool>,
|
|
|
|
|
|
|
|
|
|
/// Data used only for the worker task.
|
|
|
|
|
host: msg::Once<HostInner<T>>,
|
|
|
|
|
}
|
|
|
|
@ -80,22 +90,23 @@ struct Handle<T: Send+ chain::Chainable>
|
|
|
|
|
#[derive(Clone, Debug)]
|
|
|
|
|
pub struct ChainHandle<T: Send + chain::Chainable>(Arc<Box<Handle<T>>>);
|
|
|
|
|
|
|
|
|
|
impl<T: Send+ chain::Chainable> ChainHandle<T>
|
|
|
|
|
impl<T: Send+ chain::Chainable + 'static> ChainHandle<T>
|
|
|
|
|
{
|
|
|
|
|
#[inline] pub fn new(chain: chain::Chain<T>) -> Self
|
|
|
|
|
{
|
|
|
|
|
Self::with_settings(chain, Default::default())
|
|
|
|
|
}
|
|
|
|
|
pub fn with_settings(chain: chain::Chain<T>, opt: Settings) -> Self
|
|
|
|
|
{
|
|
|
|
|
let (shutdown_tx, shutdown) = watch::channel(false);
|
|
|
|
|
let (itx, irx) = mpsc::channel(opt.backlog);
|
|
|
|
|
Self(Arc::new(Box::new(Handle{
|
|
|
|
|
chain: RwLock::new(chain),
|
|
|
|
|
input: itx,
|
|
|
|
|
opt,
|
|
|
|
|
push_now: Arc::new(Notify::new()),
|
|
|
|
|
notify_write: Arc::new(Notify::new()),
|
|
|
|
|
shutdown: shutdown_tx,
|
|
|
|
|
|
|
|
|
|
host: msg::Once::new(HostInner{
|
|
|
|
|
input: irx,
|
|
|
|
|
shutdown,
|
|
|
|
|
})
|
|
|
|
|
})))
|
|
|
|
|
}
|
|
|
|
@ -122,10 +133,51 @@ impl<T: Send+ chain::Chainable> ChainHandle<T>
|
|
|
|
|
}
|
|
|
|
|
|
|
|
|
|
/// Send this buffer to the chain
|
|
|
|
|
pub async fn write(&self, buf: Vec<T>) -> Result<(), SendError<Vec<T>>>
|
|
|
|
|
pub fn write(&self, buf: Vec<T>) -> impl futures::Future<Output = Result<(), SendError<Vec<T>>>> + 'static
|
|
|
|
|
{
|
|
|
|
|
let mut write = self.0.input.clone();
|
|
|
|
|
async move {
|
|
|
|
|
write.send(buf).await
|
|
|
|
|
}
|
|
|
|
|
}
|
|
|
|
|
|
|
|
|
|
/// Send this stream buffer to the chain
|
|
|
|
|
pub fn write_stream<'a, I: Stream<Item=T>>(&self, buf: I) -> impl futures::Future<Output = Result<(), SendError<Vec<T>>>> + 'a
|
|
|
|
|
where I: 'a
|
|
|
|
|
{
|
|
|
|
|
let mut write = self.0.input.clone();
|
|
|
|
|
async move {
|
|
|
|
|
write.send(buf.collect().await).await
|
|
|
|
|
}
|
|
|
|
|
}
|
|
|
|
|
|
|
|
|
|
/// Send this buffer to the chain
|
|
|
|
|
pub async fn write_in_place(&self, buf: Vec<T>) -> Result<(), SendError<Vec<T>>>
|
|
|
|
|
{
|
|
|
|
|
self.0.input.clone().send(buf).await
|
|
|
|
|
}
|
|
|
|
|
|
|
|
|
|
/// A referencer for the notifier
|
|
|
|
|
pub fn notify_when(&self) -> &Arc<Notify>
|
|
|
|
|
{
|
|
|
|
|
&self.0.notify_write
|
|
|
|
|
}
|
|
|
|
|
|
|
|
|
|
/// Force the pending buffers to be written to the chain now
|
|
|
|
|
pub fn push_now(&self)
|
|
|
|
|
{
|
|
|
|
|
self.0.push_now.notify();
|
|
|
|
|
}
|
|
|
|
|
|
|
|
|
|
/// Hang the worker thread, preventing it from taking any more inputs and also flushing it.
|
|
|
|
|
///
|
|
|
|
|
/// # Panics
|
|
|
|
|
/// If there was no worker thread.
|
|
|
|
|
pub fn hang(&self)
|
|
|
|
|
{
|
|
|
|
|
trace!("Communicating hang request");
|
|
|
|
|
self.0.shutdown.broadcast(true).expect("Failed to communicate hang");
|
|
|
|
|
}
|
|
|
|
|
}
|
|
|
|
|
|
|
|
|
|
impl ChainHandle<String>
|
|
|
|
@ -157,13 +209,13 @@ impl ChainHandle<String>
|
|
|
|
|
pub async fn host(from: ChainHandle<String>)
|
|
|
|
|
{
|
|
|
|
|
let opt = from.0.opt.clone();
|
|
|
|
|
let data = from.0.host.unwrap().await;
|
|
|
|
|
let mut data = from.0.host.unwrap().await;
|
|
|
|
|
|
|
|
|
|
let (mut tx, child) = {
|
|
|
|
|
let (mut tx, mut child) = {
|
|
|
|
|
// The `real` input channel.
|
|
|
|
|
let from = from.clone();
|
|
|
|
|
let opt = opt.clone();
|
|
|
|
|
let (tx, rx) = mpsc::channel::<Vec<Vec<_>>>(opt.backlog);
|
|
|
|
|
let (tx, rx) = mpsc::channel::<Vec<Vec<_>>>(opt.internal_backlog);
|
|
|
|
|
(tx, tokio::spawn(async move {
|
|
|
|
|
let mut rx = if let Some(thr) = opt.throttle {
|
|
|
|
|
time::throttle(thr, rx).boxed()
|
|
|
|
@ -172,13 +224,18 @@ pub async fn host(from: ChainHandle<String>)
|
|
|
|
|
};
|
|
|
|
|
trace!("child: Begin waiting on parent");
|
|
|
|
|
while let Some(item) = rx.next().await {
|
|
|
|
|
let mut lock = from.0.chain.write().await;
|
|
|
|
|
for item in item.into_iter()
|
|
|
|
|
{
|
|
|
|
|
use std::ops::DerefMut;
|
|
|
|
|
for item in item.into_iter() {
|
|
|
|
|
feed::feed(lock.deref_mut(), item, &from.0.opt.bounds);
|
|
|
|
|
if item.len() > 0 {
|
|
|
|
|
info!("Write lock acq");
|
|
|
|
|
let mut lock = from.0.chain.write().await;
|
|
|
|
|
for item in item.into_iter()
|
|
|
|
|
{
|
|
|
|
|
use std::ops::DerefMut;
|
|
|
|
|
for item in item.into_iter() {
|
|
|
|
|
feed::feed(lock.deref_mut(), item, &from.0.opt.bounds);
|
|
|
|
|
}
|
|
|
|
|
}
|
|
|
|
|
trace!("Signalling write");
|
|
|
|
|
from.0.notify_write.notify();
|
|
|
|
|
}
|
|
|
|
|
}
|
|
|
|
|
trace!("child: exiting");
|
|
|
|
@ -187,44 +244,94 @@ pub async fn host(from: ChainHandle<String>)
|
|
|
|
|
|
|
|
|
|
trace!("Begin polling on child");
|
|
|
|
|
tokio::select!{
|
|
|
|
|
v = child => {
|
|
|
|
|
v = &mut child => {
|
|
|
|
|
match v {
|
|
|
|
|
#[cold] Ok(_) => {warn!("Child exited before we have? This should probably never happen.")},//Should never happen.
|
|
|
|
|
Err(e) => {error!("Child exited abnormally. Aborting: {}", e)}, //Child panic or cancel.
|
|
|
|
|
}
|
|
|
|
|
},
|
|
|
|
|
_ = async move {
|
|
|
|
|
let mut rx = data.input.chunk(opt.capacity); //we don't even need this tbh
|
|
|
|
|
let mut rx = data.input.chunk(opt.capacity); //we don't even need this tbh, oh well.
|
|
|
|
|
|
|
|
|
|
while Arc::strong_count(&from.0) > 2 {
|
|
|
|
|
tokio::select!{
|
|
|
|
|
_ = time::delay_for(opt.timeout) => {
|
|
|
|
|
rx.push_now();
|
|
|
|
|
if !data.shutdown.recv().await.unwrap_or(true) { //first shutdown we get for free
|
|
|
|
|
while Arc::strong_count(&from.0) > 2 {
|
|
|
|
|
if *data.shutdown.borrow() {
|
|
|
|
|
break;
|
|
|
|
|
}
|
|
|
|
|
Some(buffer) = rx.next() => {
|
|
|
|
|
if let Err(err) = tx.send(buffer).await {
|
|
|
|
|
// Receive closed?
|
|
|
|
|
//
|
|
|
|
|
// This probably shouldn't happen, as we `select!` for it up there and child never calls `close()` on `rx`.
|
|
|
|
|
// In any case, it means we should abort.
|
|
|
|
|
error!("Failed to send buffer: {}", err);
|
|
|
|
|
|
|
|
|
|
tokio::select!{
|
|
|
|
|
Some(true) = data.shutdown.recv() => {
|
|
|
|
|
debug!("Got shutdown (hang) request. Sending now then breaking");
|
|
|
|
|
|
|
|
|
|
let mut rest = {
|
|
|
|
|
let irx = rx.get_mut();
|
|
|
|
|
irx.close(); //accept no more inputs
|
|
|
|
|
let mut output = Vec::with_capacity(opt.capacity);
|
|
|
|
|
while let Ok(item) = irx.try_recv() {
|
|
|
|
|
output.push(item);
|
|
|
|
|
}
|
|
|
|
|
output
|
|
|
|
|
};
|
|
|
|
|
rest.extend(rx.take_now());
|
|
|
|
|
if rest.len() > 0 {
|
|
|
|
|
if let Err(err) = tx.send(rest).await {
|
|
|
|
|
error!("Failed to force send buffer, exiting now: {}", err);
|
|
|
|
|
}
|
|
|
|
|
}
|
|
|
|
|
break;
|
|
|
|
|
}
|
|
|
|
|
_ = time::delay_for(opt.timeout) => {
|
|
|
|
|
trace!("Setting push now");
|
|
|
|
|
rx.push_now();
|
|
|
|
|
}
|
|
|
|
|
_ = from.0.push_now.notified() => {
|
|
|
|
|
debug!("Got force push signal");
|
|
|
|
|
let take =rx.take_now();
|
|
|
|
|
rx.push_now();
|
|
|
|
|
if take.len() > 0 {
|
|
|
|
|
if let Err(err) = tx.send(take).await {
|
|
|
|
|
error!("Failed to force send buffer: {}", err);
|
|
|
|
|
break;
|
|
|
|
|
}
|
|
|
|
|
}
|
|
|
|
|
}
|
|
|
|
|
Some(buffer) = rx.next() => {
|
|
|
|
|
debug!("Sending {} (cap {})", buffer.len(), buffer.capacity());
|
|
|
|
|
if let Err(err) = tx.send(buffer).await {
|
|
|
|
|
// Receive closed?
|
|
|
|
|
//
|
|
|
|
|
// This probably shouldn't happen, as we `select!` for it up there and child never calls `close()` on `rx`.
|
|
|
|
|
// In any case, it means we should abort.
|
|
|
|
|
#[cold] error!("Failed to send buffer: {}", err);
|
|
|
|
|
break;
|
|
|
|
|
}
|
|
|
|
|
}
|
|
|
|
|
}
|
|
|
|
|
}
|
|
|
|
|
}
|
|
|
|
|
let last = rx.into_buffer();
|
|
|
|
|
if last.len() > 0 {
|
|
|
|
|
if let Err(err) = tx.send(last).await {
|
|
|
|
|
error!("Failed to force send last part of buffer: {}", err);
|
|
|
|
|
} else {
|
|
|
|
|
trace!("Sent rest of buffer");
|
|
|
|
|
}
|
|
|
|
|
}
|
|
|
|
|
} => {
|
|
|
|
|
// Normal exit
|
|
|
|
|
trace!("Normal exit")
|
|
|
|
|
},
|
|
|
|
|
}
|
|
|
|
|
trace!("Waiting on child");
|
|
|
|
|
// No more handles except child, no more possible inputs.
|
|
|
|
|
child.await.expect("Child panic");
|
|
|
|
|
trace!("Returning");
|
|
|
|
|
}
|
|
|
|
|
|
|
|
|
|
/// Spawn a new chain handler for this chain.
|
|
|
|
|
pub fn spawn(from: chain::Chain<String>, opt: Settings) -> (JoinHandle<()>, ChainHandle<String>)
|
|
|
|
|
{
|
|
|
|
|
debug!("Spawning with opt: {:?}", opt);
|
|
|
|
|
let handle = ChainHandle::with_settings(from, opt);
|
|
|
|
|
(tokio::spawn(host(handle.clone())), handle)
|
|
|
|
|
}
|
|
|
|
|