feat: save correct resource usage data to the db

dev-sys-do · Aug 31, 2023 · 55879b4 · 55879b4
1 parent 0c5bc6a
commit 55879b4
Show file tree

Hide file tree

Showing 5 changed files with 39 additions and 21 deletions.
diff --git a/controller/src/client.rs b/controller/src/client.rs
@@ -1,5 +1,5 @@
-use orka_proto::scheduler_controller::{scheduling_service_client::SchedulingServiceClient, self};
 use orka_proto::scheduler_controller::SchedulingRequest;
+use orka_proto::scheduler_controller::{self, scheduling_service_client::SchedulingServiceClient};
 use tonic::transport::Channel;
 use tonic::Streaming;
 

diff --git a/controller/src/main.rs b/controller/src/main.rs
@@ -4,16 +4,18 @@ mod routes;
 mod store;
 mod types;
 
-
-use orka_proto::scheduler_controller::{WorkloadInstance, self};
+use orka_proto::scheduler_controller::{self, WorkloadInstance};
 use store::kv_manager::{KeyValueBatch, KeyValueStore, DB_BATCH};
 use tokio::sync::mpsc;
 use tokio_stream::wrappers::ReceiverStream;
 use tonic::{Request, Response, Status};
 
 use axum::Router;
-use orka_proto::scheduler_controller::scheduling_service_server::{SchedulingService};
-use orka_proto::scheduler_controller::{workload_status::Status as DeploymentStatus, 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;
@@ -67,23 +69,35 @@ impl SchedulingService for Scheduler {
                         code: 0,
                         message: Some("The workload is waiting".to_string()),
                     }),
-                    ..Default::default()
+                    resource_usage: Some(Resources {
+                        cpu: 2,
+                        memory: 3,
+                        disk: 4,
+                    }),
                 },
                 WorkloadStatus {
                     instance_id: workload.instance_id.clone(),
                     status: Some(DeploymentStatus {
                         code: 1,
                         message: Some("The workload is running".to_string()),
                     }),
-                    ..Default::default()
+                    resource_usage: Some(Resources {
+                        cpu: 2,
+                        memory: 3,
+                        disk: 4,
+                    }),
                 },
                 WorkloadStatus {
                     instance_id: workload.instance_id,
                     status: Some(DeploymentStatus {
                         code: 2,
                         message: Some("The workload is terminated".to_string()),
                     }),
-                    ..Default::default()
+                    resource_usage: Some(Resources {
+                        cpu: 2,
+                        memory: 3,
+                        disk: 4,
+                    }),
                 },
             ];
 

diff --git a/controller/src/routes/instances.rs b/controller/src/routes/instances.rs
@@ -1,4 +1,3 @@
-use orka_proto::scheduler_controller::{SchedulingRequest, Workload, WorkloadInstance};
 use crate::client::Client;
 use crate::errors::ApiError;
 use crate::store::kv_manager::{KeyValueStore, DB_BATCH};
@@ -7,6 +6,7 @@ use crate::types::instance_status::InstanceStatus;
 use axum::extract::Path;
 use axum::Json;
 use log::{error, trace};
+use orka_proto::scheduler_controller::{SchedulingRequest, Workload, WorkloadInstance};
 use serde_json::{self, json, Value};
 use validator::Validate;
 

diff --git a/controller/src/types/instance_status.rs b/controller/src/types/instance_status.rs
@@ -19,11 +19,7 @@ impl From<&WorkloadStatus> for InstanceStatus {
         InstanceStatus {
             name: (*status.instance_id).to_string(),
             status_code: InstanceStatusCode::from(status.status.clone()),
-            resource_usage: InstanceResources {
-                cpu: 1,
-                memory: 1,
-                disk: 1,
-            },
+            resource_usage: InstanceResources::from(status.resource_usage.clone()),
         }
     }
 }
@@ -37,12 +33,19 @@ pub struct InstanceResources {
     pub disk: i32,
 }
 
-impl From<Resources> for InstanceResources {
-    fn from(res: Resources) -> Self {
-        InstanceResources {
-            cpu: res.cpu,
-            memory: res.memory,
-            disk: res.disk,
+impl From<Option<Resources>> for InstanceResources {
+    fn from(resources: Option<Resources>) -> Self {
+        match resources {
+            Some(res) => InstanceResources {
+                cpu: res.cpu,
+                memory: res.memory,
+                disk: res.disk,
+            },
+            None => InstanceResources {
+                cpu: 0,
+                memory: 0,
+                disk: 0,
+            },
         }
     }
 }

diff --git a/controller/src/types/workload_request.rs b/controller/src/types/workload_request.rs
@@ -3,7 +3,8 @@ use uuid::Uuid;
 use validator::{Validate, ValidationError};
 
 use orka_proto::scheduler_controller::{
-    workload::{Resources, Type}, self,
+    self,
+    workload::{Resources, Type},
 };
 
 #[derive(Debug, Validate, Deserialize, Serialize, Clone)]