mirror of
https://github.com/edera-dev/krata.git
synced 2025-11-03 23:29:39 +00:00
krata: utilize gRPC for control service
This commit is contained in:
Alex Zenla
172
daemon/src/control.rs
Normal file
172
daemon/src/control.rs
Normal file
@ -0,0 +1,172 @@
|
||||
use std::{io, pin::Pin};
|
||||
|
||||
use async_stream::try_stream;
|
||||
use futures::Stream;
|
||||
use krata::control::{
|
||||
control_service_server::ControlService, ConsoleDataReply, ConsoleDataRequest,
|
||||
DestroyGuestReply, DestroyGuestRequest, GuestInfo, LaunchGuestReply, LaunchGuestRequest,
|
||||
ListGuestsReply, ListGuestsRequest,
|
||||
};
|
||||
use tokio::{
|
||||
io::{AsyncReadExt, AsyncWriteExt},
|
||||
select,
|
||||
};
|
||||
use tokio_stream::StreamExt;
|
||||
use tonic::{Request, Response, Status, Streaming};
|
||||
|
||||
use crate::runtime::{launch::GuestLaunchRequest, Runtime};
|
||||
|
||||
pub struct ApiError {
|
||||
message: String,
|
||||
}
|
||||
|
||||
impl From<anyhow::Error> for ApiError {
|
||||
fn from(value: anyhow::Error) -> Self {
|
||||
ApiError {
|
||||
message: value.to_string(),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl From<ApiError> for Status {
|
||||
fn from(value: ApiError) -> Self {
|
||||
Status::unknown(value.message)
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Clone)]
|
||||
pub struct RuntimeControlService {
|
||||
runtime: Runtime,
|
||||
}
|
||||
|
||||
impl RuntimeControlService {
|
||||
pub fn new(runtime: Runtime) -> Self {
|
||||
Self { runtime }
|
||||
}
|
||||
}
|
||||
|
||||
enum ConsoleDataSelect {
|
||||
Read(io::Result<usize>),
|
||||
Write(Option<Result<ConsoleDataRequest, tonic::Status>>),
|
||||
}
|
||||
|
||||
#[tonic::async_trait]
|
||||
impl ControlService for RuntimeControlService {
|
||||
type ConsoleDataStream =
|
||||
Pin<Box<dyn Stream<Item = Result<ConsoleDataReply, Status>> + Send + 'static>>;
|
||||
|
||||
async fn launch_guest(
|
||||
&self,
|
||||
request: Request<LaunchGuestRequest>,
|
||||
) -> Result<Response<LaunchGuestReply>, Status> {
|
||||
let request = request.into_inner();
|
||||
let guest: GuestInfo = self
|
||||
.runtime
|
||||
.launch(GuestLaunchRequest {
|
||||
image: &request.image,
|
||||
vcpus: request.vcpus,
|
||||
mem: request.mem,
|
||||
env: empty_vec_optional(request.env),
|
||||
run: empty_vec_optional(request.run),
|
||||
debug: false,
|
||||
})
|
||||
.await
|
||||
.map_err(ApiError::from)?
|
||||
.into();
|
||||
Ok(Response::new(LaunchGuestReply { guest: Some(guest) }))
|
||||
}
|
||||
|
||||
async fn destroy_guest(
|
||||
&self,
|
||||
request: Request<DestroyGuestRequest>,
|
||||
) -> Result<Response<DestroyGuestReply>, Status> {
|
||||
let request = request.into_inner();
|
||||
self.runtime
|
||||
.destroy(&request.guest_id)
|
||||
.await
|
||||
.map_err(ApiError::from)?;
|
||||
Ok(Response::new(DestroyGuestReply {}))
|
||||
}
|
||||
|
||||
async fn list_guests(
|
||||
&self,
|
||||
request: Request<ListGuestsRequest>,
|
||||
) -> Result<Response<ListGuestsReply>, Status> {
|
||||
let _ = request.into_inner();
|
||||
let guests = self.runtime.list().await.map_err(ApiError::from)?;
|
||||
let guests = guests
|
||||
.into_iter()
|
||||
.map(GuestInfo::from)
|
||||
.collect::<Vec<GuestInfo>>();
|
||||
Ok(Response::new(ListGuestsReply { guests }))
|
||||
}
|
||||
|
||||
async fn console_data(
|
||||
&self,
|
||||
request: Request<Streaming<ConsoleDataRequest>>,
|
||||
) -> Result<Response<Self::ConsoleDataStream>, Status> {
|
||||
let mut input = request.into_inner();
|
||||
let Some(request) = input.next().await else {
|
||||
return Err(ApiError {
|
||||
message: "expected to have at least one request".to_string(),
|
||||
}
|
||||
.into());
|
||||
};
|
||||
let request = request?;
|
||||
let mut console = self
|
||||
.runtime
|
||||
.console(&request.guest)
|
||||
.await
|
||||
.map_err(ApiError::from)?;
|
||||
|
||||
let output = try_stream! {
|
||||
let mut buffer: Vec<u8> = vec![0u8; 256];
|
||||
loop {
|
||||
let what = select! {
|
||||
x = console.read_handle.read(&mut buffer) => ConsoleDataSelect::Read(x),
|
||||
x = input.next() => ConsoleDataSelect::Write(x),
|
||||
};
|
||||
|
||||
match what {
|
||||
ConsoleDataSelect::Read(result) => {
|
||||
let size = result?;
|
||||
let data = buffer[0..size].to_vec();
|
||||
yield ConsoleDataReply { data, };
|
||||
},
|
||||
|
||||
ConsoleDataSelect::Write(Some(request)) => {
|
||||
let request = request?;
|
||||
if !request.data.is_empty() {
|
||||
console.write_handle.write_all(&request.data).await?;
|
||||
}
|
||||
},
|
||||
|
||||
ConsoleDataSelect::Write(None) => {
|
||||
break;
|
||||
}
|
||||
}
|
||||
}
|
||||
};
|
||||
|
||||
Ok(Response::new(Box::pin(output) as Self::ConsoleDataStream))
|
||||
}
|
||||
}
|
||||
|
||||
impl From<crate::runtime::GuestInfo> for GuestInfo {
|
||||
fn from(value: crate::runtime::GuestInfo) -> Self {
|
||||
GuestInfo {
|
||||
id: value.uuid.to_string(),
|
||||
image: value.image,
|
||||
ipv4: value.ipv4.map(|x| x.ip().to_string()).unwrap_or_default(),
|
||||
ipv6: value.ipv6.map(|x| x.ip().to_string()).unwrap_or_default(),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
fn empty_vec_optional<T>(value: Vec<T>) -> Option<Vec<T>> {
|
||||
if value.is_empty() {
|
||||
None
|
||||
} else {
|
||||
Some(value)
|
||||
}
|
||||
}
|
||||
@ -1,91 +0,0 @@
|
||||
use anyhow::{anyhow, Result};
|
||||
use krata::control::{ConsoleStreamResponse, ConsoleStreamUpdate, Request, Response, StreamUpdate};
|
||||
use log::warn;
|
||||
use tokio::{
|
||||
io::{AsyncReadExt, AsyncWriteExt},
|
||||
select,
|
||||
};
|
||||
|
||||
use crate::{
|
||||
listen::DaemonRequestHandler,
|
||||
runtime::{console::XenConsole, Runtime},
|
||||
};
|
||||
use krata::stream::{ConnectionStreams, StreamContext};
|
||||
pub struct ConsoleStreamRequestHandler {}
|
||||
|
||||
impl Default for ConsoleStreamRequestHandler {
|
||||
fn default() -> Self {
|
||||
Self::new()
|
||||
}
|
||||
}
|
||||
|
||||
impl ConsoleStreamRequestHandler {
|
||||
pub fn new() -> Self {
|
||||
Self {}
|
||||
}
|
||||
|
||||
async fn link_console_stream(mut stream: StreamContext, mut console: XenConsole) -> Result<()> {
|
||||
loop {
|
||||
let mut buffer = vec![0u8; 256];
|
||||
select! {
|
||||
x = console.read_handle.read(&mut buffer) => match x {
|
||||
Ok(size) => {
|
||||
let data = buffer[0..size].to_vec();
|
||||
let update = StreamUpdate::ConsoleStream(ConsoleStreamUpdate {
|
||||
data,
|
||||
});
|
||||
stream.send(update).await?;
|
||||
},
|
||||
|
||||
Err(error) => {
|
||||
return Err(error.into());
|
||||
}
|
||||
},
|
||||
|
||||
x = stream.receiver.recv() => match x {
|
||||
Some(StreamUpdate::ConsoleStream(update)) => {
|
||||
console.write_handle.write_all(&update.data).await?;
|
||||
}
|
||||
|
||||
None => {
|
||||
break;
|
||||
}
|
||||
}
|
||||
};
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait::async_trait]
|
||||
impl DaemonRequestHandler for ConsoleStreamRequestHandler {
|
||||
fn accepts(&self, request: &Request) -> bool {
|
||||
matches!(request, Request::ConsoleStream(_))
|
||||
}
|
||||
|
||||
async fn handle(
|
||||
&self,
|
||||
streams: ConnectionStreams,
|
||||
runtime: Runtime,
|
||||
request: Request,
|
||||
) -> Result<Response> {
|
||||
let console_stream = match request {
|
||||
Request::ConsoleStream(stream) => stream,
|
||||
_ => return Err(anyhow!("unknown request")),
|
||||
};
|
||||
let console = runtime.console(&console_stream.guest).await?;
|
||||
let stream = streams.open().await?;
|
||||
let id = stream.id;
|
||||
tokio::task::spawn(async move {
|
||||
if let Err(error) =
|
||||
ConsoleStreamRequestHandler::link_console_stream(stream, console).await
|
||||
{
|
||||
warn!("failed to process console stream: {}", error);
|
||||
}
|
||||
});
|
||||
|
||||
Ok(Response::ConsoleStream(ConsoleStreamResponse {
|
||||
stream: id,
|
||||
}))
|
||||
}
|
||||
}
|
||||
@ -1,44 +0,0 @@
|
||||
use anyhow::{anyhow, Result};
|
||||
use krata::{
|
||||
control::{DestroyResponse, Request, Response},
|
||||
stream::ConnectionStreams,
|
||||
};
|
||||
|
||||
use crate::{listen::DaemonRequestHandler, runtime::Runtime};
|
||||
|
||||
pub struct DestroyRequestHandler {}
|
||||
|
||||
impl Default for DestroyRequestHandler {
|
||||
fn default() -> Self {
|
||||
Self::new()
|
||||
}
|
||||
}
|
||||
|
||||
impl DestroyRequestHandler {
|
||||
pub fn new() -> Self {
|
||||
Self {}
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait::async_trait]
|
||||
impl DaemonRequestHandler for DestroyRequestHandler {
|
||||
fn accepts(&self, request: &Request) -> bool {
|
||||
matches!(request, Request::Destroy(_))
|
||||
}
|
||||
|
||||
async fn handle(
|
||||
&self,
|
||||
_: ConnectionStreams,
|
||||
runtime: Runtime,
|
||||
request: Request,
|
||||
) -> Result<Response> {
|
||||
let destroy = match request {
|
||||
Request::Destroy(destroy) => destroy,
|
||||
_ => return Err(anyhow!("unknown request")),
|
||||
};
|
||||
let guest = runtime.destroy(&destroy.guest).await?;
|
||||
Ok(Response::Destroy(DestroyResponse {
|
||||
guest: guest.to_string(),
|
||||
}))
|
||||
}
|
||||
}
|
||||
@ -1,55 +0,0 @@
|
||||
use anyhow::{anyhow, Result};
|
||||
use krata::{
|
||||
control::{GuestInfo, LaunchResponse, Request, Response},
|
||||
stream::ConnectionStreams,
|
||||
};
|
||||
|
||||
use crate::{
|
||||
listen::DaemonRequestHandler,
|
||||
runtime::{launch::GuestLaunchRequest, Runtime},
|
||||
};
|
||||
|
||||
pub struct LaunchRequestHandler {}
|
||||
|
||||
impl Default for LaunchRequestHandler {
|
||||
fn default() -> Self {
|
||||
Self::new()
|
||||
}
|
||||
}
|
||||
|
||||
impl LaunchRequestHandler {
|
||||
pub fn new() -> Self {
|
||||
Self {}
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait::async_trait]
|
||||
impl DaemonRequestHandler for LaunchRequestHandler {
|
||||
fn accepts(&self, request: &Request) -> bool {
|
||||
matches!(request, Request::Launch(_))
|
||||
}
|
||||
|
||||
async fn handle(
|
||||
&self,
|
||||
_: ConnectionStreams,
|
||||
runtime: Runtime,
|
||||
request: Request,
|
||||
) -> Result<Response> {
|
||||
let launch = match request {
|
||||
Request::Launch(launch) => launch,
|
||||
_ => return Err(anyhow!("unknown request")),
|
||||
};
|
||||
let guest: GuestInfo = runtime
|
||||
.launch(GuestLaunchRequest {
|
||||
image: &launch.image,
|
||||
vcpus: launch.vcpus,
|
||||
mem: launch.mem,
|
||||
env: launch.env,
|
||||
run: launch.run,
|
||||
debug: false,
|
||||
})
|
||||
.await?
|
||||
.into();
|
||||
Ok(Response::Launch(LaunchResponse { guest }))
|
||||
}
|
||||
}
|
||||
@ -1,37 +0,0 @@
|
||||
use anyhow::Result;
|
||||
use krata::{
|
||||
control::{GuestInfo, ListResponse, Request, Response},
|
||||
stream::ConnectionStreams,
|
||||
};
|
||||
|
||||
use crate::{listen::DaemonRequestHandler, runtime::Runtime};
|
||||
|
||||
pub struct ListRequestHandler {}
|
||||
|
||||
impl Default for ListRequestHandler {
|
||||
fn default() -> Self {
|
||||
Self::new()
|
||||
}
|
||||
}
|
||||
|
||||
impl ListRequestHandler {
|
||||
pub fn new() -> Self {
|
||||
Self {}
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait::async_trait]
|
||||
impl DaemonRequestHandler for ListRequestHandler {
|
||||
fn accepts(&self, request: &Request) -> bool {
|
||||
matches!(request, Request::List(_))
|
||||
}
|
||||
|
||||
async fn handle(&self, _: ConnectionStreams, runtime: Runtime, _: Request) -> Result<Response> {
|
||||
let guests = runtime.list().await?;
|
||||
let guests = guests
|
||||
.into_iter()
|
||||
.map(GuestInfo::from)
|
||||
.collect::<Vec<GuestInfo>>();
|
||||
Ok(Response::List(ListResponse { guests }))
|
||||
}
|
||||
}
|
||||
@ -1,15 +0,0 @@
|
||||
pub mod console;
|
||||
pub mod destroy;
|
||||
pub mod launch;
|
||||
pub mod list;
|
||||
|
||||
impl From<crate::runtime::GuestInfo> for krata::control::GuestInfo {
|
||||
fn from(value: crate::runtime::GuestInfo) -> Self {
|
||||
krata::control::GuestInfo {
|
||||
id: value.uuid.to_string(),
|
||||
image: value.image.clone(),
|
||||
ipv4: value.ipv4.map(|x| x.ip().to_string()),
|
||||
ipv6: value.ipv6.map(|x| x.ip().to_string()),
|
||||
}
|
||||
}
|
||||
}
|
||||
@ -1,37 +1,74 @@
|
||||
use anyhow::Result;
|
||||
use handlers::{
|
||||
console::ConsoleStreamRequestHandler, destroy::DestroyRequestHandler,
|
||||
launch::LaunchRequestHandler, list::ListRequestHandler,
|
||||
};
|
||||
use listen::{DaemonListener, DaemonRequestHandlers};
|
||||
use runtime::Runtime;
|
||||
use tokio_listener::Listener;
|
||||
use std::{net::SocketAddr, path::PathBuf, str::FromStr};
|
||||
|
||||
pub mod handlers;
|
||||
pub mod listen;
|
||||
use anyhow::Result;
|
||||
use control::RuntimeControlService;
|
||||
use krata::{control::control_service_server::ControlServiceServer, dial::ControlDialAddress};
|
||||
use log::info;
|
||||
use runtime::Runtime;
|
||||
use tokio::net::UnixListener;
|
||||
use tokio_stream::wrappers::UnixListenerStream;
|
||||
use tonic::transport::{Identity, Server, ServerTlsConfig};
|
||||
|
||||
pub mod control;
|
||||
pub mod runtime;
|
||||
|
||||
pub struct Daemon {
|
||||
store: String,
|
||||
runtime: Runtime,
|
||||
}
|
||||
|
||||
impl Daemon {
|
||||
pub async fn new(runtime: Runtime) -> Result<Self> {
|
||||
Ok(Self { runtime })
|
||||
pub async fn new(store: String, runtime: Runtime) -> Result<Self> {
|
||||
Ok(Self { store, runtime })
|
||||
}
|
||||
|
||||
pub async fn listen(&mut self, listener: Listener) -> Result<()> {
|
||||
let handlers = DaemonRequestHandlers::new(
|
||||
self.runtime.clone(),
|
||||
vec![
|
||||
Box::new(LaunchRequestHandler::new()),
|
||||
Box::new(DestroyRequestHandler::new()),
|
||||
Box::new(ConsoleStreamRequestHandler::new()),
|
||||
Box::new(ListRequestHandler::new()),
|
||||
],
|
||||
);
|
||||
let mut listener = DaemonListener::new(listener, handlers);
|
||||
listener.handle().await?;
|
||||
pub async fn listen(&mut self, addr: ControlDialAddress) -> Result<()> {
|
||||
let control_service = RuntimeControlService::new(self.runtime.clone());
|
||||
|
||||
let mut server = Server::builder();
|
||||
|
||||
if let ControlDialAddress::Tls {
|
||||
host: _,
|
||||
port: _,
|
||||
insecure,
|
||||
} = &addr
|
||||
{
|
||||
let mut tls_config = ServerTlsConfig::new();
|
||||
if !insecure {
|
||||
let certificate_path = format!("{}/tls/daemon.pem", self.store);
|
||||
let key_path = format!("{}/tls/daemon.key", self.store);
|
||||
tls_config = tls_config.identity(Identity::from_pem(certificate_path, key_path));
|
||||
}
|
||||
server = server.tls_config(tls_config)?;
|
||||
}
|
||||
|
||||
let server = server.add_service(ControlServiceServer::new(control_service));
|
||||
info!("listening on address {}", addr);
|
||||
match addr {
|
||||
ControlDialAddress::UnixSocket { path } => {
|
||||
let path = PathBuf::from(path);
|
||||
if path.exists() {
|
||||
tokio::fs::remove_file(&path).await?;
|
||||
}
|
||||
let listener = UnixListener::bind(path)?;
|
||||
let stream = UnixListenerStream::new(listener);
|
||||
server.serve_with_incoming(stream).await?;
|
||||
}
|
||||
|
||||
ControlDialAddress::Tcp { host, port } => {
|
||||
let address = format!("{}:{}", host, port);
|
||||
server.serve(SocketAddr::from_str(&address)?).await?;
|
||||
}
|
||||
|
||||
ControlDialAddress::Tls {
|
||||
host,
|
||||
port,
|
||||
insecure: _,
|
||||
} => {
|
||||
let address = format!("{}:{}", host, port);
|
||||
server.serve(SocketAddr::from_str(&address)?).await?;
|
||||
}
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
|
||||
@ -1,228 +0,0 @@
|
||||
use std::collections::HashMap;
|
||||
use std::sync::Arc;
|
||||
|
||||
use anyhow::{anyhow, Result};
|
||||
use krata::control::{ErrorResponse, Message, Request, RequestBox, Response, ResponseBox};
|
||||
use log::trace;
|
||||
use log::warn;
|
||||
use tokio::sync::Mutex;
|
||||
use tokio::{
|
||||
io::{AsyncBufReadExt, AsyncWriteExt, BufReader},
|
||||
select,
|
||||
sync::mpsc::{channel, Receiver, Sender},
|
||||
};
|
||||
use tokio_listener::{Connection, Listener, SomeSocketAddrClonable};
|
||||
use tokio_stream::{wrappers::LinesStream, StreamExt};
|
||||
|
||||
use crate::runtime::Runtime;
|
||||
use krata::stream::ConnectionStreams;
|
||||
|
||||
const QUEUE_MAX_LEN: usize = 100;
|
||||
|
||||
#[async_trait::async_trait]
|
||||
pub trait DaemonRequestHandler: Send + Sync {
|
||||
fn accepts(&self, request: &Request) -> bool;
|
||||
async fn handle(
|
||||
&self,
|
||||
streams: ConnectionStreams,
|
||||
runtime: Runtime,
|
||||
request: Request,
|
||||
) -> Result<Response>;
|
||||
}
|
||||
|
||||
#[derive(Clone)]
|
||||
pub struct DaemonRequestHandlers {
|
||||
runtime: Runtime,
|
||||
handlers: Arc<Vec<Box<dyn DaemonRequestHandler>>>,
|
||||
}
|
||||
|
||||
impl DaemonRequestHandlers {
|
||||
pub fn new(runtime: Runtime, handlers: Vec<Box<dyn DaemonRequestHandler>>) -> Self {
|
||||
DaemonRequestHandlers {
|
||||
runtime,
|
||||
handlers: Arc::new(handlers),
|
||||
}
|
||||
}
|
||||
|
||||
async fn dispatch(&self, streams: ConnectionStreams, request: Request) -> Result<Response> {
|
||||
for handler in self.handlers.iter() {
|
||||
if handler.accepts(&request) {
|
||||
return handler.handle(streams, self.runtime.clone(), request).await;
|
||||
}
|
||||
}
|
||||
Err(anyhow!("daemon cannot handle that request"))
|
||||
}
|
||||
}
|
||||
|
||||
pub struct DaemonListener {
|
||||
listener: Listener,
|
||||
handlers: DaemonRequestHandlers,
|
||||
connections: Arc<Mutex<HashMap<u64, DaemonConnection>>>,
|
||||
next: Arc<Mutex<u64>>,
|
||||
}
|
||||
|
||||
impl DaemonListener {
|
||||
pub fn new(listener: Listener, handlers: DaemonRequestHandlers) -> DaemonListener {
|
||||
DaemonListener {
|
||||
listener,
|
||||
handlers,
|
||||
connections: Arc::new(Mutex::new(HashMap::new())),
|
||||
next: Arc::new(Mutex::new(0)),
|
||||
}
|
||||
}
|
||||
|
||||
pub async fn handle(&mut self) -> Result<()> {
|
||||
loop {
|
||||
let (connection, addr) = self.listener.accept().await?;
|
||||
let connection =
|
||||
DaemonConnection::new(connection, addr.clonable(), self.handlers.clone()).await?;
|
||||
let id = {
|
||||
let mut next = self.next.lock().await;
|
||||
let id = *next;
|
||||
*next = id + 1;
|
||||
id
|
||||
};
|
||||
trace!("new connection from {}", connection.addr);
|
||||
let tx_channel = connection.tx_sender.clone();
|
||||
let addr = connection.addr.clone();
|
||||
self.connections.lock().await.insert(id, connection);
|
||||
let connections_for_close = self.connections.clone();
|
||||
tokio::task::spawn(async move {
|
||||
tx_channel.closed().await;
|
||||
trace!("connection from {} closed", addr);
|
||||
connections_for_close.lock().await.remove(&id);
|
||||
});
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Clone)]
|
||||
pub struct DaemonConnection {
|
||||
tx_sender: Sender<Message>,
|
||||
addr: SomeSocketAddrClonable,
|
||||
handlers: DaemonRequestHandlers,
|
||||
streams: ConnectionStreams,
|
||||
}
|
||||
|
||||
impl DaemonConnection {
|
||||
pub async fn new(
|
||||
connection: Connection,
|
||||
addr: SomeSocketAddrClonable,
|
||||
handlers: DaemonRequestHandlers,
|
||||
) -> Result<Self> {
|
||||
let (tx_sender, tx_receiver) = channel::<Message>(QUEUE_MAX_LEN);
|
||||
let streams_tx_sender = tx_sender.clone();
|
||||
let instance = DaemonConnection {
|
||||
tx_sender,
|
||||
addr,
|
||||
handlers,
|
||||
streams: ConnectionStreams::new(streams_tx_sender),
|
||||
};
|
||||
|
||||
{
|
||||
let mut instance = instance.clone();
|
||||
tokio::task::spawn(async move {
|
||||
if let Err(error) = instance.process(tx_receiver, connection).await {
|
||||
warn!(
|
||||
"failed to process daemon connection for {}: {}",
|
||||
instance.addr, error
|
||||
);
|
||||
}
|
||||
});
|
||||
}
|
||||
|
||||
Ok(instance)
|
||||
}
|
||||
|
||||
async fn process(
|
||||
&mut self,
|
||||
mut tx_receiver: Receiver<Message>,
|
||||
connection: Connection,
|
||||
) -> Result<()> {
|
||||
let (read, mut write) = tokio::io::split(connection);
|
||||
let mut read = LinesStream::new(BufReader::new(read).lines());
|
||||
|
||||
loop {
|
||||
select! {
|
||||
x = read.next() => match x {
|
||||
Some(Ok(line)) => {
|
||||
let message: Message = serde_json::from_str(&line)?;
|
||||
trace!("received message '{}' from {}", serde_json::to_string(&message)?, self.addr);
|
||||
let mut context = self.clone();
|
||||
tokio::task::spawn(async move {
|
||||
if let Err(error) = context.handle_message(&message).await {
|
||||
let line = serde_json::to_string(&message).unwrap_or("<invalid>".to_string());
|
||||
warn!("failed to handle message '{}' from {}: {}", line, context.addr, error);
|
||||
}
|
||||
});
|
||||
},
|
||||
|
||||
Some(Err(error)) => {
|
||||
return Err(error.into());
|
||||
},
|
||||
|
||||
None => {
|
||||
break;
|
||||
}
|
||||
},
|
||||
|
||||
x = tx_receiver.recv() => match x {
|
||||
Some(message) => {
|
||||
if let Message::StreamUpdated(ref update) = message {
|
||||
self.streams.outgoing(update).await?;
|
||||
}
|
||||
let mut line = serde_json::to_string(&message)?;
|
||||
trace!("sending message '{}' to {}", line, self.addr);
|
||||
line.push('\n');
|
||||
write.write_all(line.as_bytes()).await?;
|
||||
},
|
||||
None => {
|
||||
break;
|
||||
}
|
||||
}
|
||||
};
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn handle_message(&mut self, message: &Message) -> Result<()> {
|
||||
match message {
|
||||
Message::Request(req) => {
|
||||
self.handle_request(req.clone()).await?;
|
||||
}
|
||||
|
||||
Message::Response(_) => {
|
||||
return Err(anyhow!(
|
||||
"received a response message from client {}, but this is the daemon",
|
||||
self.addr
|
||||
));
|
||||
}
|
||||
|
||||
Message::StreamUpdated(updated) => {
|
||||
self.streams.incoming(updated.clone()).await?;
|
||||
}
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn handle_request(&mut self, req: RequestBox) -> Result<()> {
|
||||
let id = req.id;
|
||||
let response = self
|
||||
.handlers
|
||||
.dispatch(self.streams.clone(), req.request)
|
||||
.await
|
||||
.map_err(|error| {
|
||||
Response::Error(ErrorResponse {
|
||||
message: error.to_string(),
|
||||
})
|
||||
});
|
||||
let response = if let Err(response) = response {
|
||||
response
|
||||
} else {
|
||||
response.unwrap()
|
||||
};
|
||||
let resp = ResponseBox { id, response };
|
||||
self.tx_sender.send(Message::Response(resp)).await?;
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user