You can not select more than 25 topics
Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
393 lines
9.4 KiB
393 lines
9.4 KiB
//! Chain handler.
|
|
use super::*;
|
|
use std::{
|
|
marker::Send,
|
|
sync::Weak,
|
|
num::NonZeroUsize,
|
|
task::{Poll, Context,},
|
|
pin::Pin,
|
|
};
|
|
use tokio::{
|
|
sync::{
|
|
RwLock,
|
|
RwLockReadGuard,
|
|
mpsc::{
|
|
self,
|
|
error::SendError,
|
|
},
|
|
watch,
|
|
Notify,
|
|
},
|
|
task::JoinHandle,
|
|
time::{
|
|
self,
|
|
Duration,
|
|
},
|
|
};
|
|
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>,
|
|
pub bounds: range::DynRange<usize>,
|
|
}
|
|
|
|
impl Settings
|
|
{
|
|
/// Should we keep this string.
|
|
#[inline] fn matches(&self, _s: &str) -> bool
|
|
{
|
|
true
|
|
}
|
|
}
|
|
|
|
impl Default for Settings
|
|
{
|
|
#[inline]
|
|
fn default() -> Self
|
|
{
|
|
Self {
|
|
backlog: 32,
|
|
internal_backlog: 8,
|
|
capacity: 4,
|
|
timeout: Duration::from_secs(5),
|
|
throttle: Some(Duration::from_millis(200)),
|
|
bounds: feed::DEFAULT_FEED_BOUNDS.into(),
|
|
}
|
|
}
|
|
}
|
|
|
|
|
|
#[derive(Debug)]
|
|
struct HostInner<T>
|
|
{
|
|
input: mpsc::Receiver<Vec<T>>,
|
|
shutdown: watch::Receiver<bool>,
|
|
}
|
|
|
|
#[derive(Debug)]
|
|
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>>,
|
|
}
|
|
|
|
#[derive(Clone, Debug)]
|
|
pub struct ChainHandle<T: Send + chain::Chainable>(Arc<Box<Handle<T>>>);
|
|
|
|
impl<T: Send+ chain::Chainable + 'static> ChainHandle<T>
|
|
{
|
|
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,
|
|
})
|
|
})))
|
|
}
|
|
|
|
/// Acquire the chain read lock
|
|
async fn chain(&self) -> RwLockReadGuard<'_, chain::Chain<T>>
|
|
{
|
|
self.0.chain.read().await
|
|
}
|
|
|
|
/// A reference to the chain
|
|
pub fn chain_ref(&self) -> &RwLock<chain::Chain<T>>
|
|
{
|
|
&self.0.chain
|
|
}
|
|
|
|
/// Create a stream that reads generated values forever.
|
|
pub fn read(&self) -> ChainStream<T>
|
|
{
|
|
ChainStream{
|
|
chain: Arc::downgrade(&self.0),
|
|
buffer: Vec::with_capacity(self.0.opt.backlog),
|
|
}
|
|
}
|
|
|
|
/// Send this buffer to the chain
|
|
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>
|
|
{
|
|
#[deprecated = "use read() pls"]
|
|
pub async fn generate_body(&self, state: &state::State, num: Option<usize>, mut output: mpsc::Sender<String>) -> Result<(), SendError<String>>
|
|
{
|
|
let chain = self.chain().await;
|
|
if !chain.is_empty() {
|
|
let filter = state.outbound_filter();
|
|
match num {
|
|
Some(num) if num < state.config().max_gen_size => {
|
|
//This could DoS writes, potentially.
|
|
for string in chain.str_iter_for(num) {
|
|
output.send(filter.filter_owned(string)).await?;
|
|
}
|
|
},
|
|
_ => output.send(filter.filter_owned(chain.generate_str())).await?,
|
|
}
|
|
}
|
|
Ok(())
|
|
}
|
|
}
|
|
|
|
/// Host this handle on the current task.
|
|
///
|
|
/// # Panics
|
|
/// If `from` has already been hosted.
|
|
pub async fn host(from: ChainHandle<String>)
|
|
{
|
|
let opt = from.0.opt.clone();
|
|
let mut data = from.0.host.unwrap().await;
|
|
|
|
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.internal_backlog);
|
|
(tx, tokio::spawn(async move {
|
|
let mut rx = if let Some(thr) = opt.throttle {
|
|
time::throttle(thr, rx).boxed()
|
|
} else {
|
|
rx.boxed()
|
|
};
|
|
trace!("child: Begin waiting on parent");
|
|
while let Some(item) = rx.next().await {
|
|
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");
|
|
}))
|
|
};
|
|
|
|
trace!("Begin polling on child");
|
|
tokio::select!{
|
|
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, oh well.
|
|
|
|
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;
|
|
}
|
|
|
|
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)
|
|
}
|
|
|
|
#[derive(Debug)]
|
|
pub struct ChainStream<T: Send + chain::Chainable>
|
|
{
|
|
chain: Weak<Box<Handle<T>>>,
|
|
buffer: Vec<T>,
|
|
}
|
|
|
|
impl ChainStream<String>
|
|
{
|
|
async fn try_pull(&mut self, n: usize) -> Option<NonZeroUsize>
|
|
{
|
|
if n == 0 {
|
|
return None;
|
|
}
|
|
if let Some(read) = self.chain.upgrade() {
|
|
let chain = read.chain.read().await;
|
|
if chain.is_empty() {
|
|
return None;
|
|
}
|
|
|
|
let n = if n == 1 {
|
|
self.buffer.push(chain.generate_str());
|
|
1
|
|
} else {
|
|
self.buffer.extend(chain.str_iter_for(n));
|
|
n //for now
|
|
};
|
|
Some(unsafe{NonZeroUsize::new_unchecked(n)})
|
|
} else {
|
|
None
|
|
}
|
|
}
|
|
}
|
|
|
|
impl Stream for ChainStream<String>
|
|
{
|
|
type Item = String;
|
|
fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
|
|
use futures::Future;
|
|
let this = self.get_mut();
|
|
|
|
if this.buffer.len() == 0 {
|
|
let pull = this.try_pull(this.buffer.capacity());
|
|
tokio::pin!(pull);
|
|
match pull.poll(cx) {
|
|
Poll::Ready(Some(_)) => {},
|
|
Poll::Pending => return Poll::Pending,
|
|
_ => return Poll::Ready(None),
|
|
};
|
|
}
|
|
debug_assert!(this.buffer.len()>0);
|
|
Poll::Ready(Some(this.buffer.remove(0)))
|
|
}
|
|
}
|