Feat/kv + routes

.github/workflows/controller.yaml

@@ -10,4 +10,4 @@ jobs:
   rustci:
     uses: "./.github/workflows/rust-ci.yaml"
     with:
-      directory: "controller"
+      directory: "controller"

.github/workflows/rust-ci.yaml

@@ -38,4 +38,4 @@ jobs:
         run: cargo fmt --check
 
       - name: Run lint
-        run: cargo clippy -- -D warnings
+        run: cargo clippy -- -D warnings

controller/.gitignore

@@ -1 +1,2 @@
-target/
+target/
+db/

controller/Cargo.toml

@@ -10,15 +10,21 @@ tokio = { version = "1", features = ["full"] }
 axum = "0.6.19"
 tokio-stream = "0.1.6"
 serde_json = "1.0.104"
-serde = {version = "1.0", features = ["derive"] }
+serde = { version = "1.0", features = ["derive"] }
 validator = { version = "0.16.1", features = ["derive"] }
 anyhow = "1.0.75"
 thiserror = "1.0.47"
 pretty_env_logger = "0.5.0"
 log = "0.4.20"
+kv = { version =  "0.24.0", features = ["json-value"] }
+dotenv = "0.15.0"
+once_cell = "1.18.0"
+orka-proto = { path = "../proto" }
 
-[dependencies.syn]
-version = "2.0.28"
-
-[build-dependencies]
-tonic-build = "0.9"
+[dependencies.uuid]
+version = "1.4.1"
+features = [
+    "v4",                # Lets you generate random UUIDs
+    "fast-rng",          # Use a faster (but still sufficiently random) RNG
+    "macro-diagnostics", # Enable better diagnostics for compile-time UUIDs
+]

controller/README.md

