Safely Terminate Aggregator

include/core/mission_state.hpp

@@ -98,6 +98,15 @@ class MissionState {
     std::shared_ptr<CVAggregator> getCV();
     void setCV(std::shared_ptr<CVAggregator> cv);
 
+    enum class CVStatus {
+        None = 0,
+        Validated = 1,
+        Rejected = 2,
+    };
+
+    CVStatus getCVStatus();
+    void setCVStatus(CVStatus status);
+
     /*
      * Gets a shared_ptr to the camera client, which lets you
      * take photos of ground targets.
@@ -144,6 +153,10 @@ class MissionState {
     std::shared_ptr<MavlinkClient> mav;
     std::shared_ptr<AirdropClient> airdrop;
     std::shared_ptr<CVAggregator> cv;
+
+    std::mutex cv_status_mut;
+    CVStatus cv_status = CVStatus::None;
+
     std::shared_ptr<CameraInterface> camera;
 
     std::mutex cv_mut;

include/cv/aggregator.hpp

@@ -5,13 +5,13 @@
 #include <cmath>
 #include <functional>
 #include <future>
+#include <map>
 #include <memory>
 #include <mutex>
 #include <queue>
 #include <thread>
 #include <unordered_map>
 #include <vector>
-#include <map>
 
 #include "cv/pipeline.hpp"
 #include "cv/utilities.hpp"
@@ -42,6 +42,9 @@ class CVAggregator {
     // Spawn a thread to run the pipeline on the given imageData
     void runPipeline(const ImageData& image);
 
+    // Stop accepting work, discard queued images, and wait for active workers to finish
+    void terminate();
+
     // Lockable pointer to retrieve aggregator results
     LockPtr<CVResults> getResults();
 
@@ -66,7 +69,9 @@ class CVAggregator {
     Pipeline pipeline;
 
     std::mutex mut;
-    int num_worker_threads;
+    std::atomic<int> num_worker_threads;
+    std::atomic<bool> accepting_images;
+    std::vector<std::thread> worker_threads;
 
     // For when too many pipelines are active at once
     std::queue<ImageData> overflow_queue;

include/ticks/cv_loiter.hpp

@@ -22,24 +22,10 @@
 
 class CVLoiterTick : public Tick {
  public:
-    /* 
-     * Status of the CV processing
-     * (Will be) Modified through GCS Manual Control
-    */
-    enum class Status {
-        None = 0,
-        Validated = 1,
-        Rejected = 2,
-    };
-
     explicit CVLoiterTick(std::shared_ptr<MissionState> state);
     std::chrono::milliseconds getWait() const override;
 
-    void setStatus(Status status);
     Tick* tick() override;
-
- private:
-    Status status;
 };
 
 #endif  // INCLUDE_TICKS_CV_LOITER_HPP_

src/core/mission_state.cpp

@@ -121,6 +121,16 @@ std::shared_ptr<CVAggregator> MissionState::getCV() { return this->cv; }
 
 void MissionState::setCV(std::shared_ptr<CVAggregator> cv) { this->cv = cv; }
 
+MissionState::CVStatus MissionState::getCVStatus() {
+    Lock lock(this->cv_status_mut);
+    return this->cv_status;
+}
+
+void MissionState::setCVStatus(CVStatus status) {
+    Lock lock(this->cv_status_mut);
+    this->cv_status = status;
+}
+
 std::shared_ptr<CameraInterface> MissionState::getCamera() { return this->camera; }
 
 void MissionState::setCamera(std::shared_ptr<CameraInterface> camera) { this->camera = camera; }

src/cv/aggregator.cpp

@@ -1,12 +1,15 @@
 #include "cv/aggregator.hpp"
 
