update

Author: namchuai

engine/common/download_event.h

@@ -0,0 +1,55 @@
+#pragma once
+
+#include "common/download_task.h"
+#include "common/event.h"
+
+namespace cortex::event {
+
+enum class DownloadEventType {
+  DownloadValidateStarted,
+  DownloadValidateFailed,
+  DownloadStarted,
+  DownloadPaused,
+  DownloadUpdated,
+  DownloadSuccess,
+  DownloadError,
+};
+
+namespace {
+std::string DownloadEventTypeToString(DownloadEventType type) {
+  switch (type) {
+    case DownloadEventType::DownloadValidateStarted:
+      return "DownloadValidateStarted";
+    case DownloadEventType::DownloadValidateFailed:
+      return "DownloadValidateFailed";
+    case DownloadEventType::DownloadStarted:
+      return "DownloadStarted";
+    case DownloadEventType::DownloadPaused:
+      return "DownloadPaused";
+    case DownloadEventType::DownloadUpdated:
+      return "DownloadUpdated";
+    case DownloadEventType::DownloadSuccess:
+      return "DownloadSuccess";
+    case DownloadEventType::DownloadError:
+      return "DownloadError";
+    default:
+      return "Unknown";
+  }
+}
+}  // namespace
+
+class DownloadEvent : public cortex::event::Event {
+ public:
+  explicit DownloadEvent(DownloadEventType type, const DownloadTask task)
+      : type_{type}, download_task_{task} {}
+
+  std::string ToString() const {
+    return "DownloadEvent: " + DownloadEventTypeToString(type_) +
+           " task: " + download_task_.ToString();
+  }
+
+ private:
+  DownloadEventType type_;
+  DownloadTask download_task_;
+};
+};  // namespace cortex::event

engine/common/download_task.h

@@ -0,0 +1,51 @@
+#pragma once
+
+#include <filesystem>
+#include <sstream>
+#include <string>
+
+enum class DownloadType { Model, Engine, Miscellaneous, CudaToolkit, Cortex };
+
+struct DownloadItem {
+  std::string id;
+
+  std::string downloadUrl;
+
+  /**
+   * An absolute path to where the file is located (locally).
+   */
+  std::filesystem::path localPath;
+
+  std::optional<std::string> checksum;
+
+  std::optional<uint64_t> bytes;
+
+  std::string ToString() const {
+    std::ostringstream output;
+    output << "DownloadItem{id: " << id << ", downloadUrl: " << downloadUrl
+           << ", localContainerPath: " << localPath
+           << ", checksum: " << checksum.value_or("N/A") << "}";
+    return output.str();
+  }
+};
+
+struct DownloadTask {
+  std::string id;
+
+  DownloadType type;
+
+  std::vector<DownloadItem> items;
+
+  std::string ToString() const {
+    std::ostringstream output;
+    output << "DownloadTask{id: " << id << ", type: " << static_cast<int>(type)
+           << ", items: [";
+    for (const auto& item : items) {
+      output << item.ToString() << ", ";
+    }
+    output << "]}";
+    return output.str();
+  }
+
+  // TODO: namh implement ToJson()
+};

engine/common/event.h

@@ -1,10 +1,9 @@
 #pragma once
 
