feat: 优化多线程扫描性能和线程池同步

- 改进了多线程特征码扫描的效率，通过在找到结果时提前退出，并返回最小地址。
- 调整了 `scan_dynamic_anchor` 中的 NEON 优化逻辑，以更好地处理边界条件。
- 增强了多线程操作中的内存顺序，以确保更好的内存可见性。
- 优化了线程池的任务调度和等待机制，通过精确跟踪任务数量和使用 `notify_one` 减少不必要的唤醒。

Author: SMlc666

include/ur/signature.hpp

@@ -275,27 +275,38 @@ class runtime_signature : public std::enable_shared_from_this<runtime_signature>
     std::optional<uintptr_t> scan_multithreaded(std::span<const std::byte> memory_range, scanner_func_t core_scanner) const {
         const unsigned int num_threads = std::thread::hardware_concurrency();
         const size_t chunk_size = 65536 * 4;
-        const size_t min_range_for_multithread = chunk_size * num_threads;
+        const size_t min_range_for_multithread = chunk_size;
+
         if (num_threads <= 1 || memory_range.size() < min_range_for_multithread) {
             return (this->*core_scanner)(memory_range, nullptr);
         }
+
         auto& pool = get_pool();
         std::vector<std::future<std::optional<uintptr_t>>> futures;
         auto found_flag = std::make_shared<std::atomic<bool>>(false);
         const size_t overlap = pattern_.size() > 1 ? pattern_.size() - 1 : 0;
+
         for (size_t start = 0; start < memory_range.size(); start += chunk_size) {
+            if (found_flag->load(std::memory_order_acquire)) {
+                break; 
+            }
             size_t end = std::min(start + chunk_size + overlap, memory_range.size());
             if (start >= end || (end - start) < pattern_.size()) continue;
+            
             std::span<const std::byte> chunk = memory_range.subspan(start, end - start);
             futures.push_back(pool.enqueue(core_scanner, shared_from_this(), chunk, found_flag));
         }
+
+        std::optional<uintptr_t> first_result;
         for (auto& fut : futures) {
             if (auto result = fut.get(); result.has_value()) {
-                found_flag->store(true, std::memory_order_relaxed);
-                return result;
+                if (!first_result.has_value() || result.value() < first_result.value()) {
+                    first_result = result;
+                }
             }
         }
-        return std::nullopt;
+        
+        return first_result;
     }
 #endif
 
