-
Notifications
You must be signed in to change notification settings - Fork 3
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
refactor(control): simplify endpoint definition
- Loading branch information
Showing
9 changed files
with
170 additions
and
209 deletions.
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -1,103 +1,13 @@ | ||
| use std::future::Future; | ||
| use std::path::Path; | ||
|
|
||
| use futures_util::future::BoxFuture; | ||
| use futures_util::{FutureExt, StreamExt}; | ||
|
|
||
| pub use self::client::ControlClient; | ||
| pub use self::error::{ClientError, ServerResult}; | ||
| pub use self::server::impls::*; | ||
| pub use self::server::{ | ||
| ArchiveInfo, ArchiveInfoRequest, ArchiveInfoResponse, ArchiveSliceRequest, | ||
| ArchiveSliceResponse, BlockProofRequest, BlockProofResponse, BlockRequest, BlockResponse, | ||
| ControlServer, | ||
| }; | ||
|
|
||
| pub type Context = tarpc::context::Context; | ||
| pub use self::profiler::{MemoryProfiler, StubMemoryProfiler}; | ||
| pub use self::server::{ControlEndpoint, ControlServer, ControlServerBuilder, ControlServerConfig}; | ||
|
|
||
| mod client; | ||
| mod error; | ||
| mod profiler; | ||
| mod proto; | ||
| mod server; | ||
|
|
||
| // TODO: Change the path to a more general setup. | ||
| pub const DEFAULT_SOCKET_PATH: &str = "/var/venom/data/tycho.sock"; | ||
|
|
||
| pub struct ControlEndpoint { | ||
| inner: BoxFuture<'static, ()>, | ||
| } | ||
|
|
||
| impl ControlEndpoint { | ||
| pub async fn bind<P, S>(path: P, server: S) -> std::io::Result<Self> | ||
| where | ||
| P: AsRef<Path>, | ||
| S: ControlServerExt, | ||
| { | ||
| use tarpc::tokio_serde::formats::Bincode; | ||
|
|
||
| let mut listener = tarpc::serde_transport::unix::listen(path, Bincode::default).await?; | ||
| listener.config_mut().max_frame_length(usize::MAX); | ||
|
|
||
| let inner = listener | ||
| // Ignore accept errors. | ||
| .filter_map(|r| futures_util::future::ready(r.ok())) | ||
| .map(tarpc::server::BaseChannel::with_defaults) | ||
| .map(move |channel| server.clone().execute_all(channel)) | ||
| // Max 1 channel. | ||
| .buffer_unordered(1) | ||
| .for_each(|_| async {}) | ||
| .boxed(); | ||
|
|
||
| Ok(Self { inner }) | ||
| } | ||
|
|
||
| pub async fn serve(self) { | ||
| self.inner.await; | ||
| } | ||
| } | ||
|
|
||
| // FIXME: Remove when https://github.com/google/tarpc/pull/448 is merged. | ||
| #[macro_export] | ||
| macro_rules! impl_serve { | ||
| ($ident:ident) => { | ||
| impl $crate::ControlServerExt for $ident { | ||
| fn execute_all<T>( | ||
| self, | ||
| channel: $crate::__internal::RawChannel<T>, | ||
| ) -> impl ::std::future::Future<Output = ()> + Send | ||
| where | ||
| T: tarpc::Transport< | ||
| $crate::__internal::tarpc::Response<$crate::__internal::RawResponse>, | ||
| $crate::__internal::tarpc::ClientMessage<$crate::__internal::RawRequest>, | ||
| > + Send | ||
| + 'static, | ||
| { | ||
| use $crate::__internal::futures_util::{future, StreamExt}; | ||
| use $crate::__internal::tarpc::server::Channel; | ||
|
|
||
| channel.execute(self.serve()).for_each(|t| { | ||
| $crate::__internal::tokio::spawn(t); | ||
| future::ready(()) | ||
| }) | ||
| } | ||
| } | ||
| }; | ||
| } | ||
|
|
||
| pub trait ControlServerExt: ControlServer + Clone + Send + 'static { | ||
| fn execute_all<T>(self, channel: __internal::RawChannel<T>) -> impl Future<Output = ()> + Send | ||
| where | ||
| T: tarpc::Transport< | ||
| tarpc::Response<__internal::RawResponse>, | ||
| tarpc::ClientMessage<__internal::RawRequest>, | ||
| > + Send | ||
| + 'static; | ||
| } | ||
|
|
||
| #[doc(hidden)] | ||
| pub mod __internal { | ||
| pub use {futures_util, tarpc, tokio}; | ||
|
|
||
| pub type RawChannel<T> = tarpc::server::BaseChannel<RawRequest, RawResponse, T>; | ||
| pub type RawRequest = crate::server::ControlServerRequest; | ||
| pub type RawResponse = crate::server::ControlServerResponse; | ||
| } |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,34 @@ | ||
| use std::sync::Arc; | ||
|
|
||
| use anyhow::Result; | ||
|
|
||
| #[async_trait::async_trait] | ||
| pub trait MemoryProfiler: Send + Sync + 'static { | ||
| async fn set_enabled(&self, enabled: bool) -> bool; | ||
| async fn dump(&self) -> Result<Vec<u8>>; | ||
| } | ||
|
|
||
| #[async_trait::async_trait] | ||
| impl<T: MemoryProfiler> MemoryProfiler for Arc<T> { | ||
| async fn set_enabled(&self, enabled: bool) -> bool { | ||
| T::set_enabled(self, enabled).await | ||
| } | ||
|
|
||
| async fn dump(&self) -> Result<Vec<u8>> { | ||
| T::dump(self).await | ||
| } | ||
| } | ||
|
|
||
| #[derive(Debug, Clone, Copy)] | ||
| pub struct StubMemoryProfiler; | ||
|
|
||
| #[async_trait::async_trait] | ||
| impl MemoryProfiler for StubMemoryProfiler { | ||
| async fn set_enabled(&self, _: bool) -> bool { | ||
| false | ||
| } | ||
|
|
||
| async fn dump(&self) -> Result<Vec<u8>> { | ||
| anyhow::bail!("stub memory profiler does not support dumping data") | ||
| } | ||
| } |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Oops, something went wrong.