154 lines
6.6 KiB
Rust
154 lines
6.6 KiB
Rust
use crate::command::service::CommandManagementService;
|
|
use crate::core::client_side_command::Inner;
|
|
use crate::core::command_service_server::CommandService;
|
|
use crate::core::{ClientSideCommand, Command, CommandResponse, Uuid};
|
|
use log::{error, trace};
|
|
use std::collections::HashMap;
|
|
use std::pin::Pin;
|
|
use std::sync::Arc;
|
|
use tokio::select;
|
|
use tokio::sync::mpsc;
|
|
use tokio::sync::oneshot;
|
|
use tokio_util::sync::CancellationToken;
|
|
use tonic::codegen::tokio_stream::wrappers::ReceiverStream;
|
|
use tonic::codegen::tokio_stream::{Stream, StreamExt};
|
|
use tonic::{Request, Response, Status, Streaming};
|
|
|
|
pub struct CoreCommandService {
|
|
pub command_service: Arc<CommandManagementService>,
|
|
pub cancellation_token: CancellationToken,
|
|
}
|
|
|
|
#[tonic::async_trait]
|
|
impl CommandService for CoreCommandService {
|
|
type NewCommandStream = Pin<Box<dyn Stream<Item = Result<Command, Status>> + Send>>;
|
|
|
|
async fn new_command(
|
|
&self,
|
|
request: Request<Streaming<ClientSideCommand>>,
|
|
) -> Result<Response<Self::NewCommandStream>, Status> {
|
|
trace!("CoreCommandService::new_command");
|
|
|
|
let cancel_token = self.cancellation_token.clone();
|
|
let mut in_stream = request.into_inner();
|
|
|
|
let cmd_request = select! {
|
|
_ = cancel_token.cancelled() => return Err(Status::internal("Shutting Down")),
|
|
Some(message) = in_stream.next() => {
|
|
match message {
|
|
Ok(ClientSideCommand {
|
|
inner: Some(Inner::Request(cmd_request))
|
|
}) => cmd_request,
|
|
Err(err) => {
|
|
error!("Error in Stream: {err}");
|
|
return Err(Status::cancelled("Error in Stream"));
|
|
},
|
|
_ => {
|
|
return Err(Status::invalid_argument("First message must be request"));
|
|
},
|
|
}
|
|
},
|
|
else => return Err(Status::internal("Shutting Down")),
|
|
};
|
|
|
|
let mut cmd_rx = match self.command_service.register_command(cmd_request).await {
|
|
Ok(rx) => rx,
|
|
Err(e) => {
|
|
error!("Failed to register command: {e}");
|
|
return Err(Status::internal("Failed to register command"));
|
|
}
|
|
};
|
|
|
|
let (tx, rx) = mpsc::channel(128);
|
|
|
|
tokio::spawn(async move {
|
|
let mut result = Status::resource_exhausted("End of Command Stream");
|
|
let mut in_progress = HashMap::<String, oneshot::Sender<CommandResponse>>::new();
|
|
loop {
|
|
select! {
|
|
_ = cancel_token.cancelled() => break,
|
|
_ = tx.closed() => break,
|
|
Some(message) = cmd_rx.recv() => {
|
|
match message {
|
|
None => break,
|
|
Some(message) => {
|
|
let key = message.0.uuid.clone().unwrap().value;
|
|
in_progress.insert(key.clone(), message.1);
|
|
match tx.send(Ok(message.0)).await {
|
|
Ok(()) => {},
|
|
Err(e) => {
|
|
error!("Failed to send command data: {e}");
|
|
if in_progress.remove(&key).unwrap().send(CommandResponse {
|
|
uuid: Some(Uuid::from(key)),
|
|
success: false,
|
|
response: "Failed to send command data.".to_string(),
|
|
}).is_err() {
|
|
error!("Failed to send command response on failure to send command data");
|
|
}
|
|
break;
|
|
}
|
|
}
|
|
}
|
|
}
|
|
},
|
|
Some(message) = in_stream.next() => {
|
|
match message {
|
|
Ok(message) => {
|
|
match message.inner {
|
|
Some(Inner::Response(response)) => {
|
|
if let Some(uuid) = &response.uuid {
|
|
match in_progress.remove(&uuid.value) {
|
|
Some(sender) => {
|
|
if sender.send(response).is_err() {
|
|
error!("Failed to send command response on success")
|
|
}
|
|
}
|
|
None => {
|
|
result = Status::invalid_argument("Invalid Command UUID");
|
|
break;
|
|
}
|
|
}
|
|
}
|
|
}
|
|
_ => {
|
|
result = Status::invalid_argument("Subsequent Message Must Be Command Responses");
|
|
break;
|
|
}
|
|
}
|
|
}
|
|
Err(e) => {
|
|
error!("Received error from command handler {e}");
|
|
break
|
|
},
|
|
}
|
|
}
|
|
else => break,
|
|
}
|
|
}
|
|
cmd_rx.close();
|
|
if !tx.is_closed() {
|
|
match tx.send(Err(result)).await {
|
|
Ok(()) => {}
|
|
Err(e) => {
|
|
error!("Failed to close old command sender {e}");
|
|
}
|
|
}
|
|
}
|
|
for (key, sender) in in_progress.drain() {
|
|
if sender
|
|
.send(CommandResponse {
|
|
uuid: Some(Uuid::from(key)),
|
|
success: false,
|
|
response: "Command Handler Shut Down".to_string(),
|
|
})
|
|
.is_err()
|
|
{
|
|
error!("Failed to send command response on shutdown");
|
|
}
|
|
}
|
|
});
|
|
|
|
Ok(Response::new(Box::pin(ReceiverStream::new(rx))))
|
|
}
|
|
}
|