1use std::collections::HashMap;
46use std::hash::Hash;
47use std::sync::RwLock;
48
49#[derive(Debug, Clone)]
55pub struct GraphEdge {
56 pub parent_field: String,
58 pub relation: String,
60 pub sub_graph: Option<Box<EntityGraph>>,
62}
63
64#[derive(Debug, Clone, Default)]
78pub struct EntityGraph {
79 edges: Vec<GraphEdge>,
81}
82
83impl EntityGraph {
84 pub fn new() -> Self {
86 Self { edges: Vec::new() }
87 }
88
89 pub fn add_edge(
91 &mut self,
92 parent_field: impl Into<String>,
93 relation: impl Into<String>,
94 ) -> &mut Self {
95 self.edges.push(GraphEdge {
96 parent_field: parent_field.into(),
97 relation: relation.into(),
98 sub_graph: None,
99 });
100 self
101 }
102
103 pub fn add_edge_with_graph(
105 &mut self,
106 parent_field: impl Into<String>,
107 relation: impl Into<String>,
108 sub_graph: EntityGraph,
109 ) -> &mut Self {
110 self.edges.push(GraphEdge {
111 parent_field: parent_field.into(),
112 relation: relation.into(),
113 sub_graph: Some(Box::new(sub_graph)),
114 });
115 self
116 }
117
118 pub fn edges(&self) -> &[GraphEdge] {
120 &self.edges
121 }
122
123 pub fn edge_count(&self) -> usize {
125 self.edges.len()
126 }
127
128 pub fn relations_of(&self, parent_field: &str) -> Vec<&GraphEdge> {
130 self.edges
131 .iter()
132 .filter(|e| e.parent_field == parent_field)
133 .collect()
134 }
135
136 pub fn all_relations(&self) -> Vec<String> {
138 let mut rels: Vec<String> = self.edges.iter().map(|e| e.relation.clone()).collect();
139 rels.sort();
140 rels.dedup();
141 rels
142 }
143
144 pub fn all_parent_fields(&self) -> Vec<String> {
146 let mut fields: Vec<String> = self.edges.iter().map(|e| e.parent_field.clone()).collect();
147 fields.sort();
148 fields.dedup();
149 fields
150 }
151
152 pub fn is_empty(&self) -> bool {
154 self.edges.is_empty()
155 }
156
157 pub fn all_relations_recursive(&self) -> Vec<String> {
159 let mut result = Vec::new();
160 for edge in &self.edges {
161 result.push(edge.relation.clone());
162 if let Some(sub) = &edge.sub_graph {
163 result.extend(sub.all_relations_recursive());
164 }
165 }
166 result.sort();
167 result.dedup();
168 result
169 }
170
171 pub fn detect_cycles(&self) -> Result<(), Vec<String>> {
211 let mut adj: std::collections::HashMap<String, Vec<String>> =
213 std::collections::HashMap::new();
214 self.collect_edges_recursive(&mut adj);
215
216 for neighbors in adj.values_mut() {
218 neighbors.sort();
219 }
220
221 let mut visited = std::collections::HashSet::new();
223 let mut visiting = std::collections::HashSet::new();
224 let mut path = Vec::new();
225
226 let mut sorted_nodes: Vec<String> = adj.keys().cloned().collect();
228 sorted_nodes.sort();
229 for node in &sorted_nodes {
230 if !visited.contains(node) {
231 dfs_cycle_detect(node, &adj, &mut visited, &mut visiting, &mut path)?;
232 }
233 }
234 Ok(())
235 }
236
237 fn collect_edges_recursive(&self, adj: &mut std::collections::HashMap<String, Vec<String>>) {
242 for edge in &self.edges {
243 adj.entry(edge.parent_field.clone())
244 .or_default()
245 .push(edge.relation.clone());
246 if let Some(sub) = &edge.sub_graph {
247 sub.collect_edges_recursive(adj);
248 }
249 }
250 }
251
252 pub fn detect_duplicate_edges(&self) -> Result<(), Vec<(String, String)>> {
258 let mut seen = std::collections::HashSet::new();
259 let mut duplicates = Vec::new();
260 for edge in &self.edges {
261 let key = (edge.parent_field.clone(), edge.relation.clone());
262 if !seen.insert(key.clone()) {
263 duplicates.push((edge.parent_field.clone(), edge.relation.clone()));
264 }
265 }
266 if duplicates.is_empty() {
267 Ok(())
268 } else {
269 Err(duplicates)
270 }
271 }
272
273 pub fn validate(&self) -> Result<(), String> {
277 if let Err(cycle) = self.detect_cycles() {
279 return Err(format!(
280 "EntityGraph 循环引用检测失败:{}",
281 cycle.join(" → ")
282 ));
283 }
284 if let Err(duplicates) = self.detect_duplicate_edges() {
286 let dup_str: Vec<String> = duplicates
287 .iter()
288 .map(|(p, r)| format!("({}->{})", p, r))
289 .collect();
290 return Err(format!(
291 "EntityGraph 重复边检测失败:{}",
292 dup_str.join(", ")
293 ));
294 }
295 Ok(())
296 }
297}
298
299fn dfs_cycle_detect(
304 node: &str,
305 adj: &std::collections::HashMap<String, Vec<String>>,
306 visited: &mut std::collections::HashSet<String>,
307 visiting: &mut std::collections::HashSet<String>,
308 path: &mut Vec<String>,
309) -> Result<(), Vec<String>> {
310 if visiting.contains(node) {
312 let cycle_start = path.iter().position(|n| n == node).unwrap_or(0);
313 let mut cycle = path[cycle_start..].to_vec();
314 cycle.push(node.to_string());
315 return Err(cycle);
316 }
317 if visited.contains(node) {
319 return Ok(());
320 }
321
322 visiting.insert(node.to_string());
324 path.push(node.to_string());
325
326 if let Some(neighbors) = adj.get(node) {
328 for neighbor in neighbors {
329 dfs_cycle_detect(neighbor, adj, visited, visiting, path)?;
330 }
331 }
332
333 visiting.remove(node);
335 visited.insert(node.to_string());
336 path.pop();
337 Ok(())
338}
339
340#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Default)]
346pub enum BatchStrategy {
347 #[default]
351 In,
352 Join,
356 Subquery,
360}
361
362impl BatchStrategy {
363 pub fn name(&self) -> &'static str {
365 match self {
366 BatchStrategy::In => "in",
367 BatchStrategy::Join => "join",
368 BatchStrategy::Subquery => "subquery",
369 }
370 }
371
372 pub fn render_in_clause(column: &str, placeholders: usize) -> String {
376 if placeholders == 0 {
377 return format!("{} IN ()", column);
378 }
379 let marks: Vec<&str> = vec!["?"; placeholders];
380 format!("{} IN ({})", column, marks.join(", "))
381 }
382}
383
384#[derive(Debug, Clone, Copy)]
392pub struct BatchSizeConfig {
393 pub size: usize,
395 pub strategy: BatchStrategy,
397}
398
399impl Default for BatchSizeConfig {
400 fn default() -> Self {
401 Self {
402 size: 100,
403 strategy: BatchStrategy::In,
404 }
405 }
406}
407
408impl BatchSizeConfig {
409 pub fn new(size: usize, strategy: BatchStrategy) -> Self {
411 Self { size, strategy }
412 }
413
414 pub fn with_size(size: usize) -> Self {
416 Self {
417 size,
418 strategy: BatchStrategy::In,
419 }
420 }
421
422 pub fn batch_count(&self, total: usize) -> usize {
437 if total == 0 {
438 0
439 } else {
440 total.div_ceil(self.size)
441 }
442 }
443
444 pub fn batch_range(&self, batch_index: usize, total: usize) -> std::ops::Range<usize> {
457 let start = batch_index * self.size;
458 let end = (start + self.size).min(total);
459 start..end
460 }
461}
462
463pub type BatchLoaderFn<K, V> = Box<dyn Fn(&[K]) -> HashMap<K, V> + Send + Sync>;
469
470pub struct BatchLoader<K, V>
494where
495 K: Hash + Eq + Clone + Send + Sync,
496 V: Clone + Send + Sync,
497{
498 batch_size: usize,
500 loader: BatchLoaderFn<K, V>,
502 cache: RwLock<HashMap<K, V>>,
504}
505
506impl<K, V> BatchLoader<K, V>
507where
508 K: Hash + Eq + Clone + Send + Sync,
509 V: Clone + Send + Sync,
510{
511 pub fn new(batch_size: usize, loader: BatchLoaderFn<K, V>) -> Self {
517 Self {
518 batch_size,
519 loader,
520 cache: RwLock::new(HashMap::new()),
521 }
522 }
523
524 pub fn load_many(&self, keys: &[K]) -> HashMap<K, V> {
530 let mut result: HashMap<K, V> = HashMap::new();
531
532 let mut to_load: Vec<K> = Vec::new();
534 if let Ok(cached) = self.cache.read() {
535 for k in keys {
536 if let Some(v) = cached.get(k) {
537 result.insert(k.clone(), v.clone());
538 } else {
539 to_load.push(k.clone());
540 }
541 }
542 } else {
543 to_load.extend(keys.iter().cloned());
544 }
545
546 if to_load.is_empty() {
547 return result;
548 }
549
550 let batch_size = self.batch_size.max(1);
552 let mut all_loaded: HashMap<K, V> = HashMap::new();
553 for chunk in to_load.chunks(batch_size) {
554 let loaded = (self.loader)(chunk);
555 all_loaded.extend(loaded);
556 }
557
558 if let Ok(mut cache) = self.cache.write() {
560 for (k, v) in &all_loaded {
561 cache.insert(k.clone(), v.clone());
562 }
563 }
564
565 result.extend(all_loaded);
567 result
568 }
569
570 pub fn load_one(&self, key: &K) -> Option<V> {
572 let result = self.load_many(std::slice::from_ref(key));
573 result.get(key).cloned()
574 }
575
576 pub fn clear_cache(&self) {
578 if let Ok(mut cache) = self.cache.write() {
579 cache.clear();
580 }
581 }
582
583 pub fn cache_size(&self) -> usize {
585 match self.cache.read() {
586 Ok(g) => g.len(),
587 Err(_) => 0,
588 }
589 }
590
591 pub fn batch_size(&self) -> usize {
593 self.batch_size
594 }
595}
596
597pub struct N1QueryDetector {
642 config: N1DetectionConfig,
644 counts: RwLock<HashMap<String, u64>>,
646 batch_counts: RwLock<HashMap<String, u64>>,
648 window_active: RwLock<bool>,
650 alerts: RwLock<Vec<N1Alert>>,
652 #[cfg(feature = "prod-n1-tuning")]
654 trigger_count: std::sync::atomic::AtomicU64,
655 #[cfg(feature = "prod-n1-tuning")]
657 block_count: std::sync::atomic::AtomicU64,
658}
659
660#[derive(Debug, Clone)]
662pub struct N1DetectionConfig {
663 pub threshold: u64,
665 pub enabled: bool,
667 pub window: std::time::Duration,
669 pub block: bool,
671}
672
673impl Default for N1DetectionConfig {
674 fn default() -> Self {
675 Self {
676 threshold: 5,
677 enabled: true,
678 window: std::time::Duration::from_secs(1),
679 block: false,
680 }
681 }
682}
683
684impl N1DetectionConfig {
685 pub fn new() -> Self {
687 Self::default()
688 }
689
690 pub fn with_threshold(mut self, threshold: u64) -> Self {
692 self.threshold = threshold.max(1);
693 self
694 }
695
696 pub fn with_enabled(mut self, enabled: bool) -> Self {
698 self.enabled = enabled;
699 self
700 }
701
702 #[cfg(feature = "prod-n1-tuning")]
704 pub fn with_window(mut self, window: std::time::Duration) -> Self {
705 self.window = window;
706 self
707 }
708
709 #[cfg(feature = "prod-n1-tuning")]
711 pub fn with_block(mut self, block: bool) -> Self {
712 self.block = block;
713 self
714 }
715}
716
717#[derive(Debug, Clone, PartialEq, Eq)]
719pub struct N1Alert {
720 pub relation: String,
722 pub query_count: u64,
724 pub batch_count: u64,
726 pub threshold: u64,
728}
729
730impl N1Alert {
731 pub fn no_batch_used(&self) -> bool {
733 self.batch_count == 0
734 }
735
736 pub fn suggested_batch_size(&self) -> usize {
738 let n = self.query_count as usize;
739 if n <= 50 {
740 50
741 } else if n <= 100 {
742 100
743 } else if n <= 500 {
744 500
745 } else {
746 1000
747 }
748 }
749}
750
751impl N1QueryDetector {
752 pub fn new(config: N1DetectionConfig) -> Self {
754 Self {
755 config,
756 counts: RwLock::new(HashMap::new()),
757 batch_counts: RwLock::new(HashMap::new()),
758 window_active: RwLock::new(false),
759 alerts: RwLock::new(Vec::new()),
760 #[cfg(feature = "prod-n1-tuning")]
761 trigger_count: std::sync::atomic::AtomicU64::new(0),
762 #[cfg(feature = "prod-n1-tuning")]
763 block_count: std::sync::atomic::AtomicU64::new(0),
764 }
765 }
766
767 #[cfg(feature = "prod-n1-tuning")]
769 pub fn stats(&self) -> N1DetectorStats {
770 N1DetectorStats {
771 trigger_count: self
772 .trigger_count
773 .load(std::sync::atomic::Ordering::Relaxed),
774 block_count: self.block_count.load(std::sync::atomic::Ordering::Relaxed),
775 }
776 }
777
778 pub fn with_defaults() -> Self {
780 Self::new(N1DetectionConfig::default())
781 }
782
783 pub fn is_enabled(&self) -> bool {
785 self.config.enabled
786 }
787
788 pub fn threshold(&self) -> u64 {
790 self.config.threshold
791 }
792
793 pub fn start_window(&self) {
797 if !self.config.enabled {
798 return;
799 }
800 if let Ok(mut counts) = self.counts.write() {
801 *counts = HashMap::new();
802 }
803 if let Ok(mut batch_counts) = self.batch_counts.write() {
804 *batch_counts = HashMap::new();
805 }
806 if let Ok(mut alerts) = self.alerts.write() {
807 *alerts = Vec::new();
808 }
809 if let Ok(mut window_active) = self.window_active.write() {
810 *window_active = true;
811 }
812 }
813
814 pub fn end_window(&self) -> Vec<N1Alert> {
819 if !self.config.enabled {
820 return Vec::new();
821 }
822 if let Ok(mut window_active) = self.window_active.write() {
823 *window_active = false;
824 }
825
826 let new_alerts: Vec<N1Alert> = match (self.counts.read(), self.batch_counts.read()) {
828 (Ok(counts), Ok(batch_counts)) => {
829 let mut alerts: Vec<N1Alert> = counts
830 .iter()
831 .filter_map(|(rel, &cnt)| {
832 if cnt >= self.config.threshold {
833 Some(N1Alert {
834 relation: rel.clone(),
835 query_count: cnt,
836 batch_count: *batch_counts.get(rel).unwrap_or(&0),
837 threshold: self.config.threshold,
838 })
839 } else {
840 None
841 }
842 })
843 .collect();
844 alerts.sort_by(|a, b| a.relation.cmp(&b.relation));
846 alerts
847 }
848 _ => Vec::new(),
849 };
850
851 if let Ok(mut alerts) = self.alerts.write() {
852 *alerts = new_alerts.clone();
853 }
854 new_alerts
855 }
856
857 pub fn record_single_load(&self, relation: &str) {
859 if !self.config.enabled {
860 return;
861 }
862 {
863 let active = self.window_active.read().map(|g| *g).unwrap_or(false);
864 if !active {
865 return;
866 }
867 }
868 if let Ok(mut counts) = self.counts.write() {
869 *counts.entry(relation.to_string()).or_insert(0) += 1;
870 }
871 }
872
873 pub fn record_batch_load(&self, relation: &str, _keys_count: usize) {
878 if !self.config.enabled {
879 return;
880 }
881 {
882 let active = self.window_active.read().map(|g| *g).unwrap_or(false);
883 if !active {
884 return;
885 }
886 }
887 if let Ok(mut batch_counts) = self.batch_counts.write() {
888 *batch_counts.entry(relation.to_string()).or_insert(0) += 1;
889 }
890 }
891
892 pub fn alerts(&self) -> Vec<N1Alert> {
894 self.alerts.read().map(|g| g.clone()).unwrap_or_default()
895 }
896
897 pub fn current_count(&self, relation: &str) -> u64 {
899 self.counts
900 .read()
901 .map(|g| g.get(relation).copied().unwrap_or(0))
902 .unwrap_or(0)
903 }
904
905 pub fn current_batch_count(&self, relation: &str) -> u64 {
907 self.batch_counts
908 .read()
909 .map(|g| g.get(relation).copied().unwrap_or(0))
910 .unwrap_or(0)
911 }
912
913 pub fn is_window_active(&self) -> bool {
915 self.window_active.read().map(|g| *g).unwrap_or(false)
916 }
917
918 pub fn has_n_plus_one(&self) -> bool {
920 !self.alerts().is_empty()
921 }
922}
923
924impl Default for N1QueryDetector {
925 fn default() -> Self {
926 Self::with_defaults()
927 }
928}
929
930#[cfg(feature = "prod-n1-tuning")]
932#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)]
933pub struct N1DetectorStats {
934 pub trigger_count: u64,
936 pub block_count: u64,
938}
939
940#[cfg(test)]
944mod tests {
945 use super::*;
946
947 #[test]
950 fn test_new_graph_is_empty() {
951 let g = EntityGraph::new();
952 assert!(g.is_empty());
953 assert_eq!(g.edge_count(), 0);
954 }
955
956 #[test]
957 fn test_add_edge() {
958 let mut g = EntityGraph::new();
959 g.add_edge("user", "posts");
960 assert_eq!(g.edge_count(), 1);
961 assert!(!g.is_empty());
962 }
963
964 #[test]
965 fn test_add_multiple_edges() {
966 let mut g = EntityGraph::new();
967 g.add_edge("user", "posts")
968 .add_edge("user", "profile")
969 .add_edge("user", "comments");
970 assert_eq!(g.edge_count(), 3);
971 }
972
973 #[test]
974 fn test_add_edge_with_sub_graph() {
975 let mut sub = EntityGraph::new();
976 sub.add_edge("comments", "author");
977
978 let mut g = EntityGraph::new();
979 g.add_edge_with_graph("user", "posts", sub);
980
981 assert_eq!(g.edge_count(), 1);
982 assert!(g.edges()[0].sub_graph.is_some());
983 assert_eq!(g.edges()[0].sub_graph.as_ref().unwrap().edge_count(), 1);
984 }
985
986 #[test]
987 fn test_relations_of() {
988 let mut g = EntityGraph::new();
989 g.add_edge("user", "posts")
990 .add_edge("user", "profile")
991 .add_edge("post", "comments");
992
993 let user_relations = g.relations_of("user");
994 assert_eq!(user_relations.len(), 2);
995 assert_eq!(user_relations[0].relation, "posts");
996 assert_eq!(user_relations[1].relation, "profile");
997
998 let post_relations = g.relations_of("post");
999 assert_eq!(post_relations.len(), 1);
1000
1001 let none = g.relations_of("nonexistent");
1002 assert!(none.is_empty());
1003 }
1004
1005 #[test]
1006 fn test_all_relations() {
1007 let mut g = EntityGraph::new();
1008 g.add_edge("user", "posts")
1009 .add_edge("user", "profile")
1010 .add_edge("post", "comments");
1011
1012 let rels = g.all_relations();
1013 assert_eq!(rels, vec!["comments", "posts", "profile"]);
1014 }
1015
1016 #[test]
1017 fn test_all_parent_fields() {
1018 let mut g = EntityGraph::new();
1019 g.add_edge("user", "posts")
1020 .add_edge("user", "profile")
1021 .add_edge("post", "comments");
1022
1023 let fields = g.all_parent_fields();
1024 assert_eq!(fields, vec!["post", "user"]);
1025 }
1026
1027 #[test]
1028 fn test_all_relations_recursive() {
1029 let mut sub = EntityGraph::new();
1030 sub.add_edge("comments", "author")
1031 .add_edge("comments", "likes");
1032
1033 let mut g = EntityGraph::new();
1034 g.add_edge("user", "posts")
1035 .add_edge_with_graph("user", "comments", sub);
1036
1037 let all = g.all_relations_recursive();
1038 assert!(all.contains(&"posts".to_string()));
1039 assert!(all.contains(&"comments".to_string()));
1040 assert!(all.contains(&"author".to_string()));
1041 assert!(all.contains(&"likes".to_string()));
1042 assert_eq!(all.len(), 4);
1043 }
1044
1045 #[test]
1046 fn test_default_graph_is_empty() {
1047 let g = EntityGraph::default();
1048 assert!(g.is_empty());
1049 }
1050
1051 #[test]
1054 fn test_strategy_name() {
1055 assert_eq!(BatchStrategy::In.name(), "in");
1056 assert_eq!(BatchStrategy::Join.name(), "join");
1057 assert_eq!(BatchStrategy::Subquery.name(), "subquery");
1058 }
1059
1060 #[test]
1061 fn test_strategy_default_is_in() {
1062 assert_eq!(BatchStrategy::default(), BatchStrategy::In);
1063 }
1064
1065 #[test]
1066 fn test_render_in_clause_empty() {
1067 let sql = BatchStrategy::render_in_clause("id", 0);
1068 assert_eq!(sql, "id IN ()");
1069 }
1070
1071 #[test]
1072 fn test_render_in_clause_single() {
1073 let sql = BatchStrategy::render_in_clause("id", 1);
1074 assert_eq!(sql, "id IN (?)");
1075 }
1076
1077 #[test]
1078 fn test_render_in_clause_multiple() {
1079 let sql = BatchStrategy::render_in_clause("user_id", 3);
1080 assert_eq!(sql, "user_id IN (?, ?, ?)");
1081 }
1082
1083 #[test]
1086 fn test_default_config() {
1087 let config = BatchSizeConfig::default();
1088 assert_eq!(config.size, 100);
1089 assert_eq!(config.strategy, BatchStrategy::In);
1090 }
1091
1092 #[test]
1093 fn test_with_size() {
1094 let config = BatchSizeConfig::with_size(50);
1095 assert_eq!(config.size, 50);
1096 assert_eq!(config.strategy, BatchStrategy::In);
1097 }
1098
1099 #[test]
1100 fn test_new_with_strategy() {
1101 let config = BatchSizeConfig::new(200, BatchStrategy::Join);
1102 assert_eq!(config.size, 200);
1103 assert_eq!(config.strategy, BatchStrategy::Join);
1104 }
1105
1106 #[test]
1107 fn test_batch_count_zero() {
1108 let config = BatchSizeConfig::with_size(100);
1109 assert_eq!(config.batch_count(0), 0);
1110 }
1111
1112 #[test]
1113 fn test_batch_count_exact_multiple() {
1114 let config = BatchSizeConfig::with_size(100);
1115 assert_eq!(config.batch_count(100), 1);
1116 assert_eq!(config.batch_count(200), 2);
1117 assert_eq!(config.batch_count(500), 5);
1118 }
1119
1120 #[test]
1121 fn test_batch_count_with_remainder() {
1122 let config = BatchSizeConfig::with_size(100);
1123 assert_eq!(config.batch_count(1), 1);
1124 assert_eq!(config.batch_count(99), 1);
1125 assert_eq!(config.batch_count(101), 2);
1126 assert_eq!(config.batch_count(150), 2);
1127 assert_eq!(config.batch_count(201), 3);
1128 }
1129
1130 #[test]
1131 fn test_batch_range() {
1132 let config = BatchSizeConfig::with_size(100);
1133
1134 assert_eq!(config.batch_range(0, 250), 0..100);
1135 assert_eq!(config.batch_range(1, 250), 100..200);
1136 assert_eq!(config.batch_range(2, 250), 200..250);
1137 }
1138
1139 #[test]
1140 fn test_batch_range_exact() {
1141 let config = BatchSizeConfig::with_size(100);
1142
1143 assert_eq!(config.batch_range(0, 100), 0..100);
1144 assert_eq!(config.batch_range(1, 100), 100..100); }
1146
1147 #[test]
1148 fn test_batch_range_small_batch() {
1149 let config = BatchSizeConfig::with_size(10);
1150
1151 assert_eq!(config.batch_range(0, 25), 0..10);
1152 assert_eq!(config.batch_range(1, 25), 10..20);
1153 assert_eq!(config.batch_range(2, 25), 20..25);
1154 }
1155
1156 fn make_loader() -> BatchLoader<i64, String> {
1159 let loader = Box::new(|ids: &[i64]| -> HashMap<i64, String> {
1160 ids.iter().map(|id| (*id, format!("user_{}", id))).collect()
1161 });
1162 BatchLoader::new(2, loader)
1163 }
1164
1165 #[test]
1166 fn test_batch_loader_load_many_single_batch() {
1167 let loader = make_loader();
1168 let result = loader.load_many(&[1, 2]);
1169 assert_eq!(result.len(), 2);
1170 assert_eq!(result.get(&1), Some(&"user_1".to_string()));
1171 assert_eq!(result.get(&2), Some(&"user_2".to_string()));
1172 }
1173
1174 #[test]
1175 fn test_batch_loader_load_many_multiple_batches() {
1176 let loader = make_loader();
1177 let result = loader.load_many(&[1, 2, 3, 4, 5]);
1179 assert_eq!(result.len(), 5);
1180 for id in 1..=5 {
1181 assert_eq!(
1182 result.get(&id),
1183 Some(&format!("user_{}", id)),
1184 "missing user {}",
1185 id
1186 );
1187 }
1188 }
1189
1190 #[test]
1191 fn test_batch_loader_load_one() {
1192 let loader = make_loader();
1193 let result = loader.load_one(&42);
1194 assert_eq!(result, Some("user_42".to_string()));
1195 }
1196
1197 #[test]
1198 fn test_batch_loader_load_one_missing() {
1199 let loader: BatchLoader<i64, String> =
1201 BatchLoader::new(10, Box::new(|_ids: &[i64]| HashMap::new()));
1202 let result = loader.load_one(&100);
1203 assert_eq!(result, None);
1204 }
1205
1206 #[test]
1207 fn test_batch_loader_caches_results() {
1208 let call_count = std::sync::Arc::new(std::sync::Mutex::new(0u32));
1209 let call_count_clone = call_count.clone();
1210
1211 let loader = Box::new(move |ids: &[i64]| -> HashMap<i64, String> {
1212 *call_count_clone.lock().unwrap() += 1;
1213 ids.iter().map(|id| (*id, format!("user_{}", id))).collect()
1214 });
1215
1216 let batch_loader = BatchLoader::new(100, loader);
1217
1218 batch_loader.load_many(&[1, 2, 3]);
1220 assert_eq!(*call_count.lock().unwrap(), 1);
1221
1222 batch_loader.load_many(&[1, 2, 3]);
1224 assert_eq!(*call_count.lock().unwrap(), 1); batch_loader.load_many(&[4, 5]);
1228 assert_eq!(*call_count.lock().unwrap(), 2);
1229 }
1230
1231 #[test]
1232 fn test_batch_loader_partial_cache_hit() {
1233 let call_count = std::sync::Arc::new(std::sync::Mutex::new(0u32));
1234 let call_count_clone = call_count.clone();
1235
1236 let loader = Box::new(move |ids: &[i64]| -> HashMap<i64, String> {
1237 *call_count_clone.lock().unwrap() += 1;
1238 ids.iter().map(|id| (*id, format!("user_{}", id))).collect()
1239 });
1240
1241 let batch_loader = BatchLoader::new(100, loader);
1242
1243 batch_loader.load_many(&[1, 2, 3]);
1245 assert_eq!(*call_count.lock().unwrap(), 1);
1246
1247 let result = batch_loader.load_many(&[1, 2, 3, 4, 5]);
1249 assert_eq!(result.len(), 5);
1250 assert_eq!(*call_count.lock().unwrap(), 2); assert_eq!(batch_loader.cache_size(), 5);
1254 }
1255
1256 #[test]
1257 fn test_batch_loader_clear_cache() {
1258 let loader = make_loader();
1259 loader.load_many(&[1, 2]);
1260 assert_eq!(loader.cache_size(), 2);
1261
1262 loader.clear_cache();
1263 assert_eq!(loader.cache_size(), 0);
1264 }
1265
1266 #[test]
1267 fn test_batch_loader_empty_input() {
1268 let loader = make_loader();
1269 let result = loader.load_many(&[]);
1270 assert!(result.is_empty());
1271 }
1272
1273 #[test]
1274 fn test_batch_loader_batch_size_attribute() {
1275 let loader = make_loader();
1276 assert_eq!(loader.batch_size(), 2);
1277 }
1278
1279 #[test]
1280 fn test_batch_loader_with_size_1() {
1281 let loader = BatchLoader::new(
1282 1,
1283 Box::new(|ids: &[i64]| ids.iter().map(|id| (*id, *id * 10)).collect()),
1284 );
1285 let result = loader.load_many(&[1, 2, 3]);
1286 assert_eq!(result.len(), 3);
1287 assert_eq!(result.get(&1), Some(&10));
1288 assert_eq!(result.get(&2), Some(&20));
1289 assert_eq!(result.get(&3), Some(&30));
1290 }
1291
1292 #[test]
1295 fn test_n1_config_default() {
1296 let cfg = N1DetectionConfig::default();
1297 assert_eq!(cfg.threshold, 5);
1298 assert!(cfg.enabled);
1299 }
1300
1301 #[test]
1302 fn test_n1_config_builder() {
1303 let cfg = N1DetectionConfig::new()
1304 .with_threshold(10)
1305 .with_enabled(false);
1306 assert_eq!(cfg.threshold, 10);
1307 assert!(!cfg.enabled);
1308
1309 let cfg2 = N1DetectionConfig::new().with_threshold(0);
1311 assert_eq!(cfg2.threshold, 1);
1312 }
1313
1314 #[test]
1315 fn test_n1_detector_default() {
1316 let det = N1QueryDetector::default();
1317 assert!(det.is_enabled());
1318 assert_eq!(det.threshold(), 5);
1319 assert!(!det.is_window_active());
1320 assert!(!det.has_n_plus_one());
1321 }
1322
1323 #[test]
1324 fn test_n1_detector_disabled_is_noop() {
1325 let det = N1QueryDetector::new(N1DetectionConfig::new().with_enabled(false));
1326 det.start_window();
1327 for _ in 0..100 {
1328 det.record_single_load("posts");
1329 }
1330 assert_eq!(det.current_count("posts"), 0);
1332 let alerts = det.end_window();
1333 assert!(alerts.is_empty());
1334 }
1335
1336 #[test]
1337 fn test_n1_detector_records_outside_window_ignored() {
1338 let det = N1QueryDetector::with_defaults();
1339 det.record_single_load("posts");
1341 assert_eq!(det.current_count("posts"), 0);
1342 }
1343
1344 #[test]
1345 fn test_n1_detector_below_threshold_no_alert() {
1346 let det = N1QueryDetector::with_defaults(); det.start_window();
1348 for _ in 0..4 {
1349 det.record_single_load("posts");
1350 }
1351 assert_eq!(det.current_count("posts"), 4);
1352 let alerts = det.end_window();
1353 assert!(alerts.is_empty(), "below threshold should not alert");
1354 assert!(!det.has_n_plus_one());
1355 }
1356
1357 #[test]
1358 fn test_n1_detector_at_threshold_triggers_alert() {
1359 let det = N1QueryDetector::with_defaults(); det.start_window();
1361 for _ in 0..5 {
1362 det.record_single_load("posts");
1363 }
1364 let alerts = det.end_window();
1365 assert_eq!(alerts.len(), 1);
1366 assert_eq!(alerts[0].relation, "posts");
1367 assert_eq!(alerts[0].query_count, 5);
1368 assert_eq!(alerts[0].threshold, 5);
1369 assert_eq!(alerts[0].batch_count, 0);
1370 assert!(alerts[0].no_batch_used());
1371 assert!(det.has_n_plus_one());
1372 }
1373
1374 #[test]
1375 fn test_n1_detector_above_threshold_triggers_alert() {
1376 let det = N1QueryDetector::with_defaults();
1377 det.start_window();
1378 for _ in 0..10 {
1379 det.record_single_load("posts");
1380 }
1381 let alerts = det.end_window();
1382 assert_eq!(alerts.len(), 1);
1383 assert_eq!(alerts[0].query_count, 10);
1384 assert!(alerts[0].suggested_batch_size() >= 50);
1386 }
1387
1388 #[test]
1389 fn test_n1_detector_multiple_relations() {
1390 let det = N1QueryDetector::with_defaults();
1391 det.start_window();
1392 for _ in 0..6 {
1393 det.record_single_load("posts");
1394 }
1395 for _ in 0..3 {
1396 det.record_single_load("comments"); }
1398 for _ in 0..8 {
1399 det.record_single_load("tags");
1400 }
1401 let alerts = det.end_window();
1402 assert_eq!(alerts.len(), 2);
1404 assert_eq!(alerts[0].relation, "posts");
1406 assert_eq!(alerts[0].query_count, 6);
1407 assert_eq!(alerts[1].relation, "tags");
1408 assert_eq!(alerts[1].query_count, 8);
1409 }
1410
1411 #[test]
1412 fn test_n1_detector_batch_load_recorded_separately() {
1413 let det = N1QueryDetector::with_defaults();
1414 det.start_window();
1415 for _ in 0..6 {
1417 det.record_single_load("posts");
1418 }
1419 det.record_batch_load("posts", 100);
1421 det.record_batch_load("posts", 50);
1422 let alerts = det.end_window();
1423 assert_eq!(alerts.len(), 1);
1424 assert_eq!(alerts[0].query_count, 6);
1425 assert_eq!(alerts[0].batch_count, 2);
1426 assert!(!alerts[0].no_batch_used());
1428 }
1429
1430 #[test]
1431 fn test_n1_detector_batch_only_does_not_trigger() {
1432 let det = N1QueryDetector::with_defaults();
1434 det.start_window();
1435 for _ in 0..100 {
1436 det.record_batch_load("posts", 50);
1437 }
1438 assert_eq!(det.current_batch_count("posts"), 100);
1439 assert_eq!(det.current_count("posts"), 0);
1440 let alerts = det.end_window();
1441 assert!(alerts.is_empty());
1442 }
1443
1444 #[test]
1445 fn test_n1_detector_start_window_resets() {
1446 let det = N1QueryDetector::with_defaults();
1447 det.start_window();
1448 for _ in 0..10 {
1449 det.record_single_load("posts");
1450 }
1451 let _ = det.end_window();
1452 assert_eq!(det.alerts().len(), 1);
1453
1454 det.start_window();
1456 assert_eq!(det.alerts().len(), 0);
1457 assert_eq!(det.current_count("posts"), 0);
1458 assert!(det.is_window_active());
1459 }
1460
1461 #[test]
1462 fn test_n1_detector_end_window_deactivates() {
1463 let det = N1QueryDetector::with_defaults();
1464 det.start_window();
1465 assert!(det.is_window_active());
1466 det.end_window();
1467 assert!(!det.is_window_active());
1468
1469 det.record_single_load("posts");
1471 assert_eq!(det.current_count("posts"), 0);
1472 }
1473
1474 #[test]
1475 fn test_n1_detector_custom_threshold() {
1476 let det = N1QueryDetector::new(N1DetectionConfig::new().with_threshold(100));
1477 det.start_window();
1478 for _ in 0..50 {
1479 det.record_single_load("posts");
1480 }
1481 let alerts = det.end_window();
1482 assert!(alerts.is_empty(), "below custom threshold should not alert");
1483
1484 det.start_window();
1485 for _ in 0..100 {
1486 det.record_single_load("posts");
1487 }
1488 let alerts = det.end_window();
1489 assert_eq!(alerts.len(), 1);
1490 assert_eq!(alerts[0].threshold, 100);
1491 assert_eq!(alerts[0].query_count, 100);
1492 }
1493
1494 #[test]
1495 fn test_n1_alert_suggested_batch_size() {
1496 let mk = |cnt: u64| N1Alert {
1497 relation: "x".into(),
1498 query_count: cnt,
1499 batch_count: 0,
1500 threshold: 5,
1501 };
1502 assert_eq!(mk(5).suggested_batch_size(), 50);
1503 assert_eq!(mk(50).suggested_batch_size(), 50);
1504 assert_eq!(mk(51).suggested_batch_size(), 100);
1505 assert_eq!(mk(100).suggested_batch_size(), 100);
1506 assert_eq!(mk(101).suggested_batch_size(), 500);
1507 assert_eq!(mk(500).suggested_batch_size(), 500);
1508 assert_eq!(mk(501).suggested_batch_size(), 1000);
1509 assert_eq!(mk(10000).suggested_batch_size(), 1000);
1510 }
1511
1512 #[test]
1513 fn test_n1_detector_real_n_plus_one_scenario() {
1514 let det = N1QueryDetector::with_defaults();
1516 det.start_window();
1517 let user_ids: Vec<i64> = (1..=20).collect();
1518 for _uid in &user_ids {
1519 det.record_single_load("posts");
1521 }
1522 let alerts = det.end_window();
1523 assert_eq!(alerts.len(), 1);
1524 assert_eq!(alerts[0].query_count, 20);
1525 assert!(alerts[0].no_batch_used());
1526
1527 det.start_window();
1529 det.record_batch_load("posts", 20); let alerts2 = det.end_window();
1531 assert!(alerts2.is_empty(), "batch loading should not trigger N+1");
1532 }
1533
1534 #[test]
1537 fn test_workflow_graph_and_batch_loader() {
1538 let mut graph = EntityGraph::new();
1540 graph.add_edge_with_graph("user", "posts", {
1541 let mut sub = EntityGraph::new();
1542 sub.add_edge("posts", "comments");
1543 sub
1544 });
1545 assert_eq!(graph.all_relations_recursive().len(), 2);
1546
1547 let user_loader = BatchLoader::new(
1549 50,
1550 Box::new(|ids: &[i64]| ids.iter().map(|id| (*id, format!("User#{}", id))).collect()),
1551 );
1552
1553 let user_ids: Vec<i64> = (1..=123).collect();
1555 let users = user_loader.load_many(&user_ids);
1556 assert_eq!(users.len(), 123);
1557 assert_eq!(user_loader.cache_size(), 123);
1558 }
1559
1560 #[test]
1561 fn test_n_plus_1_problem_solved() {
1562 let query_count = std::sync::Arc::new(std::sync::Mutex::new(0u32));
1567 let query_count_clone = query_count.clone();
1568
1569 let post_loader = BatchLoader::new(
1570 100,
1571 Box::new(move |user_ids: &[i64]| {
1572 *query_count_clone.lock().unwrap() += 1;
1573 user_ids
1575 .iter()
1576 .map(|uid| (*uid, format!("posts_for_user_{}", uid)))
1577 .collect()
1578 }),
1579 );
1580
1581 let user_ids: Vec<i64> = (1..=250).collect();
1583 let _posts = post_loader.load_many(&user_ids);
1584
1585 assert_eq!(*query_count.lock().unwrap(), 3);
1587 }
1588}
1589
1590#[cfg(all(test, feature = "prod-n1-tuning"))]
1591mod n1_prod_tests {
1592 use super::*;
1593
1594 #[test]
1595 fn test_n1_config_with_window() {
1596 let config = N1DetectionConfig::new().with_window(std::time::Duration::from_secs(5));
1597 assert_eq!(config.window, std::time::Duration::from_secs(5));
1598 }
1599
1600 #[test]
1601 fn test_n1_config_with_block() {
1602 let config = N1DetectionConfig::new().with_block(true);
1603 assert!(config.block);
1604 }
1605
1606 #[test]
1607 fn test_n1_config_default_window_block() {
1608 let config = N1DetectionConfig::default();
1609 assert_eq!(config.window, std::time::Duration::from_secs(1));
1610 assert!(!config.block);
1611 }
1612
1613 #[test]
1614 fn test_n1_detector_stats_initial() {
1615 let detector = N1QueryDetector::new(N1DetectionConfig::default());
1616 let stats = detector.stats();
1617 assert_eq!(stats.trigger_count, 0);
1618 assert_eq!(stats.block_count, 0);
1619 }
1620
1621 #[test]
1622 fn test_n1_config_backward_compatible() {
1623 let config = N1DetectionConfig::new()
1624 .with_threshold(10)
1625 .with_enabled(true);
1626 assert_eq!(config.threshold, 10);
1627 assert!(config.enabled);
1628 assert_eq!(config.window, std::time::Duration::from_secs(1));
1629 assert!(!config.block);
1630 }
1631}