@@ -416,55 +427,76 @@ namespace detail {
 
 inline std::optional<uintptr_t> runtime_signature::scan_dynamic_anchor(std::span<const std::byte> memory_range, std::shared_ptr<std::atomic<bool>> found_flag) const {
     if (memory_range.size() < pattern_.size()) return std::nullopt;
+
     auto frequencies = detail::calculate_dynamic_rarity(memory_range);
     auto props = detail::find_best_anchor_and_build_props(pattern_, frequencies);
     if (!props.has_anchor) {
         return scan_forward_anchor(memory_range, found_flag);
     }
+
     const uint8x16_t v_anchor = vdupq_n_u8(static_cast<uint8_t>(props.anchor_byte));
     const uint8x16_t v_pattern16 = vld1q_u8(reinterpret_cast<const uint8_t*>(props.pattern16.data()));
     const uint8x16_t v_mask16 = vld1q_u8(reinterpret_cast<const uint8_t*>(props.mask16.data()));
+
     const std::byte* current_pos = memory_range.data();
-    const std::byte* const end_pos = memory_range.data() + memory_range.size() - pattern_.size();
-    const std::byte* const fast_scan_end_pos = memory_range.data() + memory_range.size() - 16;
+    const std::byte* const range_end = memory_range.data() + memory_range.size();
+    const std::byte* const end_pos = range_end - pattern_.size();
+    const std::byte* const fast_scan_end_pos = (range_end >= memory_range.data() + 16) ? range_end - 16 : memory_range.data();
+
     while (current_pos <= fast_scan_end_pos) {
-        if (found_flag && found_flag->load(std::memory_order_relaxed)) return std::nullopt;
+        if (found_flag && found_flag->load(std::memory_order_acquire)) return std::nullopt;
+
 #ifdef UR_ENABLE_HARDWARE_PREFETCH
         __builtin_prefetch(current_pos + 64, 0, 0);
 #endif
         const uint8x16_t v_mem = vld1q_u8(reinterpret_cast<const uint8_t*>(current_pos));
         const uint8x16_t v_cmp_result = vceqq_u8(v_mem, v_anchor);
+
         if (vmaxvq_u8(v_cmp_result) == 0) {
             current_pos += 16;
             continue;
         }
+
         uint8_t result_bytes[16];
         vst1q_u8(result_bytes, v_cmp_result);
         for (int i = 0; i < 16; ++i) {
             if (result_bytes[i] == 0xFF) {
                 const std::byte* potential_start = current_pos + i - props.anchor_offset;
                 if (potential_start < memory_range.data() || potential_start > end_pos) continue;
-                const uint8x16_t v_mem16 = vld1q_u8(reinterpret_cast<const uint8_t*>(potential_start));
-                const uint8x16_t v_masked_mem = vandq_u8(v_mem16, v_mask16);
-                const uint8x16_t v_verify_result = vceqq_u8(v_masked_mem, v_pattern16);
-                if (vminvq_u8(v_verify_result) == 0xFF) {
-                    if (pattern_.size() <= 16 || full_match_at(potential_start)) {
-                        if (found_flag) found_flag->store(true, std::memory_order_relaxed);
+                if (potential_start + pattern_.size() > range_end) continue;
+
+                if (potential_start + 16 <= range_end) {
+                    const uint8x16_t v_mem16 = vld1q_u8(reinterpret_cast<const uint8_t*>(potential_start));
+                    const uint8x16_t v_masked_mem = vandq_u8(v_mem16, v_mask16);
+                    const uint8x16_t v_verify_result = vceqq_u8(v_masked_mem, v_pattern16);
+                    
+                    if (vminvq_u8(v_verify_result) == 0xFF) {
+                        if (pattern_.size() <= 16 || full_match_at(potential_start)) {
+                            if (found_flag) found_flag->store(true, std::memory_order_release);
+                            return reinterpret_cast<uintptr_t>(potential_start);
+                        }
+                    }
+                } else {
+                    if (full_match_at(potential_start)) {
+                        if (found_flag) found_flag->store(true, std::memory_order_release);
                         return reinterpret_cast<uintptr_t>(potential_start);
                     }
                 }
             }
         }
         current_pos += 16;
     }
-    std::span<const std::byte> tail_span{current_pos, static_cast<size_t>(memory_range.data() + memory_range.size() - current_pos)};
-    const auto scan_end_tail = tail_span.data() + tail_span.size() - pattern_.size();
-    for (const std::byte* p = tail_span.data(); p <= scan_end_tail; ++p) {
-        if (full_match_at(p)) {
-            if (found_flag) found_flag->store(true, std::memory_order_relaxed);
-            return reinterpret_cast<uintptr_t>(p);
+
+    const auto scan_end_tail = range_end - pattern_.size();
+    while(current_pos <= scan_end_tail) {
+        if (found_flag && found_flag->load(std::memory_order_acquire)) return std::nullopt;
+        if (full_match_at(current_pos)) {
+            if (found_flag) found_flag->store(true, std::memory_order_release);
+            return reinterpret_cast<uintptr_t>(current_pos);
         }
+        current_pos++;
     }
+
     return std::nullopt;
 }
 #endif

include/ur/thread_pool.hpp

@@ -112,11 +112,13 @@ class ThreadPool {
             throw std::runtime_error("enqueue on stopped ThreadPool");
         }
 
+        tasks_in_flight_.fetch_add(1);
+
         // Distribute tasks to worker queues round-robin
         size_t queue_idx = submission_idx_.fetch_add(1) % thread_count_;
         queues_[queue_idx].push([task]() { (*task)(); });
 
-        condition_.notify_all();
+        condition_.notify_one();
         return res;
     }
 
@@ -129,6 +131,7 @@ class ThreadPool {
 
             // First, try to pop a task from our own queue.
             if (queues_[id].pop(task)) {
+                tasks_in_flight_.fetch_sub(1);
                 task();
                 continue;
             }
@@ -137,6 +140,7 @@ class ThreadPool {
             bool stolen = false;
             for (size_t i = 1; i < thread_count_; ++i) {
                 if (queues_[(id + i) % thread_count_].steal(task)) {
+                    tasks_in_flight_.fetch_sub(1);
                     stolen = true;
                     break;
                 }
@@ -148,12 +152,7 @@ class ThreadPool {
                 // If no task was found, wait for a notification.
                 std::unique_lock<std::mutex> lock(wait_mutex_);
                 condition_.wait(lock, [this] {
-                    if (stop_.load()) return true;
-                    // Wake up if any queue has tasks. This prevents missed notifications.
-                    for (size_t i = 0; i < thread_count_; ++i) {
-                        if (!queues_[i].empty()) return true;
-                    }
-                    return false;
+                    return stop_.load() || tasks_in_flight_.load() > 0;
                 });
             }
         }
@@ -165,6 +164,7 @@ class ThreadPool {
 
     std::atomic<bool> stop_;
     std::atomic<size_t> submission_idx_{0};
+    std::atomic<size_t> tasks_in_flight_{0};
 
     std::mutex wait_mutex_;
     std::condition_variable condition_;