diff --git a/image_analysis/indexing/IndexerThreadPool.cpp b/image_analysis/indexing/IndexerThreadPool.cpp index 3d57769f..7da37f55 100644 --- a/image_analysis/indexing/IndexerThreadPool.cpp +++ b/image_analysis/indexing/IndexerThreadPool.cpp @@ -180,13 +180,33 @@ IndexerThread::~IndexerThread() { Finalize(); } -IndexerThreadPool::IndexerThreadPool(const IndexingSettings &settings) +IndexerThreadPool::IndexerThreadPool(const IndexingSettings &settings, bool blocking) : worker_busy(settings.GetIndexingThreads(), 0), - viable_cell_min_spots(settings.GetViableCellMinSpots()) { + worker_free_count(settings.GetIndexingThreads()), + viable_cell_min_spots(settings.GetViableCellMinSpots()), + blocking(blocking) { for (size_t i = 0; i < settings.GetIndexingThreads(); ++i) tasks.emplace_back(std::make_unique(std::cref(settings), i)); } +int IndexerThreadPool::GetFreeWorker() { + std::unique_lock lock(m); + + if (tasks.size() == 0) + return -1; + + if (blocking) + c.wait(lock, [this] { return worker_free_count > 0; }); + + for (int i = 0; i < tasks.size(); i++) { + if (worker_busy[i] == 0) { + worker_busy[i] = 1; + worker_free_count--; + return i; + } + } + return -1; +} IndexerResult IndexerThreadPool::Run(const DiffractionExperiment &experiment, const std::vector &recip) { if (experiment.GetIndexingAlgorithm() == IndexingAlgorithmEnum::None) @@ -195,25 +215,22 @@ IndexerResult IndexerThreadPool::Run(const DiffractionExperiment &experiment, co return IndexerResult{.lattice = {}, .indexing_time_s = 0}; // Check if there is available worker - int task = -1; - { - std::unique_lock lock(m); - for (int i = 0; i < tasks.size(); i++) { - if (worker_busy[i] == 0) { - task = i; - worker_busy[i] = 1; - break; - } - } - } + const int task = GetFreeWorker(); std::unique_ptr result; if (task >= 0) { - result = tasks[task]->Run(experiment, recip); + try { + result = tasks[task]->Run(experiment, recip); + } catch (const std::exception &e) { + spdlog::error("Indexer thread failed: {}", e.what()); + result = nullptr; + } { std::unique_lock lock(m); worker_busy[task] = 0; + worker_free_count++; } + c.notify_one(); } if (result) return *result; diff --git a/image_analysis/indexing/IndexerThreadPool.h b/image_analysis/indexing/IndexerThreadPool.h index efccc369..7f5fa978 100644 --- a/image_analysis/indexing/IndexerThreadPool.h +++ b/image_analysis/indexing/IndexerThreadPool.h @@ -21,8 +21,6 @@ #include "../common/NUMAHWPolicy.h" #include "Indexer.h" - - class IndexerThread { struct TaskInput { const DiffractionExperiment &experiment; @@ -49,12 +47,15 @@ public: class IndexerThreadPool { std::mutex m; + std::condition_variable c; std::vector worker_busy; + size_t worker_free_count; std::vector> tasks; const int64_t viable_cell_min_spots; - + const bool blocking; + int GetFreeWorker(); public: - IndexerThreadPool(const IndexingSettings& settings); + IndexerThreadPool(const IndexingSettings& settings, bool blocking = true); IndexerResult Run(const DiffractionExperiment& experiment, const std::vector& recip); }; diff --git a/receiver/JFJochReceiverService.cpp b/receiver/JFJochReceiverService.cpp index e44655f3..a320cbe3 100644 --- a/receiver/JFJochReceiverService.cpp +++ b/receiver/JFJochReceiverService.cpp @@ -329,7 +329,7 @@ JFJochReceiverService &JFJochReceiverService::Indexing(const IndexingSettings &i if (input.GetAlgorithm() != IndexingAlgorithmEnum::None) { logger.Info("Creating indexing thread pool..."); - indexer_thread_pool = std::make_unique(input); + indexer_thread_pool = std::make_unique(input, false); logger.Info(" ... done"); } return *this;