Spaces:
Running
Running
File size: 3,737 Bytes
287a0bc |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 |
use std::fmt::Debug;
use std::sync::Arc;
use futures::Stream;
use futures::StreamExt;
use tokio::runtime::Builder;
use tokio::{pin, select};
use super::ComponentRuntime;
// use super::executor::StreamComponentExecutor;
use super::sender::{self, Sender, Wrapper};
use super::{executor, ComponentContext};
use super::{executor::ComponentExecutor, Component, ComponentHandle, Handler, StreamHandler};
use std::sync::Mutex;
#[derive(Clone)]
pub(crate) struct System {
inner: Arc<Mutex<Inner>>,
}
struct Inner {}
impl System {
pub(crate) fn new() -> System {
System {
inner: Arc::new(Mutex::new(Inner {})),
}
}
pub(crate) fn start_component<C>(&mut self, mut component: C) -> ComponentHandle<C>
where
C: Component + Send + 'static,
{
let (tx, rx) = tokio::sync::mpsc::channel(component.queue_size());
let sender = Sender::new(tx);
let cancel_token = tokio_util::sync::CancellationToken::new();
let _ = component.on_start(&ComponentContext {
system: self.clone(),
sender: sender.clone(),
cancellation_token: cancel_token.clone(),
});
let mut executor = ComponentExecutor::new(
sender.clone(),
cancel_token.clone(),
component,
self.clone(),
);
match C::runtime() {
ComponentRuntime::Global => {
let join_handle = tokio::spawn(async move { executor.run(rx).await });
return ComponentHandle::new(cancel_token, Some(join_handle), sender);
}
ComponentRuntime::Dedicated => {
println!("Spawning on dedicated thread");
// Spawn on a dedicated thread
let mut rt = Builder::new_current_thread().enable_all().build().unwrap();
let join_handle = std::thread::spawn(move || {
rt.block_on(async move { executor.run(rx).await });
});
// TODO: Implement Join for dedicated threads
return ComponentHandle::new(cancel_token, None, sender);
}
}
}
pub(super) fn register_stream<C, S, M>(&self, stream: S, ctx: &ComponentContext<C>)
where
C: StreamHandler<M> + Handler<M>,
M: Send + Debug + 'static,
S: Stream + Send + Stream<Item = M> + 'static,
{
let ctx = ComponentContext {
system: self.clone(),
sender: ctx.sender.clone(),
cancellation_token: ctx.cancellation_token.clone(),
};
tokio::spawn(async move { stream_loop(stream, &ctx).await });
}
}
async fn stream_loop<C, S, M>(stream: S, ctx: &ComponentContext<C>)
where
C: StreamHandler<M> + Handler<M>,
M: Send + Debug + 'static,
S: Stream + Send + Stream<Item = M> + 'static,
{
pin!(stream);
loop {
select! {
_ = ctx.cancellation_token.cancelled() => {
break;
}
message = stream.next() => {
match message {
Some(message) => {
let res = ctx.sender.send(message).await;
match res {
Ok(_) => {}
Err(e) => {
println!("Failed to send message: {:?}", e);
// TODO: switch to logging
// Terminate the stream
break;
}
}
},
None => {
break;
}
}
}
}
}
}
|