Merge pull request #1657 from luxonis/feature/Snaps_cache_upload_during_runtime

aljazdu · web-flow · commit 6451ee2ff3d7 · 2026-02-02T14:50:27.000+01:00
Cached snaps and events are now gradually uploaded during runtime, mi…
diff --git a/include/depthai/utility/EventsManager.hpp b/include/depthai/utility/EventsManager.hpp
@@ -193,9 +193,10 @@ class EventsManager {
 
     /**
      * Fetch the configuration limits and quotas for snaps & events
+     * @param retryOnFail Retry fetching on failure; when true, keeps retrying until successful
      * @return bool
      */
-    bool fetchConfigurationLimits();
+    bool fetchConfigurationLimits(const bool retryOnFail);
     /**
      * Prepare a batch of file groups from inputSnapBatch
      */
@@ -227,10 +228,13 @@ class EventsManager {
      */
     void cacheSnapData(std::deque<std::shared_ptr<SnapData>>& inputSnapBatch);
     /**
-     * Upload cached data to the events service
-     * @return void
+     * Upload cached events to the events service
+     */
+    void uploadCachedEvents();
+    /**
+     * Upload cached snaps to the events service
      */
-    void uploadCachedData();
+    void uploadCachedSnaps();
     /**
      * Check if there's any cached data in the filesystem
      */
@@ -259,6 +263,7 @@ class EventsManager {
     std::mutex stopThreadConditionMutex;
     std::atomic<bool> stopUploadThread;
     std::atomic<bool> configurationLimitsFetched;
+    std::atomic<bool> connectionEstablished;
     std::condition_variable eventBufferCondition;
 
     uint64_t maxFileSizeBytes;
diff --git a/src/utility/EventsManager.cpp b/src/utility/EventsManager.cpp
@@ -17,6 +17,41 @@ std::string generateLocalID(int64_t timestamp) {
     oss << timestamp << "_" << std::setw(6) << std::setfill('0') << counter++;
     return oss.str();
 }
+
+const std::unordered_map<std::string, std::string> extensionToMimeTypeMap = {
+    {".html", "text/html"},
+    {".css", "text/css"},
+    {".js", "application/javascript"},
+    {".png", "image/png"},
+    {".jpg", "image/jpeg"},
+    {".jpeg", "image/jpeg"},
+    {".gif", "image/gif"},
+    {".webp", "image/webp"},
+    {".bmp", "image/bmp"},
+    {".tiff", "image/tiff"},
+    {".svg", "image/svg+xml"},
+    {".json", "application/json"},
+    {".txt", "text/plain"},
+    {".annotation", "application/x-protobuf; proto=SnapAnnotation"},
+    {"", "application/octet-stream"},
+};
+
+const std::unordered_map<std::string, std::string> mimeTypeToExtensionMap = {
+    {"text/html", ".html"},
+    {"text/css", ".css"},
+    {"application/javascript", ".js"},
+    {"image/png", ".png"},
+    {"image/jpeg", ".jpg"},
+    {"image/gif", ".gif"},
+    {"image/webp", ".webp"},
+    {"image/bmp", ".bmp"},
+    {"image/tiff", ".tiff"},
+    {"image/svg+xml", ".svg"},
+    {"application/json", ".json"},
+    {"text/plain", ".txt"},
+    {"application/x-protobuf; proto=SnapAnnotation", ".annotation"},
+    {"application/octet-stream", ""},
+};
 }  // namespace
 
 #include "Environment.hpp"
@@ -131,17 +166,6 @@ FileData::FileData(std::string data, std::string fileName, std::string mimeType)
       classification(proto::event::PrepareFileUploadClass::UNKNOWN_FILE) {}
 
 FileData::FileData(std::filesystem::path filePath, std::string fileName) : fileName(std::move(fileName)) {
-    static const std::unordered_map<std::string, std::string> mimeTypeExtensionMap = {{".html", "text/html"},
-                                                                                      {".htm", "text/html"},
-                                                                                      {".css", "text/css"},
-                                                                                      {".js", "application/javascript"},
-                                                                                      {".png", "image/png"},
-                                                                                      {".jpg", "image/jpeg"},
-                                                                                      {".jpeg", "image/jpeg"},
-                                                                                      {".gif", "image/gif"},
-                                                                                      {".svg", "image/svg+xml"},
-                                                                                      {".json", "application/json"},
-                                                                                      {".txt", "text/plain"}};
     // Read the data
     std::ifstream fileStream(filePath, std::ios::binary | std::ios::ate);
     if(!fileStream) {
@@ -154,15 +178,17 @@ FileData::FileData(std::filesystem::path filePath, std::string fileName) : fileN
     size = data.size();
     checksum = calculateSHA256Checksum(data);
     // Determine the mime type
-    auto it = mimeTypeExtensionMap.find(filePath.extension().string());
-    if(it != mimeTypeExtensionMap.end()) {
+    auto it = extensionToMimeTypeMap.find(filePath.extension().string());
+    if(it != extensionToMimeTypeMap.end()) {
         mimeType = it->second;
     } else {
         mimeType = "application/octet-stream";
     }
     static const std::unordered_set<std::string> imageMimeTypes = {"image/jpeg", "image/png", "image/webp", "image/bmp", "image/tiff"};
     if(imageMimeTypes.find(mimeType) != imageMimeTypes.end()) {
         classification = proto::event::PrepareFileUploadClass::IMAGE_COLOR;
+    } else if(mimeType == "application/x-protobuf; proto=SnapAnnotation") {
+        classification = proto::event::PrepareFileUploadClass::ANNOTATION;
     } else {
         classification = proto::event::PrepareFileUploadClass::UNKNOWN_FILE;
     }
@@ -239,7 +265,11 @@ bool FileData::toFile(const std::filesystem::path& inputPath) {
         logger::error("Filename is empty");
         return false;
     }
-    std::string extension = mimeType == "image/jpeg" ? ".jpg" : ".txt";
+    std::string extension;
+    auto mimeIt = mimeTypeToExtensionMap.find(mimeType);
+    if(mimeIt != mimeTypeToExtensionMap.end()) {
+        extension = mimeIt->second;
+    }
     // Choose a unique filename
     std::filesystem::path target = inputPath / (fileName + extension);
     for(int i = 1; std::filesystem::exists(target); ++i) {
@@ -284,18 +314,15 @@ EventsManager::EventsManager(bool uploadCachedOnStart)
     // Thread handling preparation and uploads
     uploadThread = std::make_unique<std::thread>([this]() {
         // Fetch configuration limits when starting the new thread
-        configurationLimitsFetched = fetchConfigurationLimits();
-        auto currentTime = std::chrono::steady_clock::now();
-        auto nextTime = currentTime + std::chrono::hours(1);
+        configurationLimitsFetched = fetchConfigurationLimits(true);
         while(!stopUploadThread) {
-            // Hourly check for fetching configuration and limits
-            currentTime = std::chrono::steady_clock::now();
-            if(currentTime >= nextTime) {
-                fetchConfigurationLimits();
-                nextTime += std::chrono::hours(1);
-                if(remainingStorageBytes <= warningStorageBytes) {
-                    logger::warn("Current remaining storage is running low: {} MB", remainingStorageBytes / (1024 * 1024));
-                }
+            connectionEstablished = fetchConfigurationLimits(false);
+            if(remainingStorageBytes <= warningStorageBytes) {
+                logger::warn("Current remaining storage is running low: {} MB", remainingStorageBytes / (1024 * 1024));
+            }
+            // Add cached snaps (if any) to the snapBuffer
+            if(connectionEstablished) {
+                uploadCachedSnaps();
             }
             // Prepare the batch first to reduce contention
             std::deque<std::shared_ptr<SnapData>> snapBatch;
@@ -323,10 +350,8 @@ EventsManager::EventsManager(bool uploadCachedOnStart)
             eventBufferCondition.wait_for(lock, std::chrono::seconds(static_cast<int>(this->publishInterval)), [this]() { return stopUploadThread.load(); });
         }
     });
-    // Upload or clear previously cached data on start
-    if(uploadCachedOnStart) {
-        uploadCachedData();
-    } else {
+    // If upload on start is not set, clear the previously cached data when starting
+    if(!uploadCachedOnStart) {
         clearCachedData(cacheDir);
     }
 }
@@ -339,7 +364,7 @@ EventsManager::~EventsManager() {
     }
 }
 
-bool EventsManager::fetchConfigurationLimits() {
+bool EventsManager::fetchConfigurationLimits(const bool retryOnFail) {
     logger::info("Fetching configuration limits");
     if(token.empty()) {
         logger::warn("Missing token, please set DEPTHAI_HUB_API_KEY environment variable or use the setToken method");
@@ -368,6 +393,10 @@ bool EventsManager::fetchConfigurationLimits() {
                 }));
         if(response.status_code != cpr::status::HTTP_OK) {
             logger::error("Failed to fetch configuration limits, status code: {}", response.status_code);
+            // Fetching should be retried indefinetly only on startup
+            if(!retryOnFail) {
+                return false;
+            }
 
             // Apply exponential backoff
             auto factor = std::pow(uploadRetryPolicy.factor, ++retryAttempt);
@@ -628,6 +657,11 @@ bool EventsManager::uploadFile(std::shared_ptr<FileData> fileData, std::string u
 }
 
 void EventsManager::uploadEventBatch() {
+    // Add cached events, if any
+    if(connectionEstablished) {
+        uploadCachedEvents();
+    }
+
     auto eventBatch = std::make_unique<proto::event::BatchUploadEvents>();
     {
         std::lock_guard<std::mutex> lock(eventBufferMutex);
@@ -915,8 +949,11 @@ void EventsManager::cacheEvents() {
     logger::info("Caching events");
     std::lock_guard<std::mutex> lock(eventBufferMutex);
     for(const auto& eventData : eventBuffer) {
-        std::filesystem::path path(cacheDir);
-        path = path / ("event_" + eventData->event->name() + "_" + std::to_string(eventData->event->created_at()));
+        std::filesystem::path inputPath(cacheDir);
+        std::filesystem::path path = inputPath / (fmt::format("event_{}_{}", eventData->event->name(), eventData->event->created_at()));
+        for(int i = 1; std::filesystem::exists(path); ++i) {
+            path = inputPath / (fmt::format("event_{}_{}_{}", eventData->event->name(), eventData->event->created_at(), i));
+        }
         std::string eventDir = path.string();
         logger::info("Caching event to {}", eventDir);
         if(!std::filesystem::exists(cacheDir)) {
@@ -933,8 +970,12 @@ void EventsManager::cacheSnapData(std::deque<std::shared_ptr<SnapData>>& inputSn
     // Create a unique directory and save the snapData
     logger::info("Caching snapData");
     for(const auto& snap : inputSnapBatch) {
-        std::filesystem::path path(cacheDir);
-        path = path / ("snap_" + snap->eventData->event->name() + "_" + std::to_string(snap->eventData->event->created_at()));
+        std::filesystem::path inputPath(cacheDir);
+        std::filesystem::path path = inputPath / ("snap_" + snap->eventData->event->name() + "_" + std::to_string(snap->eventData->event->created_at()));
+        for(int i = 1; std::filesystem::exists(path); ++i) {
+            path =
+                inputPath / ("snap_" + snap->eventData->event->name() + "_" + std::to_string(snap->eventData->event->created_at()) + "_" + std::to_string(i));
+        }
         std::string snapDir = path.string();
         logger::info("Caching snap to {}", snapDir);
         if(!std::filesystem::exists(cacheDir)) {
@@ -949,13 +990,12 @@ void EventsManager::cacheSnapData(std::deque<std::shared_ptr<SnapData>>& inputSn
     }
 }
 
-void EventsManager::uploadCachedData() {
-    // Iterate over the directories in cacheDir, add events and snapsData to buffers
+void EventsManager::uploadCachedEvents() {
+    // Iterate over the directories in cacheDir, add events to eventBuffer
     if(!checkForCachedData()) {
-        logger::warn("Cache directory is empty, no cached data will be uploaded");
         return;
     }
-    logger::info("Uploading cached data");
+    logger::info("Uploading cached events");
 
     for(const auto& entry : std::filesystem::directory_iterator(cacheDir)) {
         if(!entry.is_directory()) {
@@ -967,13 +1007,34 @@ void EventsManager::uploadCachedData() {
             auto eventData = std::make_unique<EventData>();
             auto event = std::make_shared<proto::event::Event>();
             event->ParseFromIstream(&eventFile);
+
+            // Cached events should be added only if there is enough space in the upcoming request
             std::lock_guard<std::mutex> lock(eventBufferMutex);
+            if(eventBuffer.size() >= eventsPerRequest) {
+                return;
+            }
+
             eventData->event = event;
             eventBuffer.push_back(std::move(eventData));
             // Clear entries added to the eventBuffer
             clearCachedData(entry.path());
+        }
+    }
+}
+
+void EventsManager::uploadCachedSnaps() {
+    // Iterate over the directories in cacheDir, add snaps to snapBuffer
+    if(!checkForCachedData()) {
+        return;
+    }
+    logger::info("Uploading cached snaps");
 
-        } else if(entry.path().filename().string().rfind("snap", 0) == 0) {
+    for(const auto& entry : std::filesystem::directory_iterator(cacheDir)) {
+        if(!entry.is_directory()) {
+            continue;
+        }
+
+        if(entry.path().filename().string().rfind("snap", 0) == 0) {
             std::ifstream snapFile(entry.path() / "snap.pb", std::ios::binary);
             auto snapData = std::make_unique<SnapData>();
             auto eventData = std::make_shared<EventData>();
@@ -982,14 +1043,20 @@ void EventsManager::uploadCachedData() {
             event->ParseFromIstream(&snapFile);
             for(const auto& fileEntry : std::filesystem::directory_iterator(entry.path())) {
                 if(fileEntry.is_regular_file() && fileEntry.path() != entry.path() / "snap.pb") {
-                    auto fileData = std::make_shared<FileData>(fileEntry.path(), fileEntry.path().filename().string());
+                    auto fileData = std::make_shared<FileData>(fileEntry.path(), fileEntry.path().stem().string());
                     fileGroup->fileData.push_back(fileData);
                 }
             }
             snapData->eventData = eventData;
             snapData->eventData->event = event;
             snapData->fileGroup = fileGroup;
+
+            // Cached snaps should be added only if there is enough space in the upcoming request
             std::lock_guard<std::mutex> lock(snapBufferMutex);
+            if(snapBuffer.size() >= maxGroupsPerBatch) {
+                return;
+            }
+
             snapBuffer.push_back(std::move(snapData));
             // Clear entries added to the snapBuffer
             clearCachedData(entry.path());
