1use crate::executor::Dataset;
38use crate::query_fingerprinting::{FingerprintConfig, QueryFingerprint, QueryFingerprinter};
39use crate::system_load_monitor::AdaptiveConcurrencyController;
40use anyhow::Result;
41use scirs2_core::metrics::{Counter, Gauge, Timer};
42use std::collections::{HashMap, VecDeque};
43use std::sync::{Arc, Mutex, RwLock};
44use std::time::{Duration, Instant};
45use tokio::task::JoinHandle;
46
47#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash, Default)]
49pub enum QueryPriority {
50 High = 3,
52 #[default]
54 Normal = 2,
55 Low = 1,
57 Background = 0,
59}
60
61#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
63pub enum BatchMode {
64 Parallel,
66 Sequential,
68 #[default]
70 Optimized,
71 Adaptive,
73}
74
75#[derive(Debug, Clone)]
77pub struct BatchConfig {
78 pub max_concurrent: usize,
80 pub memory_limit_mb: usize,
82 pub cpu_limit: f64,
84 pub mode: BatchMode,
86 pub enable_grouping: bool,
88 pub enable_caching: bool,
90 pub batch_timeout: Duration,
92 pub query_timeout: Duration,
94 pub fair_scheduling: bool,
96}
97
98impl Default for BatchConfig {
99 fn default() -> Self {
100 Self {
101 max_concurrent: std::thread::available_parallelism()
102 .map(|n| n.get())
103 .unwrap_or(1),
104 memory_limit_mb: 4096,
105 cpu_limit: 0.8,
106 mode: BatchMode::default(),
107 enable_grouping: true,
108 enable_caching: true,
109 batch_timeout: Duration::from_secs(300), query_timeout: Duration::from_secs(60), fair_scheduling: true,
112 }
113 }
114}
115
116impl BatchConfig {
117 pub fn with_max_concurrent(mut self, max: usize) -> Self {
119 self.max_concurrent = max.max(1);
120 self
121 }
122
123 pub fn with_memory_limit_mb(mut self, limit: usize) -> Self {
125 self.memory_limit_mb = limit;
126 self
127 }
128
129 pub fn with_cpu_limit(mut self, limit: f64) -> Self {
131 self.cpu_limit = limit.clamp(0.0, 1.0);
132 self
133 }
134
135 pub fn with_mode(mut self, mode: BatchMode) -> Self {
137 self.mode = mode;
138 self
139 }
140
141 pub fn with_grouping(mut self, enabled: bool) -> Self {
143 self.enable_grouping = enabled;
144 self
145 }
146
147 pub fn with_caching(mut self, enabled: bool) -> Self {
149 self.enable_caching = enabled;
150 self
151 }
152
153 pub fn with_batch_timeout(mut self, timeout: Duration) -> Self {
155 self.batch_timeout = timeout;
156 self
157 }
158
159 pub fn with_query_timeout(mut self, timeout: Duration) -> Self {
161 self.query_timeout = timeout;
162 self
163 }
164}
165
166#[derive(Debug, Clone)]
168pub struct BatchQuery {
169 pub id: String,
171 pub query: String,
173 pub priority: QueryPriority,
175 pub fingerprint: Option<QueryFingerprint>,
177 pub submitted_at: Instant,
179 pub started_at: Option<Instant>,
181 pub completed_at: Option<Instant>,
183}
184
185#[derive(Debug, Clone)]
187pub struct BatchQueryResult {
188 pub id: String,
190 pub success: bool,
192 pub results: Option<String>, pub error: Option<String>,
196 pub duration: Duration,
198 pub result_count: usize,
200}
201
202#[derive(Debug, Clone)]
204pub struct BatchStatistics {
205 pub total_queries: usize,
207 pub successful_queries: usize,
209 pub failed_queries: usize,
211 pub total_duration: Duration,
213 pub avg_duration: Duration,
215 pub min_duration: Duration,
217 pub max_duration: Duration,
219 pub total_results: usize,
221 pub throughput: f64,
223 pub peak_memory_mb: f64,
225 pub avg_cpu_usage: f64,
227 pub cached_queries: usize,
229 pub query_groups: usize,
231}
232
233impl BatchStatistics {
234 pub fn new() -> Self {
236 Self {
237 total_queries: 0,
238 successful_queries: 0,
239 failed_queries: 0,
240 total_duration: Duration::from_secs(0),
241 avg_duration: Duration::from_secs(0),
242 min_duration: Duration::MAX,
243 max_duration: Duration::from_secs(0),
244 total_results: 0,
245 throughput: 0.0,
246 peak_memory_mb: 0.0,
247 avg_cpu_usage: 0.0,
248 cached_queries: 0,
249 query_groups: 0,
250 }
251 }
252
253 pub fn calculate_derived(&mut self) {
255 if self.total_queries > 0 {
256 let total_secs = self.total_duration.as_secs_f64();
257 if total_secs > 0.0 {
258 self.throughput = self.total_queries as f64 / total_secs;
259 }
260
261 if self.successful_queries > 0 {
262 self.avg_duration = self.total_duration / self.successful_queries as u32;
263 }
264 }
265 }
266
267 pub fn success_rate(&self) -> f64 {
269 if self.total_queries == 0 {
270 return 0.0;
271 }
272 self.successful_queries as f64 / self.total_queries as f64
273 }
274
275 pub fn cache_hit_rate(&self) -> f64 {
277 if self.total_queries == 0 {
278 return 0.0;
279 }
280 self.cached_queries as f64 / self.total_queries as f64
281 }
282}
283
284impl Default for BatchStatistics {
285 fn default() -> Self {
286 Self::new()
287 }
288}
289
290pub struct QueryBatchExecutor {
292 config: BatchConfig,
294 queue: Arc<Mutex<VecDeque<BatchQuery>>>,
296 #[allow(dead_code)]
298 results: Arc<RwLock<HashMap<String, BatchQueryResult>>>,
299 fingerprinter: QueryFingerprinter,
301 cache: Arc<RwLock<HashMap<String, String>>>, stats: Arc<RwLock<BatchStatistics>>,
305 queries_executed: Counter,
307 queries_failed: Counter,
308 batch_duration: Timer,
309 active_queries: Gauge,
310}
311
312impl QueryBatchExecutor {
313 pub fn new(config: BatchConfig) -> Self {
315 Self {
316 config,
317 queue: Arc::new(Mutex::new(VecDeque::new())),
318 results: Arc::new(RwLock::new(HashMap::new())),
319 fingerprinter: QueryFingerprinter::new(FingerprintConfig::default()),
320 cache: Arc::new(RwLock::new(HashMap::new())),
321 stats: Arc::new(RwLock::new(BatchStatistics::new())),
322 queries_executed: Counter::new("batch_queries_executed".to_string()),
323 queries_failed: Counter::new("batch_queries_failed".to_string()),
324 batch_duration: Timer::new("batch_execution_duration".to_string()),
325 active_queries: Gauge::new("batch_active_queries".to_string()),
326 }
327 }
328
329 pub fn add_query(&self, query: impl Into<String>, priority: QueryPriority) -> Result<String> {
331 let query = query.into();
332 let id = format!("query_{}", uuid::Uuid::new_v4());
333
334 let fingerprint = if self.config.enable_grouping {
336 Some(self.fingerprinter.fingerprint(&query)?)
337 } else {
338 None
339 };
340
341 let batch_query = BatchQuery {
342 id: id.clone(),
343 query,
344 priority,
345 fingerprint,
346 submitted_at: Instant::now(),
347 started_at: None,
348 completed_at: None,
349 };
350
351 let mut queue = self.queue.lock().expect("lock poisoned");
352
353 if self.config.fair_scheduling {
355 let insert_pos = queue
357 .iter()
358 .rposition(|q| q.priority >= priority)
359 .map(|pos| pos + 1)
360 .unwrap_or(0);
361 queue.insert(insert_pos, batch_query);
362 } else {
363 let insert_pos = queue
365 .iter()
366 .position(|q| q.priority < priority)
367 .unwrap_or(queue.len());
368 queue.insert(insert_pos, batch_query);
369 }
370
371 Ok(id)
372 }
373
374 pub fn add_queries(&self, queries: Vec<(String, QueryPriority)>) -> Result<Vec<String>> {
376 queries
377 .into_iter()
378 .map(|(q, p)| self.add_query(q, p))
379 .collect()
380 }
381
382 pub fn queue_size(&self) -> usize {
384 self.queue.lock().expect("lock poisoned").len()
385 }
386
387 pub fn clear_queue(&self) {
389 self.queue.lock().expect("lock poisoned").clear();
390 }
391
392 pub fn statistics(&self) -> BatchStatistics {
394 self.stats.read().expect("lock poisoned").clone()
395 }
396
397 pub async fn execute_batch_async<D: Dataset + Send + Sync + 'static>(
399 &self,
400 dataset: Arc<D>,
401 ) -> Result<Vec<BatchQueryResult>> {
402 let start_time = Instant::now();
403
404 let queries: Vec<BatchQuery> = {
406 let mut queue = self.queue.lock().expect("lock poisoned");
407 queue.drain(..).collect()
408 };
409
410 if queries.is_empty() {
411 return Ok(Vec::new());
412 }
413
414 {
416 let mut stats = self.stats.write().expect("lock poisoned");
417 stats.total_queries = queries.len();
418 }
419
420 let results = match self.config.mode {
422 BatchMode::Parallel => self.execute_parallel(queries, dataset).await?,
423 BatchMode::Sequential => self.execute_sequential(queries, dataset).await?,
424 BatchMode::Optimized => self.execute_optimized(queries, dataset).await?,
425 BatchMode::Adaptive => self.execute_adaptive(queries, dataset).await?,
426 };
427
428 let duration = start_time.elapsed();
430 {
431 let mut stats = self.stats.write().expect("lock poisoned");
432 stats.total_duration = duration;
433 stats.calculate_derived();
434 }
435
436 self.batch_duration.observe(duration);
437
438 Ok(results)
439 }
440
441 async fn execute_parallel<D: Dataset + Send + Sync + 'static>(
443 &self,
444 queries: Vec<BatchQuery>,
445 dataset: Arc<D>,
446 ) -> Result<Vec<BatchQueryResult>> {
447 let semaphore = Arc::new(tokio::sync::Semaphore::new(self.config.max_concurrent));
448 let mut handles: Vec<JoinHandle<BatchQueryResult>> = Vec::new();
449
450 for query in queries {
451 let permit = semaphore.clone().acquire_owned().await?;
452 let dataset = dataset.clone();
453 let timeout = self.config.query_timeout;
454 let cache = self.cache.clone();
455 let enable_caching = self.config.enable_caching;
456 let fingerprint = query.fingerprint.clone();
457
458 self.active_queries.inc();
459
460 let handle = tokio::spawn(async move {
461 let result = Self::execute_single_query(
462 query,
463 dataset,
464 timeout,
465 cache,
466 enable_caching,
467 fingerprint,
468 )
469 .await;
470 drop(permit);
471 result
472 });
473
474 handles.push(handle);
475 }
476
477 let mut results = Vec::new();
479 for handle in handles {
480 match handle.await {
481 Ok(result) => {
482 self.update_stats(&result);
483 results.push(result);
484 }
485 Err(e) => {
486 eprintln!("Task failed: {}", e);
487 self.queries_failed.inc();
488 }
489 }
490 self.active_queries.dec();
491 }
492
493 Ok(results)
494 }
495
496 async fn execute_sequential<D: Dataset + Send + Sync + 'static>(
498 &self,
499 queries: Vec<BatchQuery>,
500 dataset: Arc<D>,
501 ) -> Result<Vec<BatchQueryResult>> {
502 let mut results = Vec::new();
503
504 for query in queries {
505 self.active_queries.inc();
506
507 let result = Self::execute_single_query(
508 query,
509 dataset.clone(),
510 self.config.query_timeout,
511 self.cache.clone(),
512 self.config.enable_caching,
513 None,
514 )
515 .await;
516
517 self.update_stats(&result);
518 results.push(result);
519
520 self.active_queries.dec();
521 }
522
523 Ok(results)
524 }
525
526 async fn execute_optimized<D: Dataset + Send + Sync + 'static>(
528 &self,
529 queries: Vec<BatchQuery>,
530 dataset: Arc<D>,
531 ) -> Result<Vec<BatchQueryResult>> {
532 let mut groups: HashMap<String, Vec<BatchQuery>> = HashMap::new();
534
535 for query in queries {
536 let key = query
537 .fingerprint
538 .as_ref()
539 .map(|f| f.hash.clone())
540 .unwrap_or_else(|| query.id.clone());
541
542 groups.entry(key).or_default().push(query);
543 }
544
545 {
547 let mut stats = self.stats.write().expect("lock poisoned");
548 stats.query_groups = groups.len();
549 }
550
551 let mut all_results = Vec::new();
553
554 for (_key, group) in groups {
555 let group_results = self.execute_parallel(group, dataset.clone()).await?;
556 all_results.extend(group_results);
557 }
558
559 Ok(all_results)
560 }
561
562 async fn execute_adaptive<D: Dataset + Send + Sync + 'static>(
564 &self,
565 queries: Vec<BatchQuery>,
566 dataset: Arc<D>,
567 ) -> Result<Vec<BatchQueryResult>> {
568 use tokio::sync::Semaphore;
569
570 let controller = Arc::new(
572 AdaptiveConcurrencyController::new(self.config.max_concurrent)
573 .with_thresholds(0.75, 0.40) .with_adjustment_interval(Duration::from_secs(5)),
575 );
576
577 let initial_permits = controller.current_concurrency();
579 let semaphore = Arc::new(Semaphore::new(initial_permits));
580
581 let controller_clone = Arc::clone(&controller);
583 let semaphore_clone = Arc::clone(&semaphore);
584 let adjustment_task = tokio::spawn(async move {
585 loop {
586 tokio::time::sleep(Duration::from_secs(5)).await;
587
588 controller_clone.update_concurrency();
590 let new_concurrency = controller_clone.current_concurrency();
591
592 let current_permits = semaphore_clone.available_permits();
594 if new_concurrency > current_permits {
595 semaphore_clone.add_permits(new_concurrency - current_permits);
597 }
598 }
600 });
601
602 let mut tasks = Vec::new();
604
605 for query in queries {
606 let dataset_clone = Arc::clone(&dataset);
607 let cache_clone = Arc::clone(&self.cache);
608 let timeout = self.config.query_timeout;
609 let enable_caching = self.config.enable_caching;
610 let fingerprint = query.fingerprint.clone();
611 let semaphore_clone = Arc::clone(&semaphore);
612
613 let task = tokio::spawn(async move {
614 let _permit = semaphore_clone.acquire().await.expect("Semaphore closed");
616
617 Self::execute_single_query(
619 query,
620 dataset_clone,
621 timeout,
622 cache_clone,
623 enable_caching,
624 fingerprint,
625 )
626 .await
627 });
628
629 tasks.push(task);
630 }
631
632 let mut results = Vec::with_capacity(tasks.len());
634 for task in tasks {
635 match task.await {
636 Ok(result) => results.push(result),
637 Err(e) => {
638 eprintln!("Task execution error: {}", e);
640 }
641 }
642 }
643
644 adjustment_task.abort();
646
647 Ok(results)
648 }
649
650 async fn execute_single_query<D: Dataset + Send + Sync + 'static>(
652 mut query: BatchQuery,
653 _dataset: Arc<D>,
654 timeout: Duration,
655 cache: Arc<RwLock<HashMap<String, String>>>,
656 enable_caching: bool,
657 _fingerprint: Option<QueryFingerprint>,
658 ) -> BatchQueryResult {
659 query.started_at = Some(Instant::now());
660 let start = Instant::now();
661
662 if enable_caching {
664 if let Some(fp) = &query.fingerprint {
665 if let Some(cached) = cache.read().expect("lock poisoned").get(&fp.hash) {
666 query.completed_at = Some(Instant::now());
667 return BatchQueryResult {
668 id: query.id,
669 success: true,
670 results: Some(cached.clone()),
671 error: None,
672 duration: start.elapsed(),
673 result_count: cached.lines().count(),
674 };
675 }
676 }
677 }
678
679 let result = tokio::time::timeout(timeout, async {
681 tokio::time::sleep(Duration::from_millis(10)).await;
684 Ok::<String, anyhow::Error>(format!("Results for: {}", query.query))
685 })
686 .await;
687
688 query.completed_at = Some(Instant::now());
689 let duration = start.elapsed();
690
691 match result {
692 Ok(Ok(results)) => {
693 if enable_caching {
695 if let Some(fp) = &query.fingerprint {
696 cache
697 .write()
698 .expect("lock poisoned")
699 .insert(fp.hash.clone(), results.clone());
700 }
701 }
702
703 BatchQueryResult {
704 id: query.id,
705 success: true,
706 results: Some(results.clone()),
707 error: None,
708 duration,
709 result_count: results.lines().count(),
710 }
711 }
712 Ok(Err(e)) => BatchQueryResult {
713 id: query.id,
714 success: false,
715 results: None,
716 error: Some(e.to_string()),
717 duration,
718 result_count: 0,
719 },
720 Err(_) => BatchQueryResult {
721 id: query.id,
722 success: false,
723 results: None,
724 error: Some("Query timeout".to_string()),
725 duration,
726 result_count: 0,
727 },
728 }
729 }
730
731 fn update_stats(&self, result: &BatchQueryResult) {
733 let mut stats = self.stats.write().expect("lock poisoned");
734
735 if result.success {
736 stats.successful_queries += 1;
737 self.queries_executed.inc();
738
739 stats.total_results += result.result_count;
740
741 if result.duration < stats.min_duration {
742 stats.min_duration = result.duration;
743 }
744 if result.duration > stats.max_duration {
745 stats.max_duration = result.duration;
746 }
747 } else {
748 stats.failed_queries += 1;
749 self.queries_failed.inc();
750 }
751 }
752}
753
754#[cfg(test)]
755mod tests {
756 use super::*;
757
758 #[test]
759 fn test_batch_config_builder() {
760 let config = BatchConfig::default()
761 .with_max_concurrent(32)
762 .with_memory_limit_mb(8192)
763 .with_cpu_limit(0.9)
764 .with_mode(BatchMode::Parallel);
765
766 assert_eq!(config.max_concurrent, 32);
767 assert_eq!(config.memory_limit_mb, 8192);
768 assert_eq!(config.cpu_limit, 0.9);
769 assert_eq!(config.mode, BatchMode::Parallel);
770 }
771
772 #[test]
773 fn test_priority_ordering() {
774 assert!(QueryPriority::High > QueryPriority::Normal);
775 assert!(QueryPriority::Normal > QueryPriority::Low);
776 assert!(QueryPriority::Low > QueryPriority::Background);
777 }
778
779 #[test]
780 fn test_batch_statistics() {
781 let mut stats = BatchStatistics::new();
782 stats.total_queries = 100;
783 stats.successful_queries = 95;
784 stats.failed_queries = 5;
785 stats.total_duration = Duration::from_secs(10);
786 stats.cached_queries = 20;
787
788 stats.calculate_derived();
789
790 assert_eq!(stats.success_rate(), 0.95);
791 assert_eq!(stats.cache_hit_rate(), 0.2);
792 assert_eq!(stats.throughput, 10.0); }
794
795 #[test]
796 fn test_add_query() {
797 let executor = QueryBatchExecutor::new(BatchConfig::default());
798
799 let id1 = executor
800 .add_query("SELECT * WHERE { ?s ?p ?o }", QueryPriority::Normal)
801 .unwrap();
802 let id2 = executor
803 .add_query("ASK { ?s a :Person }", QueryPriority::High)
804 .unwrap();
805
806 assert_eq!(executor.queue_size(), 2);
807 assert_ne!(id1, id2);
808
809 let queue = executor.queue.lock().expect("lock should not be poisoned");
811 assert_eq!(queue[0].priority, QueryPriority::High);
812 assert_eq!(queue[1].priority, QueryPriority::Normal);
813 }
814
815 #[test]
816 fn test_add_multiple_queries() {
817 let executor = QueryBatchExecutor::new(BatchConfig::default());
818
819 let queries = vec![
820 (
821 "SELECT ?s WHERE { ?s ?p ?o }".to_string(),
822 QueryPriority::Normal,
823 ),
824 (
825 "SELECT ?p WHERE { ?s ?p ?o }".to_string(),
826 QueryPriority::Low,
827 ),
828 (
829 "SELECT ?o WHERE { ?s ?p ?o }".to_string(),
830 QueryPriority::High,
831 ),
832 ];
833
834 let ids = executor.add_queries(queries).unwrap();
835 assert_eq!(ids.len(), 3);
836 assert_eq!(executor.queue_size(), 3);
837 }
838
839 #[test]
840 fn test_clear_queue() {
841 let executor = QueryBatchExecutor::new(BatchConfig::default());
842
843 executor
844 .add_query("SELECT * WHERE { ?s ?p ?o }", QueryPriority::Normal)
845 .unwrap();
846 executor
847 .add_query("ASK { ?s a :Person }", QueryPriority::High)
848 .unwrap();
849
850 assert_eq!(executor.queue_size(), 2);
851
852 executor.clear_queue();
853 assert_eq!(executor.queue_size(), 0);
854 }
855
856 #[test]
857 fn test_batch_modes() {
858 let modes = vec![
859 BatchMode::Parallel,
860 BatchMode::Sequential,
861 BatchMode::Optimized,
862 BatchMode::Adaptive,
863 ];
864
865 for mode in modes {
866 let config = BatchConfig::default().with_mode(mode);
867 assert_eq!(config.mode, mode);
868 }
869 }
870
871 #[test]
872 fn test_fair_scheduling() {
873 let config = BatchConfig {
874 fair_scheduling: true,
875 ..Default::default()
876 };
877
878 let executor = QueryBatchExecutor::new(config);
879
880 executor.add_query("Q1", QueryPriority::Normal).unwrap();
882 executor.add_query("Q2", QueryPriority::High).unwrap();
883 executor.add_query("Q3", QueryPriority::Normal).unwrap();
884 executor.add_query("Q4", QueryPriority::High).unwrap();
885
886 let queue = executor.queue.lock().expect("lock should not be poisoned");
887
888 assert_eq!(queue[0].priority, QueryPriority::High);
890 assert_eq!(queue[1].priority, QueryPriority::High);
891 assert_eq!(queue[2].priority, QueryPriority::Normal);
892 assert_eq!(queue[3].priority, QueryPriority::Normal);
893 }
894
895 #[test]
896 fn test_config_limits() {
897 let config = BatchConfig::default()
898 .with_max_concurrent(0) .with_cpu_limit(1.5); assert_eq!(config.max_concurrent, 1);
902 assert_eq!(config.cpu_limit, 1.0);
903 }
904
905 #[test]
906 fn test_batch_query_timing() {
907 let query = BatchQuery {
908 id: "test".to_string(),
909 query: "SELECT * WHERE { ?s ?p ?o }".to_string(),
910 priority: QueryPriority::Normal,
911 fingerprint: None,
912 submitted_at: Instant::now(),
913 started_at: None,
914 completed_at: None,
915 };
916
917 assert!(query.started_at.is_none());
918 assert!(query.completed_at.is_none());
919 }
920}