+#include <exception>
+
 #include "utilities/constants.hpp"
 #include "utilities/lockptr.hpp"
 #include "utilities/locks.hpp"
 #include "utilities/logging.hpp"
 
 CVAggregator::CVAggregator(Pipeline&& p) : pipeline(std::move(p)) {
-    this->num_worker_threads = 0;
+    this->num_worker_threads.store(0);
+    this->accepting_images.store(true);
     this->results = std::make_shared<CVResults>();
     this->matched_results = std::make_shared<MatchedResults>();
     this->cv_record = std::make_shared<std::map<int, IdentifiedTarget>>();
@@ -26,7 +29,7 @@ CVAggregator::CVAggregator(Pipeline&& p) : pipeline(std::move(p)) {
     this->matched_results->matched_airdrop[AirdropType::Beacon] = dummy;
 }
 
-CVAggregator::~CVAggregator() {}
+CVAggregator::~CVAggregator() { this->terminate(); }
 
 LockPtr<CVResults> CVAggregator::getResults() {
     return LockPtr<CVResults>(this->results, &this->mut);
@@ -52,19 +55,54 @@ void CVAggregator::updateRecords(std::vector<IdentifiedTarget>& new_values) {
 void CVAggregator::runPipeline(const ImageData& image) {
     Lock lock(this->mut);
 
-    if (this->num_worker_threads >= MAX_CV_PIPELINES) {
+    if (!this->accepting_images.load()) {
+        LOG_F(WARNING, "CVAggregator is not accepting new images. Dropping pipeline run.");
+        return;
+    }
+
+    const int active_workers = this->num_worker_threads.load();
+    if (active_workers >= MAX_CV_PIPELINES) {
         // If we have too many running workers, just queue the new image
         LOG_F(WARNING, "Too many CVAggregator workers (%d). Pushing to overflow queue...",
-              this->num_worker_threads);
+              active_workers);
         this->overflow_queue.push(image);
         LOG_F(WARNING, "Overflow queue size is now %ld", this->overflow_queue.size());
         return;
     }
 
     static int thread_counter = 0;
-    ++this->num_worker_threads;
-    std::thread worker_thread(&CVAggregator::worker, this, image, ++thread_counter);
-    worker_thread.detach();  // We don’t need to join in the caller
+    this->num_worker_threads.fetch_add(1);
+    try {
+        this->worker_threads.emplace_back(&CVAggregator::worker, this, image, ++thread_counter);
+    } catch (const std::exception& err) {
+        this->num_worker_threads.fetch_sub(1);
+        LOG_F(ERROR, "Failed to spawn CVAggregator worker: %s", err.what());
+    }
+}
+
+void CVAggregator::terminate() {
+    std::vector<std::thread> threads;
+    {
+        Lock lock(this->mut);
+        this->accepting_images.store(false);
+        const std::size_t dropped_images = this->overflow_queue.size();
+        std::queue<ImageData> empty_queue;
+        this->overflow_queue.swap(empty_queue);
+
+        if (this->worker_threads.empty()) {
+            return;
+        }
+
+        LOG_F(INFO, "Dropped %zu queued images. Waiting for %zu CVAggregator worker threads.",
+              dropped_images, this->worker_threads.size());
+        threads.swap(this->worker_threads);
+    }
+
+    for (std::thread& worker_thread : threads) {
+        if (worker_thread.joinable()) {
+            worker_thread.join();
+        }
+    }
 }
 
 static std::atomic<int> global_run_id{0};
@@ -107,10 +145,9 @@ void CVAggregator::worker(ImageData image, int thread_num) {
 
     // 4) Mark ourselves as finished
     {
-        Lock lock(this->mut);
+        const int active_workers = this->num_worker_threads.fetch_sub(1);
         LOG_F(INFO, "CVAggregator worker #%d terminating. Active threads: %d -> %d", thread_num,
-              this->num_worker_threads, this->num_worker_threads - 1);
-        --this->num_worker_threads;
+              active_workers, active_workers - 1);
     }
 }
 

src/network/gcs_routes.cpp

@@ -7,7 +7,6 @@
 #include <string>
 #include <vector>
 
-
 #include <nlohmann/json.hpp>
 
 #include "core/mission_state.hpp"
@@ -16,7 +15,6 @@
 #include "pathing/mission_path.hpp"
 #include "protos/obc.pb.h"
 #include "ticks/airdrop_approach.hpp"
-#include "ticks/cv_loiter.hpp"
 #include "ticks/path_gen.hpp"
 #include "ticks/path_validate.hpp"
 #include "ticks/tick.hpp"
@@ -414,29 +412,29 @@ DEF_GCS_HANDLE(Post, targets, matched) {
 
     LOG_S(INFO) << j_root;
 
-    LockPtr<MatchedResults> matched_results = state->getCV()->getMatchedResults();
+    state->getCV()->terminate();
 
-    if (matched_results.data == nullptr) {
-        LOG_S(ERROR) << "lockptr is null";
-    }
+    {
+        LockPtr<MatchedResults> matched_results = state->getCV()->getMatchedResults();
 
-    AirdropTarget returned_matched_result;
+        if (matched_results.data == nullptr) {
+            LOG_S(ERROR) << "lockptr is null";
+            return;
+        }
 
-    for (const auto& instance : j_root) {
-        LOG_S(INFO) << returned_matched_result.index();
-        google::protobuf::util::JsonStringToMessage(instance.dump(), &returned_matched_result);
-        LOG_S(WARNING) << returned_matched_result.index();
-        matched_results.data->matched_airdrop[returned_matched_result.index()] =
-            returned_matched_result;
-        LOG_S(ERROR) << returned_matched_result.index();
-    }
+        AirdropTarget returned_matched_result;
 
-    auto lock_ptr = state->getTickLockPtr<CVLoiterTick>();
-    if (!lock_ptr.has_value()) {
-        LOG_RESPONSE(WARNING, "Not currently in Loiter Tick", BAD_REQUEST);
-        return;
+        for (const auto& instance : j_root) {
+            LOG_S(INFO) << returned_matched_result.index();
+            google::protobuf::util::JsonStringToMessage(instance.dump(), &returned_matched_result);
+            LOG_S(WARNING) << returned_matched_result.index();
+            matched_results.data->matched_airdrop[returned_matched_result.index()] =
+                returned_matched_result;
+            LOG_S(ERROR) << returned_matched_result.index();
+        }
     }
-    lock_ptr->data->setStatus(CVLoiterTick::Status::Validated);
+
+    state->setCVStatus(MissionState::CVStatus::Validated);
 
     LOG_RESPONSE(INFO, "Finished setting targets (and thus loitering)", OK);
 }

src/ticks/cv_loiter.cpp

@@ -7,14 +7,10 @@
 #include "ticks/ids.hpp"
 #include "utilities/constants.hpp"
 
-CVLoiterTick::CVLoiterTick(std::shared_ptr<MissionState> state) : Tick(state, TickID::CVLoiter) {
-    this->status = CVLoiterTick::Status::None;
-}
+CVLoiterTick::CVLoiterTick(std::shared_ptr<MissionState> state) : Tick(state, TickID::CVLoiter) {}
 
 std::chrono::milliseconds CVLoiterTick::getWait() const { return CV_LOITER_TICK_WAIT; }
 
-void CVLoiterTick::setStatus(Status status) { this->status = status; }
-
 Tick* CVLoiterTick::tick() {
     // Tick is called if Search Zone coverage path is finished
 
@@ -32,8 +28,10 @@ Tick* CVLoiterTick::tick() {
     //     }
     // }
 
+    MissionState::CVStatus status = this->state->getCVStatus();
+
     // Check status of the CV Results
-    if (status == Status::Validated) {
+    if (status == MissionState::CVStatus::Validated) {
         /*
         const std::array<AirdropType, NUM_AIRDROPS> ALL_AIRDROPS = {
             AirdropType::Water, AirdropType::Beacon,
@@ -74,7 +72,7 @@ Tick* CVLoiterTick::tick() {
         // }
 
         return new AirdropPrepTick(this->state);
-    } else if (status == Status::Rejected) {
+    } else if (status == MissionState::CVStatus::Rejected) {
         // TODO: Tell Mav to restart Search Mission
         return new FlySearchTick(this->state);
     }

src/ticks/fly_search.cpp

@@ -3,10 +3,11 @@
 #include <memory>
 #include <chrono>
 
+#include "pathing/environment.hpp"
+#include "ticks/airdrop_prep.hpp"
+#include "ticks/cv_loiter.hpp"
 #include "ticks/ids.hpp"
 #include "utilities/common.hpp"
-#include "ticks/cv_loiter.hpp"
-#include "pathing/environment.hpp"
 
 using namespace std::chrono_literals; // NOLINT
 
@@ -44,6 +45,12 @@ void FlySearchTick::init() {
 }
 
 Tick* FlySearchTick::tick() {
+    MissionState::CVStatus status = this->state->getCVStatus();
+    if (status == MissionState::CVStatus::Validated) {
+        this->state->setCVStatus(MissionState::CVStatus::None);
+        return new AirdropPrepTick(this->state);
+    }
+
     if (!this->mission_started) {
         this->mission_started = this->state->getMav()->startMission();
         return nullptr;