implement command responses

This commit is contained in:
2025-12-28 16:23:42 -05:00
parent c3253f3204
commit 678b10de08
8 changed files with 242 additions and 57 deletions

View File

@@ -60,10 +60,24 @@ message CommandDefinitionRequest {
}
message Command {
Timestamp timestamp = 1;
map<string, TelemetryValue> parameters = 2;
UUID uuid = 1;
Timestamp timestamp = 2;
map<string, TelemetryValue> parameters = 3;
}
message CommandResponse {
UUID uuid = 1;
bool success = 2;
string response = 3;
}
message ClientSideCommand {
oneof inner {
CommandDefinitionRequest request = 1;
CommandResponse response = 2;
}
}
service CommandService {
rpc NewCommand (CommandDefinitionRequest) returns (stream Command);
rpc NewCommand (stream ClientSideCommand) returns (stream Command);
}

View File

@@ -20,6 +20,10 @@ pub enum Error {
NoCommandReceiver,
#[error("Failed to Send")]
FailedToSend,
#[error("Failed to Receive Command Response")]
FailedToReceiveResponse,
#[error("Command Failure: {0}")]
CommandFailure(String),
}
impl ResponseError for Error {
@@ -31,6 +35,8 @@ impl ResponseError for Error {
Error::WrongParameterType { .. } => StatusCode::BAD_REQUEST,
Error::NoCommandReceiver => StatusCode::SERVICE_UNAVAILABLE,
Error::FailedToSend => StatusCode::INTERNAL_SERVER_ERROR,
Error::FailedToReceiveResponse => StatusCode::INTERNAL_SERVER_ERROR,
Error::CommandFailure(_) => StatusCode::BAD_REQUEST,
}
}
}

View File

@@ -1,25 +1,26 @@
use crate::command::definition::CommandDefinition;
use crate::command::error::Error as CmdError;
use crate::command::error::Error::{
CommandNotFound, FailedToSend, IncorrectParameterCount, MisingParameter, NoCommandReceiver,
WrongParameterType,
CommandFailure, CommandNotFound, FailedToReceiveResponse, FailedToSend,
IncorrectParameterCount, MisingParameter, NoCommandReceiver, WrongParameterType,
};
use crate::core::telemetry_value::Value;
use crate::core::{
Command, CommandDefinitionRequest, TelemetryDataType, TelemetryValue, Timestamp,
Command, CommandDefinitionRequest, CommandResponse, TelemetryDataType, TelemetryValue,
Timestamp, Uuid,
};
use chrono::{DateTime, Utc};
use log::error;
use papaya::HashMap;
use std::collections::HashMap as StdHashMap;
use tokio::sync::mpsc;
use tokio::sync::mpsc::{Receiver, Sender};
use tokio::sync::oneshot;
#[derive(Clone)]
pub(super) struct RegisteredCommand {
pub(super) name: String,
pub(super) definition: CommandDefinitionRequest,
tx: Sender<Option<Command>>,
tx: mpsc::Sender<Option<(Command, oneshot::Sender<CommandResponse>)>>,
}
pub struct CommandManagementService {
@@ -54,7 +55,7 @@ impl CommandManagementService {
pub async fn register_command(
&self,
command: CommandDefinitionRequest,
) -> anyhow::Result<Receiver<Option<Command>>> {
) -> anyhow::Result<mpsc::Receiver<Option<(Command, oneshot::Sender<CommandResponse>)>>> {
let (tx, rx) = mpsc::channel(1);
let registered_commands = self.registered_commands.pin_owned();
@@ -77,7 +78,7 @@ impl CommandManagementService {
&self,
name: impl Into<String>,
parameters: serde_json::Map<String, serde_json::Value>,
) -> Result<(), CmdError> {
) -> Result<String, CmdError> {
let timestamp = Utc::now();
let offset_from_unix_epoch =
timestamp - DateTime::from_timestamp(0, 0).expect("Could not get Unix epoch");
@@ -127,20 +128,38 @@ impl CommandManagementService {
return Err(NoCommandReceiver);
}
let uuid = Uuid::random();
let (response_tx, response_rx) = oneshot::channel();
if let Err(e) = tx
.send(Some(Command {
timestamp: Some(Timestamp {
secs: offset_from_unix_epoch.num_seconds(),
nanos: offset_from_unix_epoch.subsec_nanos(),
}),
parameters: result_parameters,
}))
.send(Some((
Command {
uuid: Some(uuid),
timestamp: Some(Timestamp {
secs: offset_from_unix_epoch.num_seconds(),
nanos: offset_from_unix_epoch.subsec_nanos(),
}),
parameters: result_parameters,
},
response_tx,
)))
.await
{
error!("Failed to Send Command: {e}");
return Err(FailedToSend);
}
Ok(())
response_rx
.await
.map_err(|e| {
error!("Failed to Receive Command Response: {e}");
FailedToReceiveResponse
})
.and_then(|response| {
if response.success {
Ok(response.response)
} else {
Err(CommandFailure(response.response))
}
})
}
}

View File

@@ -1,15 +1,18 @@
use crate::command::service::CommandManagementService;
use crate::core::client_side_command::Inner;
use crate::core::command_service_server::CommandService;
use crate::core::{Command, CommandDefinitionRequest};
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;
use tonic::{Request, Response, Status};
use tonic::codegen::tokio_stream::{Stream, StreamExt};
use tonic::{Request, Response, Status, Streaming};
pub struct CoreCommandService {
pub command_service: Arc<CommandManagementService>,
@@ -22,16 +25,33 @@ impl CommandService for CoreCommandService {
async fn new_command(
&self,
request: Request<CommandDefinitionRequest>,
request: Request<Streaming<ClientSideCommand>>,
) -> Result<Response<Self::NewCommandStream>, Status> {
trace!("CoreCommandService::new_command");
let cancel_token = self.cancellation_token.clone();
let mut cmd_rx = match self
.command_service
.register_command(request.into_inner())
.await
{
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}");
@@ -42,6 +62,8 @@ impl CommandService for CoreCommandService {
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,
@@ -50,32 +72,77 @@ impl CommandService for CoreCommandService {
match message {
None => break,
Some(message) => {
match tx.send(Ok(message)).await {
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(Status::resource_exhausted("End of Command Stream")))
.await
{
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))))

View File

@@ -9,11 +9,11 @@ pub(super) async fn send_command(
name: web::Path<String>,
parameters: web::Json<serde_json::Map<String, serde_json::Value>>,
) -> Result<impl Responder, HttpServerResultError> {
command_service
let result = command_service
.send_command(name.to_string(), parameters.into_inner())
.await?;
Ok(web::Json("Command Sent Successfully."))
Ok(web::Json(result))
}
#[get("/cmd")]