Implement Commanding (#6)

Reviewed-on: #6
Co-authored-by: Sergey Savelyev <sergeysav.nn@gmail.com>
Co-committed-by: Sergey Savelyev <sergeysav.nn@gmail.com>

server/src/grpc/cmd.rs

@@ -0,0 +1,150 @@
+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))))
+    }
+}