@@ -32,20 +32,16 @@ $ export PATH="$PATH:$HOME/.local/bin"
 The `orka` controller uses the [pretty_env_logger](https://docs.rs/pretty_env_logger/latest/pretty_env_logger/) crate to log all helpful information. The log level can be set by setting the `RUST_LOG` environment variable. For example, to set the log level to `trace` you can run the following command :
 
 ```
-export RUST_LOG=trace
+export RUST_LOG=info
 ```
 
 
 ## Usage
 
 ### Run the server and client GRPC
 
-1. Run the server GRPC :
+1. Run the server the controller :
 ```
-cargo run --bin server
+cargo run --bin orka-controller
 ```
 
-2. Run the client GRPC :
-```
-cargo run --bin client
-```

controller/src/client.rs

@@ -1,11 +1,9 @@
-use scheduler::scheduling_service_client::SchedulingServiceClient;
-use scheduler::SchedulingRequest;
+use orka_proto::scheduler_controller::SchedulingRequest;
+use orka_proto::scheduler_controller::{self, scheduling_service_client::SchedulingServiceClient};
 use tonic::transport::Channel;
-use log::trace;
+use tonic::Streaming;
 
-pub mod scheduler {
-    tonic::include_proto!("orkascheduler");
-}
+use orka_proto::scheduler_controller::{WorkloadInstance, WorkloadStatus};
 
 pub struct Client {
     client: SchedulingServiceClient<Channel>,
@@ -20,16 +18,29 @@ impl Client {
     pub async fn schedule_workload(
         &mut self,
         scheduling_request: SchedulingRequest,
-    ) -> Result<(), Box<dyn std::error::Error>> {
-        let request = scheduling_request;
+    ) -> Result<Streaming<WorkloadStatus>, tonic::Status> {
+        let response = self.client.schedule(scheduling_request).await?;
 
-        let response = self.client.schedule(request).await?;
+        let stream = response.into_inner();
 
-        let mut stream = response.into_inner();
+        Ok(stream)
+    }
 
-        while let Some(status) = stream.message().await? {
-            trace!("STATUS={:?}", status);
-        }
-        Ok(())
+    pub async fn stop_instance(
+        &mut self,
+        instance: WorkloadInstance,
+    ) -> Result<scheduler_controller::Empty, tonic::Status> {
+        let response = self.client.stop(instance).await?;
+
+        Ok(response.into_inner())
     }
-}
+
+    pub async fn destroy_instance(
+        &mut self,
+        instance: WorkloadInstance,
+    ) -> Result<scheduler_controller::Empty, tonic::Status> {
+        let response = self.client.destroy(instance).await?;
+
+        Ok(response.into_inner())
+    }
+}

controller/src/errors.rs

@@ -14,6 +14,15 @@ pub enum ApiError {
 
     #[error("Serialization error")]
     SerializationError(#[from] serde_json::Error),
+
+    #[error("Database error")]
+    DatabaseError(#[from] kv::Error),
+
+    #[error("Scheduling error")]
+    SchedulingError(#[from] tonic::Status),
+
+    #[error("Instance not created")]
+    InstanceNotCreated { message: String },
 }
 
 impl IntoResponse for ApiError {
@@ -24,6 +33,9 @@ impl IntoResponse for ApiError {
             }
             ApiError::ClientConnectError(e) => (StatusCode::INTERNAL_SERVER_ERROR, e.to_string()),
             ApiError::SerializationError(e) => (StatusCode::BAD_REQUEST, e.to_string()),
+            ApiError::DatabaseError(e) => (StatusCode::INTERNAL_SERVER_ERROR, e.to_string()),
+            ApiError::SchedulingError(e) => (StatusCode::INTERNAL_SERVER_ERROR, e.to_string()),
+            ApiError::InstanceNotCreated { message } => (StatusCode::BAD_REQUEST, message),
         };
 
         let payload = json!({

controller/src/lib.rs

@@ -0,0 +1,5 @@
+pub mod client;
+pub mod errors;
+pub mod routes;
+pub mod store;
+pub mod types;

controller/src/main.rs

@@ -1,23 +1,32 @@
 mod client;
 mod errors;
 mod routes;
+mod store;
 mod types;
 
-use crate::client::scheduler;
-
+use orka_proto::scheduler_controller::{self, WorkloadInstance};
+use store::kv_manager::DB_STORE;
+use store::kv_manager::{KeyValueBatch, DB_BATCH};
 use tokio::sync::mpsc;
 use tokio_stream::wrappers::ReceiverStream;
-use tonic::{transport::Server, Request, Response, Status};
+use tonic::{Request, Response, Status};
 
 use axum::Router;
-use scheduler::scheduling_service_server::{SchedulingService, SchedulingServiceServer};
-use scheduler::{SchedulingRequest, WorkloadStatus};
+use orka_proto::scheduler_controller::scheduling_service_server::SchedulingService;
+use orka_proto::scheduler_controller::workload_status::Resources;
+use orka_proto::scheduler_controller::{
+    workload_status::Status as DeploymentStatus, SchedulingRequest, WorkloadStatus,
+};
 use std::net::SocketAddr;
+use std::thread;
+use std::time::Duration;
 use tokio::task;
 
 use axum::routing::{delete, post};
-use log::info;
-use routes::instances::{delete_instance, get_instances, get_specific_instance, post_instance};
+use log::{error, info};
+use routes::instances::{
+    delete_instance, delete_instance_force, get_instances, get_specific_instance, post_instance,
+};
 use routes::workloads::{delete_workload, get_specific_workload, get_workloads, post_workload};
 
 #[derive(Debug, Default)]
@@ -27,33 +36,69 @@ pub struct Scheduler {}
 impl SchedulingService for Scheduler {
     type ScheduleStream = ReceiverStream<Result<WorkloadStatus, Status>>;
 
+    async fn stop(
+        &self,
+        request: Request<WorkloadInstance>,
+    ) -> Result<Response<scheduler_controller::Empty>, Status> {
+        info!("{:?}", request);
+        Ok(Response::new(scheduler_controller::Empty {}))
+    }
+
+    async fn destroy(
+        &self,
+        request: Request<WorkloadInstance>,
+    ) -> Result<Response<scheduler_controller::Empty>, Status> {
+        info!("{:?}", request);
+        Ok(Response::new(scheduler_controller::Empty {}))
+    }
+
     async fn schedule(
         &self,
         request: Request<SchedulingRequest>,
     ) -> Result<Response<Self::ScheduleStream>, Status> {
-        info!("Got a request: {:?}", request);
+        info!("{:?}", request);
 
         let (sender, receiver) = mpsc::channel(4);
 
+        let workload = request.into_inner().workload.unwrap();
+
         tokio::spawn(async move {
             let fake_statuses_response = vec![
                 WorkloadStatus {
-                    name: "Workload 1".to_string(),
-                    status_code: 0,
-                    message: "Your workload is WAITING".to_string(),
-                    ..Default::default()
+                    instance_id: workload.instance_id.clone(),
+                    status: Some(DeploymentStatus {
+                        code: 0,
+                        message: Some("The workload is waiting".to_string()),
+                    }),
+                    resource_usage: Some(Resources {
+                        cpu: 2,
+                        memory: 3,
+                        disk: 4,
+                    }),
                 },
                 WorkloadStatus {
-                    name: "Workload 1".to_string(),
-                    status_code: 1,
-                    message: "Your workload is RUNNING".to_string(),
-                    ..Default::default()
+                    instance_id: workload.instance_id.clone(),
+                    status: Some(DeploymentStatus {
+                        code: 1,
+                        message: Some("The workload is running".to_string()),
+                    }),
+                    resource_usage: Some(Resources {
+                        cpu: 2,
+                        memory: 3,
+                        disk: 4,
+                    }),
                 },
                 WorkloadStatus {
-                    name: "Workload 2".to_string(),
-                    status_code: 2,
-                    message: "Your workload is TERMINATED".to_string(),
-                    ..Default::default()
+                    instance_id: workload.instance_id,
+                    status: Some(DeploymentStatus {
+                        code: 2,
+                        message: Some("The workload is terminated".to_string()),
+                    }),
+                    resource_usage: Some(Resources {
+                        cpu: 2,
+                        memory: 3,
+                        disk: 4,
+                    }),
                 },
             ];
 
@@ -84,18 +129,18 @@ async fn main() -> Result<(), Box<dyn std::error::Error>> {
     pretty_env_logger::init();
 
     // Initialize grpc
-    let grpc_addr = "[::1]:50051".parse()?;
-    let scheduler = Scheduler::default();
-
-    // Spawn the gRPC server as a tokio task
-    let grpc_thread = task::spawn(async move {
-        info!("gRPC server running at: {}", grpc_addr);
-        Server::builder()
-            .add_service(SchedulingServiceServer::new(scheduler))
-            .serve(grpc_addr)
-            .await
-            .unwrap();
-    });
+    // let grpc_addr = "[::1]:50051".parse()?;
+    // let scheduler = Scheduler::default();
+
+    // // Spawn the gRPC server as a tokio task
+    // let grpc_thread = task::spawn(async move {
+    //     info!("gRPC server running at: {}", grpc_addr);
+    //     Server::builder()
+    //         .add_service(SchedulingServiceServer::new(scheduler))
+    //         .serve(grpc_addr)
+    //         .await
+    //         .unwrap();
+    // });
 
     // Initialize http
     let http_addr = SocketAddr::from(([127, 0, 0, 1], 3000));
@@ -109,6 +154,10 @@ async fn main() -> Result<(), Box<dyn std::error::Error>> {
         .route(
             "/instances/:id",
             delete(delete_instance).get(get_specific_instance),
+        )
+        .route(
+            "/instances/:id/force",
+            delete(delete_instance_force).get(get_specific_instance),
         );
 
     // Spawn the HTTP server as a tokio task
@@ -120,8 +169,34 @@ async fn main() -> Result<(), Box<dyn std::error::Error>> {
             .unwrap();
     });
 
-    // Wait for both servers to finish
-    tokio::try_join!(grpc_thread, http_thread)?;
+    // Create a thread to sync the batch with the database
+    let db_thread = task::spawn(async move {
+        loop {
+            thread::sleep(Duration::from_secs(5));
+            let kv_batch = DB_BATCH.lock();
+            let kv_store = DB_STORE.lock();
+            match kv_batch {
+                Ok(mut kvbatch) => {
+                    match kv_store {
+                        Ok(store) => {
+                            store
+                                .instances_bucket()
+                                .unwrap()
+                                .batch(kvbatch.batch.clone())
+                                .unwrap();
+                            // Clear batch after update
+                            *kvbatch = KeyValueBatch::new();
+                        }
+                        Err(e) => error!("{}", e),
+                    }
+                }
+                Err(e) => error!("{}", e),
+            }
+        }
+    });
+
+    // Wait for both servers and a db thread to finish
+    tokio::try_join!(http_thread, db_thread)?;
 
     Ok(())
 }