mirror of
https://github.com/EasyTier/EasyTier.git
synced 2026-05-07 10:14:35 +00:00
Feat/web (PatchSet 1) (#436)
* move rpc-build out of easytier dir and make it a independant project * easytier core use launcher * fix flags not print on launch * allow launcher not fetch node info * abstract out peer rpc impl * fix arm gui ci. see https://github.com/actions/runner-images/pull/10807 * add easytier-web crate * fix manual_connector test case
This commit is contained in:
Sijie.Sun
@@ -1,8 +0,0 @@
|
||||
[package]
|
||||
name = "rpc_build"
|
||||
version = "0.1.0"
|
||||
edition = "2021"
|
||||
|
||||
[dependencies]
|
||||
heck = "0.5"
|
||||
prost-build = "0.13"
|
||||
@@ -1,383 +0,0 @@
|
||||
extern crate heck;
|
||||
extern crate prost_build;
|
||||
|
||||
use std::fmt;
|
||||
|
||||
const NAMESPACE: &str = "crate::proto::rpc_types";
|
||||
|
||||
/// The service generator to be used with `prost-build` to generate RPC implementations for
|
||||
/// `prost-simple-rpc`.
|
||||
///
|
||||
/// See the crate-level documentation for more info.
|
||||
#[allow(missing_copy_implementations)]
|
||||
#[derive(Clone, Debug)]
|
||||
pub struct ServiceGenerator {
|
||||
_private: (),
|
||||
}
|
||||
|
||||
impl ServiceGenerator {
|
||||
/// Create a new `ServiceGenerator` instance with the default options set.
|
||||
pub fn new() -> ServiceGenerator {
|
||||
ServiceGenerator { _private: () }
|
||||
}
|
||||
}
|
||||
|
||||
impl prost_build::ServiceGenerator for ServiceGenerator {
|
||||
fn generate(&mut self, service: prost_build::Service, mut buf: &mut String) {
|
||||
use std::fmt::Write;
|
||||
|
||||
let descriptor_name = format!("{}Descriptor", service.name);
|
||||
let server_name = format!("{}Server", service.name);
|
||||
let client_name = format!("{}Client", service.name);
|
||||
let method_descriptor_name = format!("{}MethodDescriptor", service.name);
|
||||
|
||||
let mut trait_methods = String::new();
|
||||
let mut enum_methods = String::new();
|
||||
let mut list_enum_methods = String::new();
|
||||
let mut client_methods = String::new();
|
||||
let mut client_own_methods = String::new();
|
||||
let mut match_name_methods = String::new();
|
||||
let mut match_proto_name_methods = String::new();
|
||||
let mut match_input_type_methods = String::new();
|
||||
let mut match_input_proto_type_methods = String::new();
|
||||
let mut match_output_type_methods = String::new();
|
||||
let mut match_output_proto_type_methods = String::new();
|
||||
let mut match_handle_methods = String::new();
|
||||
|
||||
let mut match_method_try_from = String::new();
|
||||
|
||||
for (idx, method) in service.methods.iter().enumerate() {
|
||||
assert!(
|
||||
!method.client_streaming,
|
||||
"Client streaming not yet supported for method {}",
|
||||
method.proto_name
|
||||
);
|
||||
assert!(
|
||||
!method.server_streaming,
|
||||
"Server streaming not yet supported for method {}",
|
||||
method.proto_name
|
||||
);
|
||||
|
||||
ServiceGenerator::write_comments(&mut trait_methods, 4, &method.comments).unwrap();
|
||||
writeln!(
|
||||
trait_methods,
|
||||
r#" async fn {name}(&self, ctrl: Self::Controller, input: {input_type}) -> {namespace}::error::Result<{output_type}>;"#,
|
||||
name = method.name,
|
||||
input_type = method.input_type,
|
||||
output_type = method.output_type,
|
||||
namespace = NAMESPACE,
|
||||
)
|
||||
.unwrap();
|
||||
|
||||
ServiceGenerator::write_comments(&mut enum_methods, 4, &method.comments).unwrap();
|
||||
writeln!(
|
||||
enum_methods,
|
||||
" {name} = {index},",
|
||||
name = method.proto_name,
|
||||
index = format!("{}", idx + 1)
|
||||
)
|
||||
.unwrap();
|
||||
|
||||
writeln!(
|
||||
match_method_try_from,
|
||||
" {index} => Ok({service_name}MethodDescriptor::{name}),",
|
||||
service_name = service.name,
|
||||
name = method.proto_name,
|
||||
index = format!("{}", idx + 1),
|
||||
)
|
||||
.unwrap();
|
||||
|
||||
writeln!(
|
||||
list_enum_methods,
|
||||
" {service_name}MethodDescriptor::{name},",
|
||||
service_name = service.name,
|
||||
name = method.proto_name
|
||||
)
|
||||
.unwrap();
|
||||
|
||||
writeln!(
|
||||
client_methods,
|
||||
r#" async fn {name}(&self, ctrl: H::Controller, input: {input_type}) -> {namespace}::error::Result<{output_type}> {{
|
||||
{client_name}::{name}_inner(self.0.clone(), ctrl, input).await
|
||||
}}"#,
|
||||
name = method.name,
|
||||
input_type = method.input_type,
|
||||
output_type = method.output_type,
|
||||
client_name = format!("{}Client", service.name),
|
||||
namespace = NAMESPACE,
|
||||
)
|
||||
.unwrap();
|
||||
|
||||
writeln!(
|
||||
client_own_methods,
|
||||
r#" async fn {name}_inner(handler: H, ctrl: H::Controller, input: {input_type}) -> {namespace}::error::Result<{output_type}> {{
|
||||
{namespace}::__rt::call_method(handler, ctrl, {method_descriptor_name}::{proto_name}, input).await
|
||||
}}"#,
|
||||
name = method.name,
|
||||
method_descriptor_name = method_descriptor_name,
|
||||
proto_name = method.proto_name,
|
||||
input_type = method.input_type,
|
||||
output_type = method.output_type,
|
||||
namespace = NAMESPACE,
|
||||
).unwrap();
|
||||
|
||||
let case = format!(
|
||||
" {service_name}MethodDescriptor::{proto_name} => ",
|
||||
service_name = service.name,
|
||||
proto_name = method.proto_name
|
||||
);
|
||||
|
||||
writeln!(match_name_methods, "{}{:?},", case, method.name).unwrap();
|
||||
writeln!(match_proto_name_methods, "{}{:?},", case, method.proto_name).unwrap();
|
||||
writeln!(
|
||||
match_input_type_methods,
|
||||
"{}::std::any::TypeId::of::<{}>(),",
|
||||
case, method.input_type
|
||||
)
|
||||
.unwrap();
|
||||
writeln!(
|
||||
match_input_proto_type_methods,
|
||||
"{}{:?},",
|
||||
case, method.input_proto_type
|
||||
)
|
||||
.unwrap();
|
||||
writeln!(
|
||||
match_output_type_methods,
|
||||
"{}::std::any::TypeId::of::<{}>(),",
|
||||
case, method.output_type
|
||||
)
|
||||
.unwrap();
|
||||
writeln!(
|
||||
match_output_proto_type_methods,
|
||||
"{}{:?},",
|
||||
case, method.output_proto_type
|
||||
)
|
||||
.unwrap();
|
||||
write!(
|
||||
match_handle_methods,
|
||||
r#"{} {{
|
||||
let decoded: {input_type} = {namespace}::__rt::decode(input)?;
|
||||
let ret = service.{name}(ctrl, decoded).await?;
|
||||
{namespace}::__rt::encode(ret)
|
||||
}}
|
||||
"#,
|
||||
case,
|
||||
input_type = method.input_type,
|
||||
name = method.name,
|
||||
namespace = NAMESPACE,
|
||||
)
|
||||
.unwrap();
|
||||
}
|
||||
|
||||
ServiceGenerator::write_comments(&mut buf, 0, &service.comments).unwrap();
|
||||
write!(
|
||||
buf,
|
||||
r#"
|
||||
#[async_trait::async_trait]
|
||||
#[auto_impl::auto_impl(&, Arc, Box)]
|
||||
pub trait {name} {{
|
||||
type Controller: {namespace}::controller::Controller;
|
||||
|
||||
{trait_methods}
|
||||
}}
|
||||
|
||||
/// A service descriptor for a `{name}`.
|
||||
#[derive(Clone, Debug, Eq, Ord, PartialEq, PartialOrd, Default)]
|
||||
pub struct {descriptor_name};
|
||||
|
||||
/// Methods available on a `{name}`.
|
||||
///
|
||||
/// This can be used as a key when routing requests for servers/clients of a `{name}`.
|
||||
#[derive(Clone, Copy, Debug, Eq, Ord, PartialEq, PartialOrd)]
|
||||
#[repr(u8)]
|
||||
pub enum {method_descriptor_name} {{
|
||||
{enum_methods}
|
||||
}}
|
||||
|
||||
impl std::convert::TryFrom<u8> for {method_descriptor_name} {{
|
||||
type Error = {namespace}::error::Error;
|
||||
fn try_from(value: u8) -> {namespace}::error::Result<Self> {{
|
||||
match value {{
|
||||
{match_method_try_from}
|
||||
_ => Err({namespace}::error::Error::InvalidMethodIndex(value, "{name}".to_string())),
|
||||
}}
|
||||
}}
|
||||
}}
|
||||
|
||||
/// A client for a `{name}`.
|
||||
///
|
||||
/// This implements the `{name}` trait by dispatching all method calls to the supplied `Handler`.
|
||||
#[derive(Clone, Debug)]
|
||||
pub struct {client_name}<H>(H) where H: {namespace}::handler::Handler;
|
||||
|
||||
impl<H> {client_name}<H> where H: {namespace}::handler::Handler<Descriptor = {descriptor_name}> {{
|
||||
/// Creates a new client instance that delegates all method calls to the supplied handler.
|
||||
pub fn new(handler: H) -> {client_name}<H> {{
|
||||
{client_name}(handler)
|
||||
}}
|
||||
}}
|
||||
|
||||
impl<H> {client_name}<H> where H: {namespace}::handler::Handler<Descriptor = {descriptor_name}> {{
|
||||
{client_own_methods}
|
||||
}}
|
||||
|
||||
#[async_trait::async_trait]
|
||||
impl<H> {name} for {client_name}<H> where H: {namespace}::handler::Handler<Descriptor = {descriptor_name}> {{
|
||||
type Controller = H::Controller;
|
||||
|
||||
{client_methods}
|
||||
}}
|
||||
|
||||
pub struct {client_name}Factory<C: {namespace}::controller::Controller>(std::marker::PhantomData<C>);
|
||||
|
||||
impl<C: {namespace}::controller::Controller> Clone for {client_name}Factory<C> {{
|
||||
fn clone(&self) -> Self {{
|
||||
Self(std::marker::PhantomData)
|
||||
}}
|
||||
}}
|
||||
|
||||
impl<C> {namespace}::__rt::RpcClientFactory for {client_name}Factory<C> where C: {namespace}::controller::Controller {{
|
||||
type Descriptor = {descriptor_name};
|
||||
type ClientImpl = Box<dyn {name}<Controller = C> + Send + 'static>;
|
||||
type Controller = C;
|
||||
|
||||
fn new(handler: impl {namespace}::handler::Handler<Descriptor = Self::Descriptor, Controller = Self::Controller>) -> Self::ClientImpl {{
|
||||
Box::new({client_name}::new(handler))
|
||||
}}
|
||||
}}
|
||||
|
||||
/// A server for a `{name}`.
|
||||
///
|
||||
/// This implements the `Server` trait by handling requests and dispatch them to methods on the
|
||||
/// supplied `{name}`.
|
||||
#[derive(Clone, Debug)]
|
||||
pub struct {server_name}<A>(A) where A: {name} + Clone + Send + 'static;
|
||||
|
||||
impl<A> {server_name}<A> where A: {name} + Clone + Send + 'static {{
|
||||
/// Creates a new server instance that dispatches all calls to the supplied service.
|
||||
pub fn new(service: A) -> {server_name}<A> {{
|
||||
{server_name}(service)
|
||||
}}
|
||||
|
||||
async fn call_inner(
|
||||
service: A,
|
||||
method: {method_descriptor_name},
|
||||
ctrl: A::Controller,
|
||||
input: ::bytes::Bytes)
|
||||
-> {namespace}::error::Result<::bytes::Bytes> {{
|
||||
match method {{
|
||||
{match_handle_methods}
|
||||
}}
|
||||
}}
|
||||
}}
|
||||
|
||||
impl {namespace}::descriptor::ServiceDescriptor for {descriptor_name} {{
|
||||
type Method = {method_descriptor_name};
|
||||
fn name(&self) -> &'static str {{ {name:?} }}
|
||||
fn proto_name(&self) -> &'static str {{ {proto_name:?} }}
|
||||
fn package(&self) -> &'static str {{ {package:?} }}
|
||||
fn methods(&self) -> &'static [Self::Method] {{
|
||||
&[ {list_enum_methods} ]
|
||||
}}
|
||||
}}
|
||||
|
||||
#[async_trait::async_trait]
|
||||
impl<A> {namespace}::handler::Handler for {server_name}<A>
|
||||
where
|
||||
A: {name} + Clone + Send + Sync + 'static {{
|
||||
type Descriptor = {descriptor_name};
|
||||
type Controller = A::Controller;
|
||||
|
||||
async fn call(
|
||||
&self,
|
||||
ctrl: A::Controller,
|
||||
method: {method_descriptor_name},
|
||||
input: ::bytes::Bytes)
|
||||
-> {namespace}::error::Result<::bytes::Bytes> {{
|
||||
{server_name}::call_inner(self.0.clone(), method, ctrl, input).await
|
||||
}}
|
||||
}}
|
||||
|
||||
impl {namespace}::descriptor::MethodDescriptor for {method_descriptor_name} {{
|
||||
fn name(&self) -> &'static str {{
|
||||
match *self {{
|
||||
{match_name_methods}
|
||||
}}
|
||||
}}
|
||||
|
||||
fn proto_name(&self) -> &'static str {{
|
||||
match *self {{
|
||||
{match_proto_name_methods}
|
||||
}}
|
||||
}}
|
||||
|
||||
fn input_type(&self) -> ::std::any::TypeId {{
|
||||
match *self {{
|
||||
{match_input_type_methods}
|
||||
}}
|
||||
}}
|
||||
|
||||
fn input_proto_type(&self) -> &'static str {{
|
||||
match *self {{
|
||||
{match_input_proto_type_methods}
|
||||
}}
|
||||
}}
|
||||
|
||||
fn output_type(&self) -> ::std::any::TypeId {{
|
||||
match *self {{
|
||||
{match_output_type_methods}
|
||||
}}
|
||||
}}
|
||||
|
||||
fn output_proto_type(&self) -> &'static str {{
|
||||
match *self {{
|
||||
{match_output_proto_type_methods}
|
||||
}}
|
||||
}}
|
||||
|
||||
fn index(&self) -> u8 {{
|
||||
*self as u8
|
||||
}}
|
||||
}}
|
||||
"#,
|
||||
name = service.name,
|
||||
descriptor_name = descriptor_name,
|
||||
server_name = server_name,
|
||||
client_name = client_name,
|
||||
method_descriptor_name = method_descriptor_name,
|
||||
proto_name = service.proto_name,
|
||||
package = service.package,
|
||||
trait_methods = trait_methods,
|
||||
enum_methods = enum_methods,
|
||||
list_enum_methods = list_enum_methods,
|
||||
client_own_methods = client_own_methods,
|
||||
client_methods = client_methods,
|
||||
match_name_methods = match_name_methods,
|
||||
match_proto_name_methods = match_proto_name_methods,
|
||||
match_input_type_methods = match_input_type_methods,
|
||||
match_input_proto_type_methods = match_input_proto_type_methods,
|
||||
match_output_type_methods = match_output_type_methods,
|
||||
match_output_proto_type_methods = match_output_proto_type_methods,
|
||||
match_handle_methods = match_handle_methods,
|
||||
namespace = NAMESPACE,
|
||||
).unwrap();
|
||||
}
|
||||
}
|
||||
|
||||
impl ServiceGenerator {
|
||||
fn write_comments<W>(
|
||||
mut write: W,
|
||||
indent: usize,
|
||||
comments: &prost_build::Comments,
|
||||
) -> fmt::Result
|
||||
where
|
||||
W: fmt::Write,
|
||||
{
|
||||
for comment in &comments.leading {
|
||||
for line in comment.lines().filter(|s| !s.is_empty()) {
|
||||
writeln!(write, "{}///{}", " ".repeat(indent), line)?;
|
||||
}
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,164 @@
|
||||
use std::sync::{Arc, Mutex};
|
||||
|
||||
use futures::{SinkExt as _, StreamExt};
|
||||
use tokio::{task::JoinSet, time::timeout};
|
||||
|
||||
use crate::{
|
||||
proto::rpc_types::error::Error,
|
||||
tunnel::{packet_def::PacketType, ring::create_ring_tunnel_pair, Tunnel},
|
||||
};
|
||||
|
||||
use super::{client::Client, server::Server};
|
||||
|
||||
pub struct BidirectRpcManager {
|
||||
rpc_client: Client,
|
||||
rpc_server: Server,
|
||||
|
||||
rx_timeout: Option<std::time::Duration>,
|
||||
error: Arc<Mutex<Option<Error>>>,
|
||||
tunnel: Mutex<Option<Box<dyn Tunnel>>>,
|
||||
|
||||
tasks: Mutex<Option<JoinSet<()>>>,
|
||||
}
|
||||
|
||||
impl BidirectRpcManager {
|
||||
pub fn new() -> Self {
|
||||
Self {
|
||||
rpc_client: Client::new(),
|
||||
rpc_server: Server::new(),
|
||||
|
||||
rx_timeout: None,
|
||||
error: Arc::new(Mutex::new(None)),
|
||||
tunnel: Mutex::new(None),
|
||||
|
||||
tasks: Mutex::new(None),
|
||||
}
|
||||
}
|
||||
|
||||
pub fn set_rx_timeout(mut self, timeout: Option<std::time::Duration>) -> Self {
|
||||
self.rx_timeout = timeout;
|
||||
self
|
||||
}
|
||||
|
||||
pub fn run_and_create_tunnel(&self) -> Box<dyn Tunnel> {
|
||||
let (ret, inner) = create_ring_tunnel_pair();
|
||||
self.run_with_tunnel(inner);
|
||||
ret
|
||||
}
|
||||
|
||||
pub fn run_with_tunnel(&self, inner: Box<dyn Tunnel>) {
|
||||
let mut tasks = JoinSet::new();
|
||||
self.rpc_client.run();
|
||||
self.rpc_server.run();
|
||||
|
||||
let (server_tx, mut server_rx) = (
|
||||
self.rpc_server.get_transport_sink(),
|
||||
self.rpc_server.get_transport_stream(),
|
||||
);
|
||||
let (client_tx, mut client_rx) = (
|
||||
self.rpc_client.get_transport_sink(),
|
||||
self.rpc_client.get_transport_stream(),
|
||||
);
|
||||
|
||||
let (mut inner_rx, mut inner_tx) = inner.split();
|
||||
self.tunnel.lock().unwrap().replace(inner);
|
||||
|
||||
let e_clone = self.error.clone();
|
||||
tasks.spawn(async move {
|
||||
loop {
|
||||
let packet = tokio::select! {
|
||||
Some(Ok(packet)) = server_rx.next() => {
|
||||
tracing::trace!(?packet, "recv rpc packet from server");
|
||||
packet
|
||||
}
|
||||
Some(Ok(packet)) = client_rx.next() => {
|
||||
tracing::trace!(?packet, "recv rpc packet from client");
|
||||
packet
|
||||
}
|
||||
else => {
|
||||
tracing::warn!("rpc transport read aborted, exiting");
|
||||
break;
|
||||
}
|
||||
};
|
||||
|
||||
if let Err(e) = inner_tx.send(packet).await {
|
||||
tracing::error!(error = ?e, "send to peer failed");
|
||||
e_clone.lock().unwrap().replace(Error::from(e));
|
||||
}
|
||||
}
|
||||
});
|
||||
|
||||
let recv_timeout = self.rx_timeout;
|
||||
let e_clone = self.error.clone();
|
||||
tasks.spawn(async move {
|
||||
loop {
|
||||
let ret = if let Some(recv_timeout) = recv_timeout {
|
||||
match timeout(recv_timeout, inner_rx.next()).await {
|
||||
Ok(ret) => ret,
|
||||
Err(e) => {
|
||||
e_clone.lock().unwrap().replace(e.into());
|
||||
break;
|
||||
}
|
||||
}
|
||||
} else {
|
||||
inner_rx.next().await
|
||||
};
|
||||
|
||||
let o = match ret {
|
||||
Some(Ok(o)) => o,
|
||||
Some(Err(e)) => {
|
||||
tracing::error!(error = ?e, "recv from peer failed");
|
||||
e_clone.lock().unwrap().replace(Error::from(e));
|
||||
break;
|
||||
}
|
||||
None => {
|
||||
tracing::warn!("peer rpc transport read aborted, exiting");
|
||||
e_clone.lock().unwrap().replace(Error::Shutdown);
|
||||
break;
|
||||
}
|
||||
};
|
||||
|
||||
if o.peer_manager_header().unwrap().packet_type == PacketType::RpcReq as u8 {
|
||||
server_tx.send(o).await.unwrap();
|
||||
continue;
|
||||
} else if o.peer_manager_header().unwrap().packet_type == PacketType::RpcResp as u8
|
||||
{
|
||||
client_tx.send(o).await.unwrap();
|
||||
continue;
|
||||
}
|
||||
}
|
||||
});
|
||||
|
||||
self.tasks.lock().unwrap().replace(tasks);
|
||||
}
|
||||
|
||||
pub fn rpc_client(&self) -> &Client {
|
||||
&self.rpc_client
|
||||
}
|
||||
|
||||
pub fn rpc_server(&self) -> &Server {
|
||||
&self.rpc_server
|
||||
}
|
||||
|
||||
pub async fn stop(&self) {
|
||||
let Some(mut tasks) = self.tasks.lock().unwrap().take() else {
|
||||
return;
|
||||
};
|
||||
tasks.abort_all();
|
||||
while let Some(_) = tasks.join_next().await {}
|
||||
}
|
||||
|
||||
pub fn take_error(&self) -> Option<Error> {
|
||||
self.error.lock().unwrap().take()
|
||||
}
|
||||
|
||||
pub async fn wait(&self) {
|
||||
let Some(mut tasks) = self.tasks.lock().unwrap().take() else {
|
||||
return;
|
||||
};
|
||||
while let Some(_) = tasks.join_next().await {
|
||||
// when any task is done, abort all tasks
|
||||
tasks.abort_all();
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -2,6 +2,7 @@ use crate::tunnel::{mpsc::MpscTunnel, Tunnel};
|
||||
|
||||
pub type RpcController = super::rpc_types::controller::BaseController;
|
||||
|
||||
pub mod bidirect;
|
||||
pub mod client;
|
||||
pub mod packet;
|
||||
pub mod server;
|
||||
|
||||
@@ -59,6 +59,14 @@ impl ServiceRegistry {
|
||||
}
|
||||
}
|
||||
|
||||
pub fn replace_registry(&self, registry: &ServiceRegistry) {
|
||||
self.table.clear();
|
||||
for item in registry.table.iter() {
|
||||
let (k, v) = item.pair();
|
||||
self.table.insert(k.clone(), v.clone());
|
||||
}
|
||||
}
|
||||
|
||||
pub fn register<H: Handler<Controller = RpcController>>(&self, h: H, domain_name: &str) {
|
||||
let desc = h.service_descriptor();
|
||||
let key = ServiceKey {
|
||||
|
||||
@@ -4,66 +4,18 @@ use std::{
|
||||
};
|
||||
|
||||
use anyhow::Context as _;
|
||||
use futures::{SinkExt as _, StreamExt};
|
||||
use tokio::task::JoinSet;
|
||||
|
||||
use crate::{
|
||||
common::join_joinset_background,
|
||||
proto::rpc_types::{__rt::RpcClientFactory, error::Error},
|
||||
proto::{
|
||||
rpc_impl::bidirect::BidirectRpcManager,
|
||||
rpc_types::{__rt::RpcClientFactory, error::Error},
|
||||
},
|
||||
tunnel::{Tunnel, TunnelConnector, TunnelListener},
|
||||
};
|
||||
|
||||
use super::{client::Client, server::Server, service_registry::ServiceRegistry};
|
||||
|
||||
struct StandAloneServerOneTunnel {
|
||||
tunnel: Box<dyn Tunnel>,
|
||||
rpc_server: Server,
|
||||
}
|
||||
|
||||
impl StandAloneServerOneTunnel {
|
||||
pub fn new(tunnel: Box<dyn Tunnel>, registry: Arc<ServiceRegistry>) -> Self {
|
||||
let rpc_server = Server::new_with_registry(registry);
|
||||
StandAloneServerOneTunnel { tunnel, rpc_server }
|
||||
}
|
||||
|
||||
pub async fn run(self) {
|
||||
use tokio_stream::StreamExt as _;
|
||||
|
||||
let (tunnel_rx, tunnel_tx) = self.tunnel.split();
|
||||
let (rpc_rx, rpc_tx) = (
|
||||
self.rpc_server.get_transport_stream(),
|
||||
self.rpc_server.get_transport_sink(),
|
||||
);
|
||||
|
||||
let mut tasks = JoinSet::new();
|
||||
|
||||
tasks.spawn(async move {
|
||||
let ret = tunnel_rx.timeout(Duration::from_secs(60));
|
||||
tokio::pin!(ret);
|
||||
while let Ok(Some(Ok(p))) = ret.try_next().await {
|
||||
if let Err(e) = rpc_tx.send(p).await {
|
||||
tracing::error!("tunnel_rx send to rpc_tx error: {:?}", e);
|
||||
break;
|
||||
}
|
||||
}
|
||||
tracing::info!("forward tunnel_rx to rpc_tx done");
|
||||
});
|
||||
|
||||
tasks.spawn(async move {
|
||||
let ret = rpc_rx.forward(tunnel_tx).await;
|
||||
tracing::info!("rpc_rx forward tunnel_tx done: {:?}", ret);
|
||||
});
|
||||
|
||||
self.rpc_server.run();
|
||||
|
||||
while let Some(ret) = tasks.join_next().await {
|
||||
self.rpc_server.close();
|
||||
tracing::info!("task done: {:?}", ret);
|
||||
}
|
||||
|
||||
tracing::info!("all tasks done");
|
||||
}
|
||||
}
|
||||
use super::service_registry::ServiceRegistry;
|
||||
|
||||
pub struct StandAloneServer<L> {
|
||||
registry: Arc<ServiceRegistry>,
|
||||
@@ -102,11 +54,15 @@ impl<L: TunnelListener + 'static> StandAloneServer<L> {
|
||||
|
||||
self.tasks.lock().unwrap().spawn(async move {
|
||||
while let Ok(tunnel) = listener.accept().await {
|
||||
let server = StandAloneServerOneTunnel::new(tunnel, registry.clone());
|
||||
let registry = registry.clone();
|
||||
let inflight_server = inflight_server.clone();
|
||||
inflight_server.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
|
||||
tasks.lock().unwrap().spawn(async move {
|
||||
server.run().await;
|
||||
let server =
|
||||
BidirectRpcManager::new().set_rx_timeout(Some(Duration::from_secs(60)));
|
||||
server.rpc_server().registry().replace_registry(®istry);
|
||||
server.run_with_tunnel(tunnel);
|
||||
server.wait().await;
|
||||
inflight_server.fetch_sub(1, std::sync::atomic::Ordering::Relaxed);
|
||||
});
|
||||
}
|
||||
@@ -122,86 +78,9 @@ impl<L: TunnelListener + 'static> StandAloneServer<L> {
|
||||
}
|
||||
}
|
||||
|
||||
struct StandAloneClientOneTunnel {
|
||||
rpc_client: Client,
|
||||
tasks: Arc<Mutex<JoinSet<()>>>,
|
||||
error: Arc<Mutex<Option<Error>>>,
|
||||
}
|
||||
|
||||
impl StandAloneClientOneTunnel {
|
||||
pub fn new(tunnel: Box<dyn Tunnel>) -> Self {
|
||||
let rpc_client = Client::new();
|
||||
let (mut rpc_rx, rpc_tx) = (
|
||||
rpc_client.get_transport_stream(),
|
||||
rpc_client.get_transport_sink(),
|
||||
);
|
||||
let tasks = Arc::new(Mutex::new(JoinSet::new()));
|
||||
|
||||
let (mut tunnel_rx, mut tunnel_tx) = tunnel.split();
|
||||
|
||||
let error_store = Arc::new(Mutex::new(None));
|
||||
|
||||
let error = error_store.clone();
|
||||
tasks.lock().unwrap().spawn(async move {
|
||||
while let Some(p) = rpc_rx.next().await {
|
||||
match p {
|
||||
Ok(p) => {
|
||||
if let Err(e) = tunnel_tx
|
||||
.send(p)
|
||||
.await
|
||||
.with_context(|| "failed to send packet")
|
||||
{
|
||||
*error.lock().unwrap() = Some(e.into());
|
||||
}
|
||||
}
|
||||
Err(e) => {
|
||||
*error.lock().unwrap() = Some(anyhow::Error::from(e).into());
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
*error.lock().unwrap() = Some(anyhow::anyhow!("rpc_rx next exit").into());
|
||||
});
|
||||
|
||||
let error = error_store.clone();
|
||||
tasks.lock().unwrap().spawn(async move {
|
||||
while let Some(p) = tunnel_rx.next().await {
|
||||
match p {
|
||||
Ok(p) => {
|
||||
if let Err(e) = rpc_tx
|
||||
.send(p)
|
||||
.await
|
||||
.with_context(|| "failed to send packet")
|
||||
{
|
||||
*error.lock().unwrap() = Some(e.into());
|
||||
}
|
||||
}
|
||||
Err(e) => {
|
||||
*error.lock().unwrap() = Some(anyhow::Error::from(e).into());
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
*error.lock().unwrap() = Some(anyhow::anyhow!("tunnel_rx next exit").into());
|
||||
});
|
||||
|
||||
rpc_client.run();
|
||||
|
||||
StandAloneClientOneTunnel {
|
||||
rpc_client,
|
||||
tasks,
|
||||
error: error_store,
|
||||
}
|
||||
}
|
||||
|
||||
pub fn take_error(&self) -> Option<Error> {
|
||||
self.error.lock().unwrap().take()
|
||||
}
|
||||
}
|
||||
|
||||
pub struct StandAloneClient<C: TunnelConnector> {
|
||||
connector: C,
|
||||
client: Option<StandAloneClientOneTunnel>,
|
||||
client: Option<BidirectRpcManager>,
|
||||
}
|
||||
|
||||
impl<C: TunnelConnector> StandAloneClient<C> {
|
||||
@@ -230,7 +109,9 @@ impl<C: TunnelConnector> StandAloneClient<C> {
|
||||
if c.is_none() || error.is_some() {
|
||||
tracing::info!("reconnect due to error: {:?}", error);
|
||||
let tunnel = self.connect().await?;
|
||||
c = Some(StandAloneClientOneTunnel::new(tunnel));
|
||||
let mgr = BidirectRpcManager::new().set_rx_timeout(Some(Duration::from_secs(60)));
|
||||
mgr.run_with_tunnel(tunnel);
|
||||
c = Some(mgr);
|
||||
}
|
||||
|
||||
self.client = c;
|
||||
@@ -239,7 +120,7 @@ impl<C: TunnelConnector> StandAloneClient<C> {
|
||||
.client
|
||||
.as_ref()
|
||||
.unwrap()
|
||||
.rpc_client
|
||||
.rpc_client()
|
||||
.scoped_client::<F>(1, 1, domain_name))
|
||||
}
|
||||
}
|
||||
|
||||
@@ -29,6 +29,9 @@ pub enum Error {
|
||||
|
||||
#[error("Tunnel error: {0}")]
|
||||
TunnelError(#[from] crate::tunnel::TunnelError),
|
||||
|
||||
#[error("Shutdown")]
|
||||
Shutdown,
|
||||
}
|
||||
|
||||
pub type Result<T> = result::Result<T, Error>;
|
||||
|
||||
@@ -300,3 +300,75 @@ async fn standalone_rpc_test() {
|
||||
tokio::time::sleep(std::time::Duration::from_secs(1)).await;
|
||||
assert_eq!(0, server.inflight_server());
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_bidirect_rpc_manager() {
|
||||
use crate::common::scoped_task::ScopedTask;
|
||||
use crate::proto::rpc_impl::bidirect::BidirectRpcManager;
|
||||
use crate::tunnel::tcp::{TcpTunnelConnector, TcpTunnelListener};
|
||||
use crate::tunnel::{TunnelConnector, TunnelListener};
|
||||
|
||||
let c = BidirectRpcManager::new();
|
||||
let s = BidirectRpcManager::new();
|
||||
|
||||
let service = GreetingServer::new(GreetingService {
|
||||
delay_ms: 0,
|
||||
prefix: "Hello Client".to_string(),
|
||||
});
|
||||
c.rpc_server().registry().register(service, "test");
|
||||
|
||||
let service = GreetingServer::new(GreetingService {
|
||||
delay_ms: 0,
|
||||
prefix: "Hello Server".to_string(),
|
||||
});
|
||||
s.rpc_server().registry().register(service, "test");
|
||||
|
||||
let mut tcp_listener = TcpTunnelListener::new("tcp://0.0.0.0:55443".parse().unwrap());
|
||||
let s_task: ScopedTask<()> = tokio::spawn(async move {
|
||||
tcp_listener.listen().await.unwrap();
|
||||
let tunnel = tcp_listener.accept().await.unwrap();
|
||||
s.run_with_tunnel(tunnel);
|
||||
|
||||
let s_c = s
|
||||
.rpc_client()
|
||||
.scoped_client::<GreetingClientFactory<RpcController>>(1, 1, "test".to_string());
|
||||
let ret = s_c
|
||||
.say_hello(
|
||||
RpcController::default(),
|
||||
SayHelloRequest {
|
||||
name: "world".to_string(),
|
||||
},
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
assert_eq!(ret.greeting, "Hello Client world!");
|
||||
println!("server done, {:?}", ret);
|
||||
|
||||
s.wait().await;
|
||||
})
|
||||
.into();
|
||||
|
||||
tokio::time::sleep(std::time::Duration::from_secs(1)).await;
|
||||
|
||||
let mut tcp_connector = TcpTunnelConnector::new("tcp://0.0.0.0:55443".parse().unwrap());
|
||||
let c_tunnel = tcp_connector.connect().await.unwrap();
|
||||
c.run_with_tunnel(c_tunnel);
|
||||
|
||||
let c_c = c
|
||||
.rpc_client()
|
||||
.scoped_client::<GreetingClientFactory<RpcController>>(1, 1, "test".to_string());
|
||||
let ret = c_c
|
||||
.say_hello(
|
||||
RpcController::default(),
|
||||
SayHelloRequest {
|
||||
name: "world".to_string(),
|
||||
},
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
assert_eq!(ret.greeting, "Hello Server world!");
|
||||
println!("client done, {:?}", ret);
|
||||
|
||||
drop(c);
|
||||
s_task.await.unwrap();
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user