-#include <string>
-
 namespace cortex::event {
 
-struct Event {
-  std::string message;
+class Event {
+ public:
+  virtual ~Event() {}
 };
 }  // namespace cortex::event

engine/common/hardward_event.h

@@ -0,0 +1,10 @@
+#pragma once
+
+#include "common/event.h"
+
+namespace cortex::event {
+class HardwareEvent : public cortex::event::Event {
+ public:
+  explicit HardwareEvent() {}
+};
+};  // namespace cortex::event

engine/common/model_event.h

@@ -0,0 +1,10 @@
+#pragma once
+
+#include "common/event.h"
+
+namespace cortex::event {
+class ModelEvent : public cortex::event::Event {
+ public:
+  explicit ModelEvent() {}
+};
+};  // namespace cortex::event

engine/controllers/events.cc

@@ -8,15 +8,15 @@ void Events::broadcast(const std::string& message) {
 
 void Events::handleNewMessage(const WebSocketConnectionPtr& wsConnPtr,
                               std::string&& message,
-                              const WebSocketMessageType& type) {}
+                              const WebSocketMessageType& type) {
+  // ignore message sent from client
+}
 
 void Events::handleNewConnection(const HttpRequestPtr& req,
-                                 const WebSocketConnectionPtr& wsConnPtr) {
-  LOG_DEBUG << "New connection";
-  connections_.insert(wsConnPtr);
+                                 const WebSocketConnectionPtr& ws_conn_ptr) {
+  connections_.insert(ws_conn_ptr);
 }
 
-void Events::handleConnectionClosed(const WebSocketConnectionPtr& wsConnPtr) {
-  LOG_DEBUG << "Connection closed";
-  connections_.erase(wsConnPtr);
+void Events::handleConnectionClosed(const WebSocketConnectionPtr& ws_conn_ptr) {
+  connections_.erase(ws_conn_ptr);
 }

engine/controllers/events.h

@@ -3,12 +3,14 @@
 #include <drogon/PubSubService.h>
 #include <drogon/WebSocketController.h>
 #include <eventpp/eventqueue.h>
+#include "common/download_event.h"
 #include "common/event.h"
 
 using namespace drogon;
 
 using Event = cortex::event::Event;
-using EventQueue = eventpp::EventQueue<std::string, void(const Event&)>;
+using DownloadEvent = cortex::event::DownloadEvent;
+using EventQueue = eventpp::EventQueue<std::string, void(DownloadEvent)>;
 
 class Events : public drogon::WebSocketController<Events, false> {
 
@@ -19,8 +21,9 @@ class Events : public drogon::WebSocketController<Events, false> {
 
   explicit Events(std::shared_ptr<EventQueue> event_queue)
       : event_queue_{event_queue} {
-    event_queue_->appendListener("download-started", [this](const Event& e) {
-      this->broadcast(e.message);
+    // TODO: namh make a list of event
+    event_queue_->appendListener("download-update", [this](DownloadEvent e) {
+      this->broadcast(e.ToString());
     });
   };
 

engine/main.cc

@@ -10,6 +10,7 @@
 #include "utils/archive_utils.h"
 #include "utils/cortex_utils.h"
 #include "utils/dylib.h"
+#include "utils/event_processor.h"
 #include "utils/file_logger.h"
 #include "utils/file_manager_utils.h"
 #include "utils/logging_utils.h"
@@ -73,21 +74,22 @@ void RunServer() {
 #endif
 
   using Event = cortex::event::Event;
-  using EventQueue = eventpp::EventQueue<std::string, void(const Event&)>;
+  using EventQueue = eventpp::EventQueue<std::string, void(DownloadEvent)>;
 
   auto event_queue_ptr = std::make_shared<EventQueue>();
+  cortex::event::EventProcessor event_processor(event_queue_ptr);
+
   auto download_service = std::make_shared<DownloadService>(event_queue_ptr);
+  auto engine_service = std::make_shared<EngineService>(download_service);
+  auto model_service = std::make_shared<ModelService>(download_service);
 
   // initialize custom controllers
-  auto engine_service = std::make_shared<EngineService>(download_service);
   auto engine_ctl = std::make_shared<Engines>(engine_service);
-  drogon::app().registerController(engine_ctl);
-
-  auto model_service = std::make_shared<ModelService>(download_service);
   auto model_ctl = std::make_shared<Models>(model_service);
-  drogon::app().registerController(model_ctl);
-
   auto event_ctl = std::make_shared<Events>(event_queue_ptr);
+
+  drogon::app().registerController(engine_ctl);
+  drogon::app().registerController(model_ctl);
   drogon::app().registerController(event_ctl);
 
   LOG_INFO << "Server started, listening at: " << config.apiServerHost << ":"
@@ -176,4 +178,4 @@ int main(int argc, char* argv[]) {
   CommandLineParser clp;
   clp.SetupCommand(argc, argv);
   return 0;
-}
+}

engine/services/download_service.cc

@@ -115,39 +115,48 @@ cpp::result<uint64_t, std::string> DownloadService::GetFileSize(
 cpp::result<bool, std::string> DownloadService::AddAsyncDownloadTask(
     DownloadTask& task,
     std::optional<OnDownloadTaskSuccessfully> callback) noexcept {
-  auto verifying_result = VerifyDownloadTask(task);
-  if (verifying_result.has_error()) {
-    return cpp::fail(verifying_result.error());
+
+  if (std::find(download_task_list_.begin(), download_task_list_.end(),
+                task.id) != download_task_list_.end()) {
+    return cpp::fail("Download task already exists: " + task.id);
+  }
+
+  download_task_list_.push_back(task.id);
+  download_task_map_.insert({task.id, task});
+
+  {
+    // verify download task
+    auto result = VerifyDownloadTask(task);
+    if (result.has_error()) {
+      CleanUp(task.id);
+      return cpp::fail(result.error());
+    }
   }
 
   auto execute_download_async = [&, task, callback]() {
-    std::optional<std::string> dl_err_msg = std::nullopt;
+    active_download_task_id_ = task.id;
+    std::optional<std::string> err_msg = std::nullopt;
     for (const auto& item : task.items) {
+      active_download_item_id_ = item.id;
       CTL_INF("Start downloading: " + item.localPath.filename().string());
-
-      if (event_queue_.has_value()) {
-        event_queue_.value()->enqueue("download-started",
-                                      Event{.message = task.ToString()});
-        event_queue_.value()->process();
-        CLI_LOG("Enqueued download-started event: " << task.ToString());
-      }
-
       auto result = Download(task.id, item, false);
       if (result.has_error()) {
-        dl_err_msg = result.error();
+        err_msg = result.error();
         break;
       }
     }
 
-    if (dl_err_msg.has_value()) {
-      CTL_ERR(dl_err_msg.value());
+    if (err_msg.has_value()) {
+      CTL_ERR(err_msg.value());
+      CleanUp(task.id);
       return;
     }
 
     if (callback.has_value()) {
       CTL_INF("Download success, executing post download lambda!");
       callback.value()(task);
     }
+    CleanUp(task.id);
   };
 
   std::thread t(execute_download_async);
@@ -223,6 +232,9 @@ cpp::result<bool, std::string> DownloadService::Download(
   curl_easy_setopt(curl, CURLOPT_NOPROGRESS, 0L);
   curl_easy_setopt(curl, CURLOPT_FOLLOWLOCATION, 1L);
 
+  curl_easy_setopt(curl, CURLOPT_PROGRESSFUNCTION, progressCallback);
+  curl_easy_setopt(curl, CURLOPT_PROGRESSDATA, this);
+
   if (mode == "ab") {
     auto local_file_size = GetLocalFileSize(download_item.localPath);
     if (local_file_size != -1) {
@@ -260,3 +272,17 @@ curl_off_t DownloadService::GetLocalFileSize(
   fclose(file);
   return file_size;
 }
+
+void DownloadService::CleanUp(const std::string& task_id) {
+  CTL_INF("Cleaning up download task: " << task_id);
+  // TODO: might need to be wrap with mutex
+  // remove from task list
+  download_task_list_.erase(std::remove(download_task_list_.begin(),
+                                        download_task_list_.end(), task_id),
+                            download_task_list_.end());
+  // remove from task map
+  download_task_map_.erase(task_id);
+
+  active_download_task_id_ = std::nullopt;
+  active_download_item_id_ = std::nullopt;
+}