From b3229c129511b725dfda6d4dca03ad39b6d6a9e9 Mon Sep 17 00:00:00 2001 From: Mingwei Zhang Date: Wed, 1 Jul 2026 13:33:56 -0700 Subject: [PATCH 1/5] fix: parallelize SSE search and add concurrency controls --- CHANGELOG.md | 16 +- src/bin/commands/config.rs | 20 ++ src/bin/commands/search.rs | 23 ++ src/bin/monocle.rs | 14 ++ src/config.rs | 51 ++++ src/lens/search/mod.rs | 485 ++++++++++++++++++++++++++++--------- src/server/README.md | 4 +- src/server/http.rs | 9 + src/server/mod.rs | 5 + src/server/search.rs | 312 ++++++++---------------- 10 files changed, 608 insertions(+), 331 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 2834e6a..7fc2a59 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -54,14 +54,20 @@ All notable changes to this project will be documented in this file. element batches, cancellation on disconnect, and max-results limit * Added server configuration fields to `MonocleConfig`: `server_address`, `server_port`, `server_max_search_batch_size`, `server_max_search_results`, - `server_search_timeout_secs`. All configurable via `monocle.toml` and - `MONOCLE_*` environment variables. -* SSE search streaming uses sequential file processing with `Arc` - cancellation — no new lens method needed. The server calls existing - `SearchFilters` utility methods (`to_broker_items`, `to_parser`) directly. + `server_search_timeout_secs`, `server_max_concurrent_searches`, and shared + `search_concurrency`. All configurable via `monocle.toml` and `MONOCLE_*` + environment variables. +* Parallelized SSE search streaming through the shared search executor while + preserving cancellation, bounded-channel backpressure, max-results handling, + timeout handling, and the single-terminal-event invariant. * Bounded mpsc channel (capacity 32) with backpressure: element batches are never dropped; progress events may be coalesced under backpressure. * Terminal event invariant: exactly one of `completed`, `cancelled`, or `error`. +* Added `--concurrency` to local search and server commands. Explicit values use + local rayon thread pools; unset/0 keeps rayon defaults including + `RAYON_NUM_THREADS`. +* Added server admission control for SSE search requests via + `server_max_concurrent_searches` (default 3); excess requests return HTTP 429. * Added Phase 2 REST API endpoints: - Tier 1 (stateless): `POST /time/parse`, `POST /country/lookup`, `POST /ip/lookup`, `GET /ip/public` diff --git a/src/bin/commands/config.rs b/src/bin/commands/config.rs index 1e25038..ecb4708 100644 --- a/src/bin/commands/config.rs +++ b/src/bin/commands/config.rs @@ -104,7 +104,9 @@ struct ServerDefaults { port: u16, max_search_batch_size: usize, max_search_results: u64, + search_concurrency: usize, search_timeout_secs: u64, + max_concurrent_searches: usize, auth_enabled: bool, } @@ -115,7 +117,9 @@ impl From<&MonocleConfig> for ServerDefaults { port: config.server_port, max_search_batch_size: config.server_max_search_batch_size, max_search_results: config.server_max_search_results, + search_concurrency: config.search_concurrency, search_timeout_secs: config.server_search_timeout_secs, + max_concurrent_searches: config.server_max_concurrent_searches, auth_enabled: config.server_auth_enabled, } } @@ -371,10 +375,26 @@ fn print_config_table(info: &ConfigInfo, verbose: bool) { " Search max results: {} (0 = unlimited)", info.server_defaults.max_search_results ); + println!( + " Search concurrency: {}", + if info.server_defaults.search_concurrency == 0 { + "auto".to_string() + } else { + info.server_defaults.search_concurrency.to_string() + } + ); println!( " Search timeout: {} seconds (0 = no timeout)", info.server_defaults.search_timeout_secs ); + println!( + " Max SSE searches: {}", + if info.server_defaults.max_concurrent_searches == 0 { + "unlimited".to_string() + } else { + info.server_defaults.max_concurrent_searches.to_string() + } + ); println!(" Auth enabled: {}", info.server_defaults.auth_enabled); if verbose { diff --git a/src/bin/commands/search.rs b/src/bin/commands/search.rs index d8481a8..a7e2281 100644 --- a/src/bin/commands/search.rs +++ b/src/bin/commands/search.rs @@ -88,6 +88,10 @@ pub struct SearchArgs { #[clap(flatten)] pub filters: SearchFilters, + /// Search concurrency (0 = auto/rayon default; overrides config) + #[clap(long)] + pub concurrency: Option, + /// Remote Monocle server URL (e.g., http://localhost:8080/api/v1/search/stream). /// When set, search runs against the remote server instead of locally. #[clap(long, value_name = "URL")] @@ -554,6 +558,24 @@ fn fetch_broker_items_cached( } pub fn run(config: &MonocleConfig, args: SearchArgs, output_format: OutputFormat) { + let concurrency = args.concurrency.unwrap_or(config.search_concurrency); + if concurrency > 0 { + match rayon::ThreadPoolBuilder::new() + .num_threads(concurrency) + .build() + { + Ok(pool) => pool.install(|| run_inner(config, args, output_format)), + Err(e) => { + eprintln!("Failed to create rayon thread pool: {e}"); + std::process::exit(1); + } + } + } else { + run_inner(config, args, output_format); + } +} + +fn run_inner(config: &MonocleConfig, args: SearchArgs, output_format: OutputFormat) { let SearchArgs { dry_run, sqlite_path, @@ -569,6 +591,7 @@ pub fn run(config: &MonocleConfig, args: SearchArgs, output_format: OutputFormat use_cache, cache_dir, mut filters, + concurrency: _, remote_url, remote_token, } = args; diff --git a/src/bin/monocle.rs b/src/bin/monocle.rs index 674b78a..dfad481 100644 --- a/src/bin/monocle.rs +++ b/src/bin/monocle.rs @@ -115,10 +115,18 @@ struct ServerArgs { #[clap(long)] max_search_results: Option, + /// Search concurrency (0 = auto/rayon default, overrides config) + #[clap(long)] + concurrency: Option, + /// Search timeout in seconds (0 = no timeout, overrides config) #[clap(long)] search_timeout_secs: Option, + /// Maximum concurrent SSE search requests (0 = unlimited, overrides config) + #[clap(long)] + max_concurrent_searches: Option, + /// Enable token auth for /api/v1/* endpoints (overrides config) #[clap(long)] auth_enabled: Option, @@ -203,9 +211,15 @@ fn main() { if let Some(v) = args.max_search_results { server_config.server_max_search_results = v; } + if let Some(v) = args.concurrency { + server_config.search_concurrency = v; + } if let Some(v) = args.search_timeout_secs { server_config.server_search_timeout_secs = v; } + if let Some(v) = args.max_concurrent_searches { + server_config.server_max_concurrent_searches = v; + } if let Some(v) = args.auth_enabled { server_config.server_auth_enabled = v; } diff --git a/src/config.rs b/src/config.rs index 46ba570..9aad06a 100644 --- a/src/config.rs +++ b/src/config.rs @@ -24,6 +24,12 @@ pub const DEFAULT_SERVER_MAX_SEARCH_RESULTS: u64 = 0; /// Default search timeout in seconds (0 = no timeout) pub const DEFAULT_SERVER_SEARCH_TIMEOUT_SECS: u64 = 0; +/// Default search concurrency (0 = rayon default / CPU count) +pub const DEFAULT_SEARCH_CONCURRENCY: usize = 0; + +/// Default maximum concurrent SSE search requests +pub const DEFAULT_SERVER_MAX_CONCURRENT_SEARCHES: usize = 3; + /// Default: auth disabled pub const DEFAULT_SERVER_AUTH_ENABLED: bool = false; @@ -69,9 +75,15 @@ pub struct MonocleConfig { /// Maximum search results per request, 0 = unlimited (default: 0) pub server_max_search_results: u64, + /// Search concurrency, 0 = rayon default / CPU count (default: 0) + pub search_concurrency: usize, + /// Search timeout in seconds, 0 = no timeout (default: 0) pub server_search_timeout_secs: u64, + /// Maximum concurrent SSE search requests, 0 = unlimited (default: 3) + pub server_max_concurrent_searches: usize, + /// Enable token auth for /api/v1/* endpoints (default: false) pub server_auth_enabled: bool, @@ -101,6 +113,10 @@ const EMPTY_CONFIG: &str = r#"### monocle configuration file ### If true, error out instead of falling back to Cloudflare when RTR fails # rpki_rtr_no_fallback = false +### Search execution configuration +### Search concurrency; 0 = auto/rayon default. Can also be set with MONOCLE_SEARCH_CONCURRENCY. +# search_concurrency = 0 + ### HTTP service configuration ### These settings control the monocle HTTP/SSE server. # server_address = "127.0.0.1" @@ -111,6 +127,8 @@ const EMPTY_CONFIG: &str = r#"### monocle configuration file # server_max_search_results = 0 ### Search timeout in seconds (0 = no timeout) # server_search_timeout_secs = 0 +### Maximum concurrent SSE search requests (0 = unlimited) +# server_max_concurrent_searches = 3 ### Auth (token-based, disabled by default) ### When enabled, requests to /api/v1/* require Authorization: Bearer @@ -228,7 +246,9 @@ impl Default for MonocleConfig { server_port: DEFAULT_SERVER_PORT, server_max_search_batch_size: DEFAULT_SERVER_MAX_SEARCH_BATCH_SIZE, server_max_search_results: DEFAULT_SERVER_MAX_SEARCH_RESULTS, + search_concurrency: DEFAULT_SEARCH_CONCURRENCY, server_search_timeout_secs: DEFAULT_SERVER_SEARCH_TIMEOUT_SECS, + server_max_concurrent_searches: DEFAULT_SERVER_MAX_CONCURRENT_SEARCHES, server_auth_enabled: DEFAULT_SERVER_AUTH_ENABLED, server_auth_token: String::new(), } @@ -359,10 +379,18 @@ impl MonocleConfig { .get("server_max_search_results") .and_then(|s| s.parse().ok()) .unwrap_or(DEFAULT_SERVER_MAX_SEARCH_RESULTS); + let search_concurrency = config + .get("search_concurrency") + .and_then(|s| s.parse().ok()) + .unwrap_or(DEFAULT_SEARCH_CONCURRENCY); let server_search_timeout_secs = config .get("server_search_timeout_secs") .and_then(|s| s.parse().ok()) .unwrap_or(DEFAULT_SERVER_SEARCH_TIMEOUT_SECS); + let server_max_concurrent_searches = config + .get("server_max_concurrent_searches") + .and_then(|s| s.parse().ok()) + .unwrap_or(DEFAULT_SERVER_MAX_CONCURRENT_SEARCHES); // Parse auth configuration let server_auth_enabled = config @@ -385,7 +413,9 @@ impl MonocleConfig { server_port, server_max_search_batch_size, server_max_search_results, + search_concurrency, server_search_timeout_secs, + server_max_concurrent_searches, server_auth_enabled, server_auth_token, }) @@ -461,10 +491,26 @@ impl MonocleConfig { "Search Max Results: {}", self.server_max_search_results )); + lines.push(format!( + "Search Concurrency: {}", + if self.search_concurrency == 0 { + "auto".to_string() + } else { + self.search_concurrency.to_string() + } + )); lines.push(format!( "Search Timeout: {} seconds", self.server_search_timeout_secs )); + lines.push(format!( + "Max SSE Searches: {}", + if self.server_max_concurrent_searches == 0 { + "unlimited".to_string() + } else { + self.server_max_concurrent_searches.to_string() + } + )); lines.push(format!("Auth Enabled: {}", self.server_auth_enabled)); // Check if cache directories exist and show status @@ -941,10 +987,15 @@ mod tests { config.server_max_search_results, DEFAULT_SERVER_MAX_SEARCH_RESULTS ); + assert_eq!(config.search_concurrency, DEFAULT_SEARCH_CONCURRENCY); assert_eq!( config.server_search_timeout_secs, DEFAULT_SERVER_SEARCH_TIMEOUT_SECS ); + assert_eq!( + config.server_max_concurrent_searches, + DEFAULT_SERVER_MAX_CONCURRENT_SEARCHES + ); assert!(!config.server_auth_enabled); assert_eq!(config.server_auth_token, ""); } diff --git a/src/lens/search/mod.rs b/src/lens/search/mod.rs index d73d64c..d21e221 100644 --- a/src/lens/search/mod.rs +++ b/src/lens/search/mod.rs @@ -46,9 +46,9 @@ use bgpkit_parser::BgpkitParser; use rayon::prelude::*; use serde::{Deserialize, Serialize}; use std::io::Read; -use std::sync::atomic::{AtomicU64, Ordering}; +use std::sync::atomic::{AtomicBool, AtomicU64, Ordering}; use std::sync::Arc; -use std::time::Instant; +use std::time::{Duration, Instant}; #[cfg(feature = "cli")] use clap::{Args, ValueEnum}; @@ -145,6 +145,59 @@ pub type SearchProgressCallback = Arc; /// Called for each BGP element found during search, along with the collector ID. pub type ElementHandler = Arc; +/// Runtime options for executing a search. +#[derive(Clone, Default)] +pub struct SearchExecutionOptions { + /// Number of rayon worker threads for this search. `None` or `Some(0)` uses rayon's default. + pub concurrency: Option, + /// Maximum number of elements to emit. `None` means unlimited. + pub max_results: Option, + /// Maximum wall-clock runtime after broker query starts. `None` means no timeout. + pub timeout: Option, + /// Optional external cancellation flag. + pub cancel_flag: Option>, + /// Number of elements per batch passed to [`SearchSink::on_elements`]. + pub batch_size: usize, +} + +/// Reason a search execution stopped. +#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)] +pub enum SearchExitReason { + Completed, + Cancelled, + Timeout, + MaxResultsReached, +} + +/// Search execution result including both summary and stop reason. +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct SearchOutcome { + pub summary: SearchSummary, + pub exit_reason: SearchExitReason, +} + +/// A batch of matched elements from a single MRT file. +pub struct SearchElementBatch { + pub file_index: usize, + pub file_url: String, + pub collector: String, + pub elements: Vec, +} + +/// Control returned by a search sink after receiving elements. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum SearchControl { + Continue, + Stop, +} + +/// Sink used by the shared search executor. +pub trait SearchSink: Send + Sync { + fn on_progress(&self, _progress: SearchProgress) {} + + fn on_elements(&self, batch: SearchElementBatch) -> SearchControl; +} + // ============================================================================= // Types // ============================================================================= @@ -383,142 +436,109 @@ impl SearchLens { progress_callback: Option, element_handler: ElementHandler, ) -> Result { - // Notify broker query starting - if let Some(ref cb) = progress_callback { - cb(SearchProgress::QueryingBroker); + struct CallbackSink { + progress_callback: Option, + element_handler: ElementHandler, } - // Query broker - let items = self.query_broker(filters)?; - let total_files = items.len(); + impl SearchSink for CallbackSink { + fn on_progress(&self, progress: SearchProgress) { + if let Some(ref cb) = self.progress_callback { + cb(progress); + } + } - // Notify files found - if let Some(ref cb) = progress_callback { - cb(SearchProgress::FilesFound { count: total_files }); + fn on_elements(&self, batch: SearchElementBatch) -> SearchControl { + for elem in batch.elements { + (self.element_handler)(elem, batch.collector.clone()); + } + SearchControl::Continue + } } - if total_files == 0 { - if let Some(ref cb) = progress_callback { - cb(SearchProgress::Completed { - total_files: 0, - successful_files: 0, - failed_files: 0, - total_messages: 0, - duration_secs: 0.0, - files_per_sec: None, - }); - } + let sink = Arc::new(CallbackSink { + progress_callback, + element_handler, + }); + let outcome = self.search_with_options(filters, SearchExecutionOptions::default(), sink)?; + Ok(outcome.summary) + } - return Ok(SearchSummary { + /// Search BGP messages using the shared executor. + pub fn search_with_options( + &self, + filters: &SearchFilters, + options: SearchExecutionOptions, + sink: Arc, + ) -> Result { + sink.on_progress(SearchProgress::QueryingBroker); + + let items = self.query_broker(filters)?; + let total_files = items.len(); + sink.on_progress(SearchProgress::FilesFound { count: total_files }); + + if total_files == 0 { + let summary = SearchSummary { + total_files: 0, + successful_files: 0, + failed_files: 0, + total_messages: 0, + duration_secs: 0.0, + }; + sink.on_progress(SearchProgress::Completed { total_files: 0, successful_files: 0, failed_files: 0, total_messages: 0, duration_secs: 0.0, + files_per_sec: None, + }); + return Ok(SearchOutcome { + summary, + exit_reason: SearchExitReason::Completed, }); } let start_time = Instant::now(); + let batch_size = options.batch_size.max(1); + let max_results = options.max_results.filter(|max| *max > 0); + let deadline = options.timeout.map(|timeout| start_time + timeout); + let external_cancel = options.cancel_flag.clone(); + let stop_flag = AtomicBool::new(false); + let exit_reason = AtomicU64::new(0); let files_completed = AtomicU64::new(0); let successful_files = AtomicU64::new(0); let failed_files = AtomicU64::new(0); let total_messages = AtomicU64::new(0); - // Process files in parallel - items.into_par_iter().enumerate().for_each(|(index, item)| { - let url = item.url.clone(); - let collector = item.collector_id.clone(); - - // Notify file started - if let Some(ref cb) = progress_callback { - cb(SearchProgress::FileStarted { - file_index: index, + let run = || { + items.into_par_iter().enumerate().for_each(|(index, item)| { + let state = SearchWorkerState { total_files, - file_url: url.clone(), - collector: collector.clone(), - }); - } + start_time, + deadline, + batch_size, + max_results, + external_cancel: external_cancel.as_deref(), + stop_flag: &stop_flag, + exit_reason: &exit_reason, + files_completed: &files_completed, + successful_files: &successful_files, + failed_files: &failed_files, + total_messages: &total_messages, + sink: sink.as_ref(), + }; + process_search_item(filters, index, item, state); + }); + }; - // Try to parse the file - match filters.to_parser(url.as_str()) { - Ok(parser) => { - let mut file_messages: u64 = 0; - - for elem in parser { - element_handler(elem, collector.clone()); - file_messages += 1; - } - - total_messages.fetch_add(file_messages, Ordering::Relaxed); - successful_files.fetch_add(1, Ordering::Relaxed); - let completed = files_completed.fetch_add(1, Ordering::Relaxed) + 1; - - // Notify file completed - if let Some(ref cb) = progress_callback { - cb(SearchProgress::FileCompleted { - file_index: index, - total_files, - messages_found: file_messages, - success: true, - error: None, - }); - - // Send progress update - let elapsed = start_time.elapsed().as_secs_f64(); - let percent = completed as f64 / total_files as f64 * 100.0; - let eta = if completed > 0 && percent < 100.0 { - let rate = elapsed / completed as f64; - Some(rate * (total_files as u64 - completed) as f64) - } else { - None - }; - - cb(SearchProgress::ProgressUpdate { - files_completed: completed as usize, - total_files, - total_messages: total_messages.load(Ordering::Relaxed), - percent_complete: percent, - elapsed_secs: elapsed, - eta_secs: eta, - }); - } - } - Err(e) => { - failed_files.fetch_add(1, Ordering::Relaxed); - let completed = files_completed.fetch_add(1, Ordering::Relaxed) + 1; - - // Notify file failed - if let Some(ref cb) = progress_callback { - cb(SearchProgress::FileCompleted { - file_index: index, - total_files, - messages_found: 0, - success: false, - error: Some(e.to_string()), - }); - - // Send progress update - let elapsed = start_time.elapsed().as_secs_f64(); - let percent = completed as f64 / total_files as f64 * 100.0; - let eta = if completed > 0 && percent < 100.0 { - let rate = elapsed / completed as f64; - Some(rate * (total_files as u64 - completed) as f64) - } else { - None - }; - - cb(SearchProgress::ProgressUpdate { - files_completed: completed as usize, - total_files, - total_messages: total_messages.load(Ordering::Relaxed), - percent_complete: percent, - elapsed_secs: elapsed, - eta_secs: eta, - }); - } - } + match options.concurrency { + Some(n) if n > 0 => { + let pool = rayon::ThreadPoolBuilder::new().num_threads(n).build()?; + pool.install(run); } - }); + _ => run(), + } let duration_secs = start_time.elapsed().as_secs_f64(); let final_successful = successful_files.load(Ordering::Relaxed) as usize; @@ -530,9 +550,18 @@ impl SearchLens { None }; - // Notify completion - if let Some(ref cb) = progress_callback { - cb(SearchProgress::Completed { + let exit_reason = match exit_reason.load(Ordering::Relaxed) { + 1 => SearchExitReason::Cancelled, + 2 => SearchExitReason::Timeout, + 3 => SearchExitReason::MaxResultsReached, + _ => SearchExitReason::Completed, + }; + + if matches!( + exit_reason, + SearchExitReason::Completed | SearchExitReason::MaxResultsReached + ) { + sink.on_progress(SearchProgress::Completed { total_files, successful_files: final_successful, failed_files: final_failed, @@ -542,12 +571,17 @@ impl SearchLens { }); } - Ok(SearchSummary { + let summary = SearchSummary { total_files, successful_files: final_successful, failed_files: final_failed, total_messages: final_messages, duration_secs, + }; + + Ok(SearchOutcome { + summary, + exit_reason, }) } @@ -592,6 +626,219 @@ impl SearchLens { } } +struct SearchWorkerState<'a> { + total_files: usize, + start_time: Instant, + deadline: Option, + batch_size: usize, + max_results: Option, + external_cancel: Option<&'a AtomicBool>, + stop_flag: &'a AtomicBool, + exit_reason: &'a AtomicU64, + files_completed: &'a AtomicU64, + successful_files: &'a AtomicU64, + failed_files: &'a AtomicU64, + total_messages: &'a AtomicU64, + sink: &'a dyn SearchSink, +} + +fn process_search_item( + filters: &SearchFilters, + index: usize, + item: BrokerItem, + state: SearchWorkerState<'_>, +) { + if search_should_stop(&state) { + return; + } + + let url = item.url.clone(); + let collector = item.collector_id.clone(); + + state.sink.on_progress(SearchProgress::FileStarted { + file_index: index, + total_files: state.total_files, + file_url: url.clone(), + collector: collector.clone(), + }); + + let parser = match filters.to_parser(url.as_str()) { + Ok(parser) => parser, + Err(e) => { + state.failed_files.fetch_add(1, Ordering::Relaxed); + let completed = state.files_completed.fetch_add(1, Ordering::Relaxed) + 1; + state.sink.on_progress(SearchProgress::FileCompleted { + file_index: index, + total_files: state.total_files, + messages_found: 0, + success: false, + error: Some(e.to_string()), + }); + send_progress_update(&state, completed); + return; + } + }; + + let mut file_messages = 0_u64; + let mut batch = Vec::with_capacity(state.batch_size); + + for elem in parser { + if search_should_stop(&state) { + break; + } + + batch.push(elem); + if batch.len() >= state.batch_size { + let accepted = emit_search_batch(&state, index, &url, &collector, &mut batch); + file_messages += accepted; + } + } + + if !batch.is_empty() && !search_should_stop(&state) { + let accepted = emit_search_batch(&state, index, &url, &collector, &mut batch); + file_messages += accepted; + } + + state.successful_files.fetch_add(1, Ordering::Relaxed); + let completed = state.files_completed.fetch_add(1, Ordering::Relaxed) + 1; + state.sink.on_progress(SearchProgress::FileCompleted { + file_index: index, + total_files: state.total_files, + messages_found: file_messages, + success: true, + error: None, + }); + send_progress_update(&state, completed); +} + +fn search_should_stop(state: &SearchWorkerState<'_>) -> bool { + if state.stop_flag.load(Ordering::Relaxed) { + return true; + } + + if let Some(cancel) = state.external_cancel { + if cancel.load(Ordering::Relaxed) { + state + .exit_reason + .compare_exchange(0, 1, Ordering::Relaxed, Ordering::Relaxed) + .ok(); + state.stop_flag.store(true, Ordering::Relaxed); + return true; + } + } + + if let Some(deadline) = state.deadline { + if Instant::now() >= deadline { + state + .exit_reason + .compare_exchange(0, 2, Ordering::Relaxed, Ordering::Relaxed) + .ok(); + state.stop_flag.store(true, Ordering::Relaxed); + return true; + } + } + + false +} + +fn emit_search_batch( + state: &SearchWorkerState<'_>, + file_index: usize, + file_url: &str, + collector: &str, + batch: &mut Vec, +) -> u64 { + let original_len = batch.len(); + let accepted = reserve_result_slots(state, original_len as u64); + if accepted == 0 { + batch.clear(); + return 0; + } + + if accepted < original_len as u64 { + batch.truncate(accepted as usize); + } + + let elements = std::mem::take(batch); + let control = state.sink.on_elements(SearchElementBatch { + file_index, + file_url: file_url.to_string(), + collector: collector.to_string(), + elements, + }); + + if control == SearchControl::Stop { + state + .exit_reason + .compare_exchange(0, 1, Ordering::Relaxed, Ordering::Relaxed) + .ok(); + state.stop_flag.store(true, Ordering::Relaxed); + } + + accepted +} + +fn reserve_result_slots(state: &SearchWorkerState<'_>, wanted: u64) -> u64 { + match state.max_results { + None => { + state.total_messages.fetch_add(wanted, Ordering::Relaxed); + wanted + } + Some(max) => loop { + let current = state.total_messages.load(Ordering::Relaxed); + if current >= max { + state + .exit_reason + .compare_exchange(0, 3, Ordering::Relaxed, Ordering::Relaxed) + .ok(); + state.stop_flag.store(true, Ordering::Relaxed); + return 0; + } + + let allowed = wanted.min(max - current); + if state + .total_messages + .compare_exchange( + current, + current + allowed, + Ordering::Relaxed, + Ordering::Relaxed, + ) + .is_ok() + { + if allowed < wanted || current + allowed >= max { + state + .exit_reason + .compare_exchange(0, 3, Ordering::Relaxed, Ordering::Relaxed) + .ok(); + state.stop_flag.store(true, Ordering::Relaxed); + } + return allowed; + } + }, + } +} + +fn send_progress_update(state: &SearchWorkerState<'_>, completed: u64) { + let elapsed = state.start_time.elapsed().as_secs_f64(); + let percent = completed as f64 / state.total_files as f64 * 100.0; + let eta = if completed > 0 && percent < 100.0 { + let rate = elapsed / completed as f64; + Some(rate * (state.total_files as u64 - completed) as f64) + } else { + None + }; + + state.sink.on_progress(SearchProgress::ProgressUpdate { + files_completed: completed as usize, + total_files: state.total_files, + total_messages: state.total_messages.load(Ordering::Relaxed), + percent_complete: percent, + elapsed_secs: elapsed, + eta_secs: eta, + }); +} + impl Default for SearchLens { fn default() -> Self { Self::new() diff --git a/src/server/README.md b/src/server/README.md index 26402a6..9569390 100644 --- a/src/server/README.md +++ b/src/server/README.md @@ -290,11 +290,13 @@ Server settings are part of `MonocleConfig` and configurable via `monocle.toml` or `MONOCLE_*` environment variables. See `monocle.toml.example` for all options. ```toml +search_concurrency = 0 # 0 = auto/rayon default server_address = "127.0.0.1" server_port = 8080 server_max_search_batch_size = 100 server_max_search_results = 0 # 0 = unlimited server_search_timeout_secs = 0 # 0 = no timeout +server_max_concurrent_searches = 3 # 0 = unlimited; excess requests get 429 server_auth_enabled = false server_auth_token = "" ``` @@ -302,7 +304,7 @@ server_auth_token = "" CLI flags override config values: ```bash -monocle server --address 0.0.0.0 --port 9000 --max-search-batch-size 50 +monocle server --address 0.0.0.0 --port 9000 --max-search-batch-size 50 --concurrency 4 ``` ## Docker Deployment diff --git a/src/server/http.rs b/src/server/http.rs index ecba1c0..276be19 100644 --- a/src/server/http.rs +++ b/src/server/http.rs @@ -71,6 +71,8 @@ pub enum ApiErrorCode { Cancelled, /// Search failed during execution SearchFailed, + /// Too many concurrent requests + TooManyRequests, /// Required local data not initialized NotInitialized, /// Unexpected server error @@ -107,6 +109,13 @@ impl ApiError { ApiErrorResponse::internal(message), ) } + + pub fn too_many_requests(message: impl Into) -> Self { + Self::new( + StatusCode::TOO_MANY_REQUESTS, + ApiErrorResponse::new(ApiErrorCode::TooManyRequests, message), + ) + } } impl IntoResponse for ApiError { diff --git a/src/server/mod.rs b/src/server/mod.rs index 1ddbc5f..dfa81a5 100644 --- a/src/server/mod.rs +++ b/src/server/mod.rs @@ -29,6 +29,7 @@ use axum::middleware::from_fn_with_state; use axum::routing::get; use axum::Router as AxumRouter; use std::sync::Arc; +use tokio::sync::Semaphore; use tower_http::cors::{Any, CorsLayer}; use crate::config::MonocleConfig; @@ -41,6 +42,7 @@ use crate::config::MonocleConfig; #[derive(Clone)] pub struct ServerState { pub config: Arc, + pub search_permits: Arc, } // ============================================================================= @@ -54,8 +56,10 @@ pub async fn start_server(config: MonocleConfig) -> anyhow::Result<()> { let auth_enabled = config.server_auth_enabled; let auth_token = config.server_auth_token.trim().to_string(); + let max_concurrent_searches = config.server_max_concurrent_searches; let state = ServerState { config: Arc::new(config), + search_permits: Arc::new(Semaphore::new(max_concurrent_searches)), }; let cors = CorsLayer::new() @@ -113,6 +117,7 @@ mod tests { let config = MonocleConfig::default(); let state = ServerState { config: Arc::new(config), + search_permits: Arc::new(Semaphore::new(3)), }; let _cloned = state.clone(); } diff --git a/src/server/search.rs b/src/server/search.rs index 7465b6e..b1a006c 100644 --- a/src/server/search.rs +++ b/src/server/search.rs @@ -2,16 +2,16 @@ //! //! `POST /api/v1/search/stream` accepts a JSON request body and returns a //! `text/event-stream` response. The server runs a sequential, cancellable -//! search loop in a blocking task, streaming progress and element batch -//! events to the client. Cancellation is triggered by closing the HTTP +//! parallel search executor in a blocking task, streaming progress and element +//! batch events to the client. Cancellation is triggered by closing the HTTP //! connection — when the SSE response is dropped, the cancellation flag is //! set and the worker stops. //! //! See `SSE_SERVICE_OVERHAUL_DESIGN.md` Section 6 for the algorithm. use std::sync::atomic::{AtomicBool, Ordering}; -use std::sync::Arc; -use std::time::{Duration, Instant}; +use std::sync::{Arc, Mutex}; +use std::time::Duration; use axum::extract::State; use axum::response::sse::{Event as SseEvent, Sse}; @@ -23,7 +23,10 @@ use tokio::sync::mpsc; use tokio_stream::wrappers::ReceiverStream; use crate::lens::parse::ParseElemType; -use crate::lens::search::{SearchFilters, SearchProgress, SearchSummary}; +use crate::lens::search::{ + SearchControl, SearchElementBatch, SearchExecutionOptions, SearchExitReason, SearchFilters, + SearchLens, SearchProgress, SearchSink, SearchSummary, +}; use crate::server::http::{ApiError, ApiErrorCode, ApiErrorResponse}; use crate::server::ServerState; @@ -161,6 +164,7 @@ pub struct SearchStarted { pub struct ElementsBatch { pub total_so_far: u64, pub collector: Option, + pub file_url: String, pub elements: Vec, } @@ -227,7 +231,8 @@ pub async fn stream_search( .min(config.server_max_search_batch_size) .max(1); - let max_results = match (request.max_results, config.server_max_search_results) { + let requested_max_results = request.max_results.filter(|max| *max > 0); + let max_results = match (requested_max_results, config.server_max_search_results) { (Some(r), 0) => Some(r), // server unlimited (Some(r), limit) => Some(r.min(limit)), // clamp to server limit (None, 0) => None, // both unlimited @@ -240,20 +245,42 @@ pub async fn stream_search( None }; - // 4. Create bounded channel and cancellation flag + // 4. Enforce server-side concurrent search limit. Requests are rejected + // immediately instead of queued so clients get deterministic feedback. + let search_permit = if config.server_max_concurrent_searches == 0 { + None + } else { + Some( + state + .search_permits + .clone() + .try_acquire_owned() + .map_err(|_| ApiError::too_many_requests("too many concurrent search requests"))?, + ) + }; + + let concurrency = if config.search_concurrency > 0 { + Some(config.search_concurrency) + } else { + None + }; + + // 5. Create bounded channel and cancellation flag let (tx, rx) = mpsc::channel::(32); let cancel_flag = Arc::new(AtomicBool::new(false)); - // 5. Spawn the search worker in a blocking task + // 6. Spawn the search worker in a blocking task let worker_cancel_flag = cancel_flag.clone(); let worker_tx = tx.clone(); tokio::task::spawn_blocking(move || { + let _search_permit = search_permit; run_search_worker( filters, batch_size, max_results, timeout_secs, + concurrency, worker_cancel_flag, worker_tx, ); @@ -293,10 +320,10 @@ fn run_search_worker( batch_size: usize, max_results: Option, timeout_secs: Option, + concurrency: Option, cancel_flag: Arc, event_tx: mpsc::Sender, ) { - // Send Started event let _ = send_event( &event_tx, SearchStreamEvent::Started(SearchStarted { @@ -306,19 +333,18 @@ fn run_search_worker( }), ); - let start_time = Instant::now(); - let deadline = timeout_secs.map(|s| start_time + Duration::from_secs(s)); - - let is_cancelled = || cancel_flag.load(Ordering::Relaxed); - - // 1. Query broker for files - let _ = send_event( - &event_tx, - SearchStreamEvent::Progress(SearchProgress::QueryingBroker), - ); + let sink = Arc::new(SseSearchSink::new(event_tx.clone(), cancel_flag.clone())); + let options = SearchExecutionOptions { + concurrency, + max_results, + timeout: timeout_secs.map(Duration::from_secs), + cancel_flag: Some(cancel_flag), + batch_size, + }; - let items = match filters.to_broker_items() { - Ok(items) => items, + let lens = SearchLens::new(); + let outcome = match lens.search_with_options(&filters, options, sink) { + Ok(outcome) => outcome, Err(e) => { let _ = send_event( &event_tx, @@ -331,204 +357,78 @@ fn run_search_worker( } }; - let total_files = items.len(); - let _ = send_event( - &event_tx, - SearchStreamEvent::Progress(SearchProgress::FilesFound { count: total_files }), - ); - - if total_files == 0 { - let _ = send_event( - &event_tx, - SearchStreamEvent::Completed(SearchSummary { - total_files: 0, - successful_files: 0, - failed_files: 0, - total_messages: 0, - duration_secs: start_time.elapsed().as_secs_f64(), - }), - ); - return; - } - - let mut successful_files: usize = 0; - let mut failed_files: usize = 0; - let mut total_messages: u64 = 0; - let mut total_elements_sent: u64 = 0; - - // Process files sequentially - for (index, item) in items.into_iter().enumerate() { - if is_cancelled() { - break; + match outcome.exit_reason { + SearchExitReason::Cancelled => { + let _ = send_event(&event_tx, SearchStreamEvent::Cancelled); } - - // Check timeout - if let Some(dl) = deadline { - if Instant::now() >= dl { - let _ = send_event( - &event_tx, - SearchStreamEvent::Error(ApiErrorResponse::new( - ApiErrorCode::SearchFailed, - format!( - "Search timed out after {} seconds", - timeout_secs.unwrap_or(0) - ), - )), - ); - return; - } - } - - let url = &item.url; - let collector = item.collector_id.clone(); - - let _ = send_event( - &event_tx, - SearchStreamEvent::Progress(SearchProgress::FileStarted { - file_index: index, - total_files, - file_url: url.clone(), - collector: collector.clone(), - }), - ); - - let parser = match filters.to_parser(url) { - Ok(p) => p, - Err(e) => { - failed_files += 1; - let _ = send_event( - &event_tx, - SearchStreamEvent::Progress(SearchProgress::FileCompleted { - file_index: index, - total_files, - messages_found: 0, - success: false, - error: Some(e.to_string()), - }), - ); - continue; - } - }; - - let mut file_messages: u64 = 0; - let mut batch: Vec = Vec::with_capacity(batch_size); - let mut max_reached = false; - - for elem in parser { - if is_cancelled() { - break; - } - - file_messages += 1; - total_elements_sent += 1; - batch.push(elem); - - if batch.len() >= batch_size { - let batch_event = SearchStreamEvent::Elements(ElementsBatch { - total_so_far: total_elements_sent, - collector: Some(collector.clone()), - elements: std::mem::take(&mut batch), - }); - - if send_event(&event_tx, batch_event).is_err() { - // Channel closed (client gone) — cancel - cancel_flag.store(true, Ordering::Relaxed); - break; - } - } - - // Check max_results - if let Some(max) = max_results { - if total_elements_sent >= max { - max_reached = true; - break; - } - } - } - - // Flush remaining partial batch for this file - if !batch.is_empty() { + SearchExitReason::Timeout => { let _ = send_event( &event_tx, - SearchStreamEvent::Elements(ElementsBatch { - total_so_far: total_elements_sent, - collector: Some(collector.clone()), - elements: std::mem::take(&mut batch), - }), + SearchStreamEvent::Error(ApiErrorResponse::new( + ApiErrorCode::SearchFailed, + format!( + "Search timed out after {} seconds", + timeout_secs.unwrap_or(0) + ), + )), ); } - - if max_reached { - successful_files += 1; - total_messages += file_messages; - let _ = send_event( - &event_tx, - SearchStreamEvent::Completed(SearchSummary { - total_files, - successful_files, - failed_files, - total_messages, - duration_secs: start_time.elapsed().as_secs_f64(), - }), - ); - return; + SearchExitReason::Completed | SearchExitReason::MaxResultsReached => { + let _ = send_event(&event_tx, SearchStreamEvent::Completed(outcome.summary)); } + } +} + +struct SseSearchState { + total_so_far: u64, +} + +struct SseSearchSink { + state: Mutex, + event_tx: mpsc::Sender, + cancel_flag: Arc, +} - if is_cancelled() { - break; +impl SseSearchSink { + fn new(event_tx: mpsc::Sender, cancel_flag: Arc) -> Self { + Self { + state: Mutex::new(SseSearchState { total_so_far: 0 }), + event_tx, + cancel_flag, } + } +} - successful_files += 1; - total_messages += file_messages; - - let _ = send_event( - &event_tx, - SearchStreamEvent::Progress(SearchProgress::FileCompleted { - file_index: index, - total_files, - messages_found: file_messages, - success: true, - error: None, - }), - ); +impl SearchSink for SseSearchSink { + fn on_progress(&self, progress: SearchProgress) { + if send_event(&self.event_tx, SearchStreamEvent::Progress(progress)).is_err() { + self.cancel_flag.store(true, Ordering::Relaxed); + } + } - let completed = index + 1; - let elapsed = start_time.elapsed().as_secs_f64(); - let percent = completed as f64 / total_files as f64 * 100.0; - let eta = if completed < total_files && percent < 100.0 { - let rate = elapsed / completed as f64; - Some(rate * (total_files - completed) as f64) - } else { - None + fn on_elements(&self, batch: SearchElementBatch) -> SearchControl { + let mut state = match self.state.lock() { + Ok(state) => state, + Err(_) => { + self.cancel_flag.store(true, Ordering::Relaxed); + return SearchControl::Stop; + } }; - let _ = send_event( - &event_tx, - SearchStreamEvent::Progress(SearchProgress::ProgressUpdate { - files_completed: completed, - total_files, - total_messages, - percent_complete: percent, - elapsed_secs: elapsed, - eta_secs: eta, - }), - ); - } + state.total_so_far += batch.elements.len() as u64; + let event = SearchStreamEvent::Elements(ElementsBatch { + total_so_far: state.total_so_far, + collector: Some(batch.collector), + file_url: batch.file_url, + elements: batch.elements, + }); - // Send terminal event - if is_cancelled() { - let _ = send_event(&event_tx, SearchStreamEvent::Cancelled); - } else { - let _ = send_event( - &event_tx, - SearchStreamEvent::Completed(SearchSummary { - total_files, - successful_files, - failed_files, - total_messages, - duration_secs: start_time.elapsed().as_secs_f64(), - }), - ); + if send_event(&self.event_tx, event).is_err() { + self.cancel_flag.store(true, Ordering::Relaxed); + SearchControl::Stop + } else { + SearchControl::Continue + } } } From de872e1bfc597d102e50f21254e1e811785e2122 Mon Sep 17 00:00:00 2001 From: Mingwei Zhang Date: Wed, 1 Jul 2026 14:08:08 -0700 Subject: [PATCH 2/5] test: cover search executor stop conditions --- src/lens/search/mod.rs | 148 +++++++++++++++++++++++++++++++++++++++-- 1 file changed, 144 insertions(+), 4 deletions(-) diff --git a/src/lens/search/mod.rs b/src/lens/search/mod.rs index d21e221..6059633 100644 --- a/src/lens/search/mod.rs +++ b/src/lens/search/mod.rs @@ -699,14 +699,29 @@ fn process_search_item( file_messages += accepted; } - state.successful_files.fetch_add(1, Ordering::Relaxed); + let stopped_reason = current_exit_reason(&state); + let stopped_before_success = matches!( + stopped_reason, + Some(SearchExitReason::Cancelled | SearchExitReason::Timeout) + ); + + if stopped_before_success { + state.failed_files.fetch_add(1, Ordering::Relaxed); + } else { + state.successful_files.fetch_add(1, Ordering::Relaxed); + } + let completed = state.files_completed.fetch_add(1, Ordering::Relaxed) + 1; state.sink.on_progress(SearchProgress::FileCompleted { file_index: index, total_files: state.total_files, messages_found: file_messages, - success: true, - error: None, + success: !stopped_before_success, + error: stopped_reason.and_then(|reason| match reason { + SearchExitReason::Cancelled => Some("search cancelled".to_string()), + SearchExitReason::Timeout => Some("search timed out".to_string()), + _ => None, + }), }); send_progress_update(&state, completed); } @@ -759,7 +774,8 @@ fn emit_search_batch( batch.truncate(accepted as usize); } - let elements = std::mem::take(batch); + let mut elements = Vec::with_capacity(state.batch_size); + std::mem::swap(&mut elements, batch); let control = state.sink.on_elements(SearchElementBatch { file_index, file_url: file_url.to_string(), @@ -778,6 +794,15 @@ fn emit_search_batch( accepted } +fn current_exit_reason(state: &SearchWorkerState<'_>) -> Option { + match state.exit_reason.load(Ordering::Relaxed) { + 1 => Some(SearchExitReason::Cancelled), + 2 => Some(SearchExitReason::Timeout), + 3 => Some(SearchExitReason::MaxResultsReached), + _ => None, + } +} + fn reserve_result_slots(state: &SearchWorkerState<'_>, wanted: u64) -> u64 { match state.max_results { None => { @@ -1008,6 +1033,121 @@ mod tests { assert!(json.contains("percent_complete")); } + struct CountingSink { + batches: AtomicU64, + elements: AtomicU64, + } + + impl CountingSink { + fn new() -> Self { + Self { + batches: AtomicU64::new(0), + elements: AtomicU64::new(0), + } + } + } + + impl SearchSink for CountingSink { + fn on_elements(&self, batch: SearchElementBatch) -> SearchControl { + self.batches.fetch_add(1, Ordering::Relaxed); + self.elements + .fetch_add(batch.elements.len() as u64, Ordering::Relaxed); + SearchControl::Continue + } + } + + fn with_worker_state( + max_results: Option, + cancel_flag: Option<&AtomicBool>, + deadline: Option, + sink: &dyn SearchSink, + f: F, + ) where + F: FnOnce(SearchWorkerState<'_>), + { + let stop_flag = AtomicBool::new(false); + let exit_reason = AtomicU64::new(0); + let files_completed = AtomicU64::new(0); + let successful_files = AtomicU64::new(0); + let failed_files = AtomicU64::new(0); + let total_messages = AtomicU64::new(0); + + f(SearchWorkerState { + total_files: 1, + start_time: Instant::now(), + deadline, + batch_size: 4, + max_results, + external_cancel: cancel_flag, + stop_flag: &stop_flag, + exit_reason: &exit_reason, + files_completed: &files_completed, + successful_files: &successful_files, + failed_files: &failed_files, + total_messages: &total_messages, + sink, + }); + } + + #[test] + fn test_reserve_result_slots_truncates_at_max_results() { + let sink = CountingSink::new(); + with_worker_state(Some(3), None, None, &sink, |state| { + assert_eq!(reserve_result_slots(&state, 2), 2); + assert_eq!(state.total_messages.load(Ordering::Relaxed), 2); + assert_eq!(reserve_result_slots(&state, 2), 1); + assert_eq!(state.total_messages.load(Ordering::Relaxed), 3); + assert_eq!( + current_exit_reason(&state), + Some(SearchExitReason::MaxResultsReached) + ); + assert!(state.stop_flag.load(Ordering::Relaxed)); + }); + } + + #[test] + fn test_emit_search_batch_reuses_batch_capacity() { + let sink = CountingSink::new(); + with_worker_state(Some(3), None, None, &sink, |state| { + let mut batch = vec![ + BgpElem::default(), + BgpElem::default(), + BgpElem::default(), + BgpElem::default(), + ]; + let accepted = emit_search_batch(&state, 0, "file", "collector", &mut batch); + assert_eq!(accepted, 3); + assert_eq!(sink.batches.load(Ordering::Relaxed), 1); + assert_eq!(sink.elements.load(Ordering::Relaxed), 3); + assert!(batch.is_empty()); + assert!(batch.capacity() >= state.batch_size); + }); + } + + #[test] + fn test_search_should_stop_cancel_and_timeout() { + let sink = CountingSink::new(); + let cancel = AtomicBool::new(true); + with_worker_state(None, Some(&cancel), None, &sink, |state| { + assert!(search_should_stop(&state)); + assert_eq!( + current_exit_reason(&state), + Some(SearchExitReason::Cancelled) + ); + }); + + with_worker_state( + None, + None, + Some(Instant::now() - Duration::from_secs(1)), + &sink, + |state| { + assert!(search_should_stop(&state)); + assert_eq!(current_exit_reason(&state), Some(SearchExitReason::Timeout)); + }, + ); + } + #[test] fn test_search_filters_validate_with_communities() { let filters = SearchFilters { From 3729772055ad91927f560b5ec5dca7ab95e78701 Mon Sep 17 00:00:00 2001 From: Mingwei Zhang Date: Wed, 1 Jul 2026 19:51:10 -0700 Subject: [PATCH 3/5] fix: include broker query time in search timeout --- src/lens/search/mod.rs | 23 +++++++++++++++++++---- 1 file changed, 19 insertions(+), 4 deletions(-) diff --git a/src/lens/search/mod.rs b/src/lens/search/mod.rs index 6059633..61966ee 100644 --- a/src/lens/search/mod.rs +++ b/src/lens/search/mod.rs @@ -471,26 +471,43 @@ impl SearchLens { options: SearchExecutionOptions, sink: Arc, ) -> Result { + let start_time = Instant::now(); + let deadline = options.timeout.map(|timeout| start_time + timeout); + sink.on_progress(SearchProgress::QueryingBroker); let items = self.query_broker(filters)?; let total_files = items.len(); sink.on_progress(SearchProgress::FilesFound { count: total_files }); + if deadline.is_some_and(|dl| Instant::now() >= dl) { + return Ok(SearchOutcome { + summary: SearchSummary { + total_files, + successful_files: 0, + failed_files: 0, + total_messages: 0, + duration_secs: start_time.elapsed().as_secs_f64(), + }, + exit_reason: SearchExitReason::Timeout, + }); + } + if total_files == 0 { + let duration_secs = start_time.elapsed().as_secs_f64(); let summary = SearchSummary { total_files: 0, successful_files: 0, failed_files: 0, total_messages: 0, - duration_secs: 0.0, + duration_secs, }; sink.on_progress(SearchProgress::Completed { total_files: 0, successful_files: 0, failed_files: 0, total_messages: 0, - duration_secs: 0.0, + duration_secs, files_per_sec: None, }); return Ok(SearchOutcome { @@ -499,10 +516,8 @@ impl SearchLens { }); } - let start_time = Instant::now(); let batch_size = options.batch_size.max(1); let max_results = options.max_results.filter(|max| *max > 0); - let deadline = options.timeout.map(|timeout| start_time + timeout); let external_cancel = options.cancel_flag.clone(); let stop_flag = AtomicBool::new(false); let exit_reason = AtomicU64::new(0); From 2997feba4875bb28b1624f6d438fa81983aef700 Mon Sep 17 00:00:00 2001 From: Mingwei Zhang Date: Wed, 1 Jul 2026 20:47:23 -0700 Subject: [PATCH 4/5] fix: refine search concurrency and max result handling --- src/lens/search/mod.rs | 36 ++++++++++++++++++++++++++++++------ src/server/mod.rs | 20 +++++++++++++++++--- src/server/search.rs | 41 ++++++++++++++++++++++++++++------------- 3 files changed, 75 insertions(+), 22 deletions(-) diff --git a/src/lens/search/mod.rs b/src/lens/search/mod.rs index 61966ee..898d3c8 100644 --- a/src/lens/search/mod.rs +++ b/src/lens/search/mod.rs @@ -150,6 +150,8 @@ pub type ElementHandler = Arc; pub struct SearchExecutionOptions { /// Number of rayon worker threads for this search. `None` or `Some(0)` uses rayon's default. pub concurrency: Option, + /// Reusable rayon thread pool supplied by the caller. Takes precedence over `concurrency`. + pub thread_pool: Option>, /// Maximum number of elements to emit. `None` means unlimited. pub max_results: Option, /// Maximum wall-clock runtime after broker query starts. `None` means no timeout. @@ -188,6 +190,9 @@ pub struct SearchElementBatch { #[derive(Debug, Clone, Copy, PartialEq, Eq)] pub enum SearchControl { Continue, + /// Stop the search early. The shared executor reports this as + /// [`SearchExitReason::Cancelled`] because the sink is the active consumer + /// of search results and no longer wants more data. Stop, } @@ -547,12 +552,16 @@ impl SearchLens { }); }; - match options.concurrency { - Some(n) if n > 0 => { - let pool = rayon::ThreadPoolBuilder::new().num_threads(n).build()?; - pool.install(run); + if let Some(pool) = options.thread_pool.as_ref() { + pool.install(run); + } else { + match options.concurrency { + Some(n) if n > 0 => { + let pool = rayon::ThreadPoolBuilder::new().num_threads(n).build()?; + pool.install(run); + } + _ => run(), } - _ => run(), } let duration_secs = start_time.elapsed().as_secs_f64(); @@ -703,7 +712,7 @@ fn process_search_item( } batch.push(elem); - if batch.len() >= state.batch_size { + if should_emit_batch(&state, batch.len()) { let accepted = emit_search_batch(&state, index, &url, &collector, &mut batch); file_messages += accepted; } @@ -771,6 +780,20 @@ fn search_should_stop(state: &SearchWorkerState<'_>) -> bool { false } +fn should_emit_batch(state: &SearchWorkerState<'_>, batch_len: usize) -> bool { + if batch_len >= state.batch_size { + return true; + } + + if let Some(max) = state.max_results { + let emitted = state.total_messages.load(Ordering::Relaxed); + let remaining = max.saturating_sub(emitted); + return remaining == 0 || batch_len as u64 >= remaining; + } + + false +} + fn emit_search_batch( state: &SearchWorkerState<'_>, file_index: usize, @@ -1110,6 +1133,7 @@ mod tests { with_worker_state(Some(3), None, None, &sink, |state| { assert_eq!(reserve_result_slots(&state, 2), 2); assert_eq!(state.total_messages.load(Ordering::Relaxed), 2); + assert!(should_emit_batch(&state, 1)); assert_eq!(reserve_result_slots(&state, 2), 1); assert_eq!(state.total_messages.load(Ordering::Relaxed), 3); assert_eq!( diff --git a/src/server/mod.rs b/src/server/mod.rs index dfa81a5..3d37b41 100644 --- a/src/server/mod.rs +++ b/src/server/mod.rs @@ -42,7 +42,8 @@ use crate::config::MonocleConfig; #[derive(Clone)] pub struct ServerState { pub config: Arc, - pub search_permits: Arc, + pub search_permits: Option>, + pub search_pool: Option>, } // ============================================================================= @@ -57,9 +58,21 @@ pub async fn start_server(config: MonocleConfig) -> anyhow::Result<()> { let auth_token = config.server_auth_token.trim().to_string(); let max_concurrent_searches = config.server_max_concurrent_searches; + let search_concurrency = config.search_concurrency; + let search_pool = if search_concurrency > 0 { + Some(Arc::new( + rayon::ThreadPoolBuilder::new() + .num_threads(search_concurrency) + .build()?, + )) + } else { + None + }; let state = ServerState { config: Arc::new(config), - search_permits: Arc::new(Semaphore::new(max_concurrent_searches)), + search_permits: (max_concurrent_searches > 0) + .then(|| Arc::new(Semaphore::new(max_concurrent_searches))), + search_pool, }; let cors = CorsLayer::new() @@ -117,7 +130,8 @@ mod tests { let config = MonocleConfig::default(); let state = ServerState { config: Arc::new(config), - search_permits: Arc::new(Semaphore::new(3)), + search_permits: Some(Arc::new(Semaphore::new(3))), + search_pool: None, }; let _cloned = state.clone(); } diff --git a/src/server/search.rs b/src/server/search.rs index b1a006c..295b115 100644 --- a/src/server/search.rs +++ b/src/server/search.rs @@ -247,16 +247,14 @@ pub async fn stream_search( // 4. Enforce server-side concurrent search limit. Requests are rejected // immediately instead of queued so clients get deterministic feedback. - let search_permit = if config.server_max_concurrent_searches == 0 { - None - } else { - Some( - state - .search_permits + let search_permit = match state.search_permits.as_ref() { + Some(search_permits) => Some( + search_permits .clone() .try_acquire_owned() .map_err(|_| ApiError::too_many_requests("too many concurrent search requests"))?, - ) + ), + None => None, }; let concurrency = if config.search_concurrency > 0 { @@ -264,6 +262,7 @@ pub async fn stream_search( } else { None }; + let search_pool = state.search_pool.clone(); // 5. Create bounded channel and cancellation flag let (tx, rx) = mpsc::channel::(32); @@ -275,15 +274,16 @@ pub async fn stream_search( tokio::task::spawn_blocking(move || { let _search_permit = search_permit; - run_search_worker( + run_search_worker(SearchWorkerConfig { filters, batch_size, max_results, timeout_secs, concurrency, - worker_cancel_flag, - worker_tx, - ); + search_pool, + cancel_flag: worker_cancel_flag, + event_tx: worker_tx, + }); }); // 6. Build SSE stream from the channel receiver. @@ -315,15 +315,29 @@ pub async fn stream_search( // Search Worker (runs in spawn_blocking) // ============================================================================= -fn run_search_worker( +struct SearchWorkerConfig { filters: SearchFilters, batch_size: usize, max_results: Option, timeout_secs: Option, concurrency: Option, + search_pool: Option>, cancel_flag: Arc, event_tx: mpsc::Sender, -) { +} + +fn run_search_worker(config: SearchWorkerConfig) { + let SearchWorkerConfig { + filters, + batch_size, + max_results, + timeout_secs, + concurrency, + search_pool, + cancel_flag, + event_tx, + } = config; + let _ = send_event( &event_tx, SearchStreamEvent::Started(SearchStarted { @@ -336,6 +350,7 @@ fn run_search_worker( let sink = Arc::new(SseSearchSink::new(event_tx.clone(), cancel_flag.clone())); let options = SearchExecutionOptions { concurrency, + thread_pool: search_pool, max_results, timeout: timeout_secs.map(Duration::from_secs), cancel_flag: Some(cancel_flag), From c0a233869f2442a7f77a7aa4a95e0642c5242ea2 Mon Sep 17 00:00:00 2001 From: Mingwei Zhang Date: Wed, 1 Jul 2026 21:07:56 -0700 Subject: [PATCH 5/5] fix: set default search batch_size, suppress redundant SSE completion, update config example - Set SearchExecutionOptions::default() batch_size to 64 instead of 0 (which was clamped to 1), reducing per-element allocations in the library search path. - Suppress SearchProgress::Completed in SseSearchSink::on_progress to avoid sending a redundant progress-level completion event before the terminal SearchStreamEvent::Completed. - Add search_concurrency and server_max_concurrent_searches to monocle.toml.example to match documentation in config.rs and server/README.md. --- monocle.toml.example | 6 ++++++ src/lens/search/mod.rs | 18 +++++++++++++++++- src/server/search.rs | 7 +++++++ 3 files changed, 30 insertions(+), 1 deletion(-) diff --git a/monocle.toml.example b/monocle.toml.example index 6e2c5f6..0f225e2 100644 --- a/monocle.toml.example +++ b/monocle.toml.example @@ -21,9 +21,15 @@ server_max_search_batch_size = 100 # Maximum search results per request (0 = unlimited) server_max_search_results = 10000 +# Search concurrency; 0 = auto/rayon default (default: 0) +search_concurrency = 0 + # Search timeout in seconds (0 = no timeout) server_search_timeout_secs = 300 +# Maximum concurrent SSE search requests (0 = unlimited, default: 3) +server_max_concurrent_searches = 3 + # ============================================================================= # Auth (token-based, disabled by default) # ============================================================================= diff --git a/src/lens/search/mod.rs b/src/lens/search/mod.rs index 898d3c8..219956d 100644 --- a/src/lens/search/mod.rs +++ b/src/lens/search/mod.rs @@ -145,8 +145,11 @@ pub type SearchProgressCallback = Arc; /// Called for each BGP element found during search, along with the collector ID. pub type ElementHandler = Arc; +/// Default number of elements per batch when no explicit `batch_size` is set. +const DEFAULT_SEARCH_BATCH_SIZE: usize = 64; + /// Runtime options for executing a search. -#[derive(Clone, Default)] +#[derive(Clone)] pub struct SearchExecutionOptions { /// Number of rayon worker threads for this search. `None` or `Some(0)` uses rayon's default. pub concurrency: Option, @@ -162,6 +165,19 @@ pub struct SearchExecutionOptions { pub batch_size: usize, } +impl Default for SearchExecutionOptions { + fn default() -> Self { + Self { + concurrency: None, + thread_pool: None, + max_results: None, + timeout: None, + cancel_flag: None, + batch_size: DEFAULT_SEARCH_BATCH_SIZE, + } + } +} + /// Reason a search execution stopped. #[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)] pub enum SearchExitReason { diff --git a/src/server/search.rs b/src/server/search.rs index 295b115..cff31a4 100644 --- a/src/server/search.rs +++ b/src/server/search.rs @@ -416,6 +416,13 @@ impl SseSearchSink { impl SearchSink for SseSearchSink { fn on_progress(&self, progress: SearchProgress) { + // The executor emits SearchProgress::Completed as a progress event, + // but run_search_worker sends the terminal SearchStreamEvent::Completed + // (or Cancelled/Error) itself. Suppress the redundant progress-level + // Completed to avoid sending two completion events to the client. + if matches!(progress, SearchProgress::Completed { .. }) { + return; + } if send_event(&self.event_tx, SearchStreamEvent::Progress(progress)).is_err() { self.cancel_flag.store(true, Ordering::Relaxed); }