forked from Orchid/orchid
59 lines
2.0 KiB
Rust
59 lines
2.0 KiB
Rust
use crate::entrypoint::ExtensionData;
|
|
|
|
#[cfg(feature = "tokio")]
|
|
pub async fn tokio_main(data: ExtensionData) {
|
|
use std::io::Write;
|
|
use std::mem;
|
|
use std::pin::{Pin, pin};
|
|
use std::rc::Rc;
|
|
|
|
use async_once_cell::OnceCell;
|
|
use futures::StreamExt;
|
|
use futures::future::LocalBoxFuture;
|
|
use futures::lock::Mutex;
|
|
use futures::stream::FuturesUnordered;
|
|
use orchid_api_traits::{Decode, Encode};
|
|
use orchid_base::msg::{recv_msg, send_msg};
|
|
use tokio::io;
|
|
use tokio::io::Stdout;
|
|
use tokio::task::{LocalSet, spawn_local};
|
|
use tokio_util::compat::{Compat, TokioAsyncReadCompatExt, TokioAsyncWriteCompatExt};
|
|
|
|
use crate::api;
|
|
use crate::entrypoint::extension_init;
|
|
|
|
let local_set = LocalSet::new();
|
|
local_set.spawn_local(async {
|
|
let host_header = api::HostHeader::decode(Pin::new(&mut tokio::io::stdin().compat())).await;
|
|
let init =
|
|
Rc::new(extension_init(data, host_header, Rc::new(|fut| mem::drop(spawn_local(fut)))));
|
|
let mut buf = Vec::new();
|
|
init.header.encode(Pin::new(&mut buf)).await;
|
|
std::io::stdout().write_all(&buf).unwrap();
|
|
std::io::stdout().flush().unwrap();
|
|
// These are concurrent processes that never exit, so if the FuturesUnordered
|
|
// produces any result the extension should exit
|
|
let mut io = FuturesUnordered::<LocalBoxFuture<()>>::new();
|
|
io.push(Box::pin(async {
|
|
loop {
|
|
match recv_msg(pin!(io::stdin().compat())).await {
|
|
Ok(msg) => init.send(&msg[..]).await,
|
|
Err(e) if e.kind() == io::ErrorKind::BrokenPipe => break,
|
|
Err(e) if e.kind() == io::ErrorKind::UnexpectedEof => break,
|
|
Err(e) => panic!("{e}"),
|
|
}
|
|
}
|
|
}));
|
|
io.push(Box::pin(async {
|
|
while let Some(msg) = init.recv().await {
|
|
static STDOUT: OnceCell<Mutex<Compat<Stdout>>> = OnceCell::new();
|
|
let stdout_lk = STDOUT.get_or_init(async { Mutex::new(io::stdout().compat_write()) }).await;
|
|
let mut stdout_g = stdout_lk.lock().await;
|
|
send_msg(pin!(&mut *stdout_g), &msg[..]).await.expect("Parent pipe broken");
|
|
}
|
|
}));
|
|
io.next().await;
|
|
});
|
|
local_set.await;
|
|
}
|