1use std::collections::HashMap;
46use std::hash::Hash;
47use std::sync::{Arc, 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 detector: Option<Arc<N1QueryDetector>>,
506 relation_name: String,
508}
509
510impl<K, V> BatchLoader<K, V>
511where
512 K: Hash + Eq + Clone + Send + Sync,
513 V: Clone + Send + Sync,
514{
515 pub fn new(batch_size: usize, loader: BatchLoaderFn<K, V>) -> Self {
521 Self {
522 batch_size,
523 loader,
524 cache: RwLock::new(HashMap::new()),
525 detector: None,
526 relation_name: String::new(),
527 }
528 }
529
530 pub fn with_detector(
536 mut self,
537 detector: Arc<N1QueryDetector>,
538 relation: impl Into<String>,
539 ) -> Self {
540 self.detector = Some(detector);
541 self.relation_name = relation.into();
542 self
543 }
544
545 pub fn load_many(&self, keys: &[K]) -> HashMap<K, V> {
551 let mut result: HashMap<K, V> = HashMap::new();
552
553 let mut to_load: Vec<K> = Vec::new();
555 if let Ok(cached) = self.cache.read() {
556 for k in keys {
557 if let Some(v) = cached.get(k) {
558 result.insert(k.clone(), v.clone());
559 } else {
560 to_load.push(k.clone());
561 }
562 }
563 } else {
564 to_load.extend(keys.iter().cloned());
565 }
566
567 if to_load.is_empty() {
568 return result;
569 }
570
571 let batch_size = self.batch_size.max(1);
573 let mut all_loaded: HashMap<K, V> = HashMap::new();
574 for chunk in to_load.chunks(batch_size) {
575 let loaded = (self.loader)(chunk);
576 all_loaded.extend(loaded);
577 }
578
579 if let Ok(mut cache) = self.cache.write() {
581 for (k, v) in &all_loaded {
582 cache.insert(k.clone(), v.clone());
583 }
584 }
585
586 if let Some(detector) = &self.detector {
588 detector.record_batch_load(&self.relation_name, to_load.len());
589 }
590
591 result.extend(all_loaded);
593 result
594 }
595
596 pub fn load_one(&self, key: &K) -> Option<V> {
598 if let Some(detector) = &self.detector {
600 detector.record_single_load(&self.relation_name);
601 }
602 let result = self.load_many(std::slice::from_ref(key));
603 result.get(key).cloned()
604 }
605
606 pub fn clear_cache(&self) {
608 if let Ok(mut cache) = self.cache.write() {
609 cache.clear();
610 }
611 }
612
613 pub fn cache_size(&self) -> usize {
615 match self.cache.read() {
616 Ok(g) => g.len(),
617 Err(_) => 0,
618 }
619 }
620
621 pub fn batch_size(&self) -> usize {
623 self.batch_size
624 }
625}
626
627pub struct N1QueryDetector {
672 config: N1DetectionConfig,
674 counts: RwLock<HashMap<String, u64>>,
676 batch_counts: RwLock<HashMap<String, u64>>,
678 window_active: RwLock<bool>,
680 alerts: RwLock<Vec<N1Alert>>,
682 #[cfg(feature = "prod-n1-tuning")]
684 trigger_count: std::sync::atomic::AtomicU64,
685 #[cfg(feature = "prod-n1-tuning")]
687 block_count: std::sync::atomic::AtomicU64,
688}
689
690#[derive(Debug, Clone)]
692pub struct N1DetectionConfig {
693 pub threshold: u64,
695 pub enabled: bool,
697 pub window: std::time::Duration,
699 pub block: bool,
701}
702
703impl Default for N1DetectionConfig {
704 fn default() -> Self {
705 Self {
706 threshold: 5,
707 enabled: true,
708 window: std::time::Duration::from_secs(1),
709 block: false,
710 }
711 }
712}
713
714impl N1DetectionConfig {
715 pub fn new() -> Self {
717 Self::default()
718 }
719
720 pub fn with_threshold(mut self, threshold: u64) -> Self {
722 self.threshold = threshold.max(1);
723 self
724 }
725
726 pub fn with_enabled(mut self, enabled: bool) -> Self {
728 self.enabled = enabled;
729 self
730 }
731
732 #[cfg(feature = "prod-n1-tuning")]
734 pub fn with_window(mut self, window: std::time::Duration) -> Self {
735 self.window = window;
736 self
737 }
738
739 #[cfg(feature = "prod-n1-tuning")]
741 pub fn with_block(mut self, block: bool) -> Self {
742 self.block = block;
743 self
744 }
745}
746
747#[derive(Debug, Clone, PartialEq, Eq)]
749pub struct N1Alert {
750 pub relation: String,
752 pub query_count: u64,
754 pub batch_count: u64,
756 pub threshold: u64,
758}
759
760impl N1Alert {
761 pub fn no_batch_used(&self) -> bool {
763 self.batch_count == 0
764 }
765
766 pub fn suggested_batch_size(&self) -> usize {
768 let n = self.query_count as usize;
769 if n <= 50 {
770 50
771 } else if n <= 100 {
772 100
773 } else if n <= 500 {
774 500
775 } else {
776 1000
777 }
778 }
779}
780
781impl N1QueryDetector {
782 pub fn new(config: N1DetectionConfig) -> Self {
784 Self {
785 config,
786 counts: RwLock::new(HashMap::new()),
787 batch_counts: RwLock::new(HashMap::new()),
788 window_active: RwLock::new(false),
789 alerts: RwLock::new(Vec::new()),
790 #[cfg(feature = "prod-n1-tuning")]
791 trigger_count: std::sync::atomic::AtomicU64::new(0),
792 #[cfg(feature = "prod-n1-tuning")]
793 block_count: std::sync::atomic::AtomicU64::new(0),
794 }
795 }
796
797 #[cfg(feature = "prod-n1-tuning")]
799 pub fn stats(&self) -> N1DetectorStats {
800 N1DetectorStats {
801 trigger_count: self
802 .trigger_count
803 .load(std::sync::atomic::Ordering::Relaxed),
804 block_count: self.block_count.load(std::sync::atomic::Ordering::Relaxed),
805 }
806 }
807
808 pub fn with_defaults() -> Self {
810 Self::new(N1DetectionConfig::default())
811 }
812
813 pub fn is_enabled(&self) -> bool {
815 self.config.enabled
816 }
817
818 pub fn threshold(&self) -> u64 {
820 self.config.threshold
821 }
822
823 pub fn start_window(&self) {
827 if !self.config.enabled {
828 return;
829 }
830 if let Ok(mut counts) = self.counts.write() {
831 *counts = HashMap::new();
832 }
833 if let Ok(mut batch_counts) = self.batch_counts.write() {
834 *batch_counts = HashMap::new();
835 }
836 if let Ok(mut alerts) = self.alerts.write() {
837 *alerts = Vec::new();
838 }
839 if let Ok(mut window_active) = self.window_active.write() {
840 *window_active = true;
841 }
842 }
843
844 pub fn end_window(&self) -> Vec<N1Alert> {
849 if !self.config.enabled {
850 return Vec::new();
851 }
852 if let Ok(mut window_active) = self.window_active.write() {
853 *window_active = false;
854 }
855
856 let new_alerts: Vec<N1Alert> = match (self.counts.read(), self.batch_counts.read()) {
858 (Ok(counts), Ok(batch_counts)) => {
859 let mut alerts: Vec<N1Alert> = counts
860 .iter()
861 .filter_map(|(rel, &cnt)| {
862 if cnt >= self.config.threshold {
863 Some(N1Alert {
864 relation: rel.clone(),
865 query_count: cnt,
866 batch_count: *batch_counts.get(rel).unwrap_or(&0),
867 threshold: self.config.threshold,
868 })
869 } else {
870 None
871 }
872 })
873 .collect();
874 alerts.sort_by(|a, b| a.relation.cmp(&b.relation));
876 alerts
877 }
878 _ => Vec::new(),
879 };
880
881 if let Ok(mut alerts) = self.alerts.write() {
882 *alerts = new_alerts.clone();
883 }
884 new_alerts
885 }
886
887 pub fn record_single_load(&self, relation: &str) {
889 if !self.config.enabled {
890 return;
891 }
892 {
893 let active = self.window_active.read().map(|g| *g).unwrap_or(false);
894 if !active {
895 return;
896 }
897 }
898 if let Ok(mut counts) = self.counts.write() {
899 *counts.entry(relation.to_string()).or_insert(0) += 1;
900 }
901 }
902
903 pub fn record_batch_load(&self, relation: &str, _keys_count: usize) {
908 if !self.config.enabled {
909 return;
910 }
911 {
912 let active = self.window_active.read().map(|g| *g).unwrap_or(false);
913 if !active {
914 return;
915 }
916 }
917 if let Ok(mut batch_counts) = self.batch_counts.write() {
918 *batch_counts.entry(relation.to_string()).or_insert(0) += 1;
919 }
920 }
921
922 pub fn alerts(&self) -> Vec<N1Alert> {
924 self.alerts.read().map(|g| g.clone()).unwrap_or_default()
925 }
926
927 pub fn current_count(&self, relation: &str) -> u64 {
929 self.counts
930 .read()
931 .map(|g| g.get(relation).copied().unwrap_or(0))
932 .unwrap_or(0)
933 }
934
935 pub fn current_batch_count(&self, relation: &str) -> u64 {
937 self.batch_counts
938 .read()
939 .map(|g| g.get(relation).copied().unwrap_or(0))
940 .unwrap_or(0)
941 }
942
943 pub fn is_window_active(&self) -> bool {
945 self.window_active.read().map(|g| *g).unwrap_or(false)
946 }
947
948 pub fn has_n_plus_one(&self) -> bool {
950 !self.alerts().is_empty()
951 }
952}
953
954impl Default for N1QueryDetector {
955 fn default() -> Self {
956 Self::with_defaults()
957 }
958}
959
960#[cfg(feature = "prod-n1-tuning")]
962#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)]
963pub struct N1DetectorStats {
964 pub trigger_count: u64,
966 pub block_count: u64,
968}
969
970#[cfg(test)]
974mod tests {
975 use super::*;
976
977 #[test]
980 fn test_new_graph_is_empty() {
981 let g = EntityGraph::new();
982 assert!(g.is_empty());
983 assert_eq!(g.edge_count(), 0);
984 }
985
986 #[test]
987 fn test_add_edge() {
988 let mut g = EntityGraph::new();
989 g.add_edge("user", "posts");
990 assert_eq!(g.edge_count(), 1);
991 assert!(!g.is_empty());
992 }
993
994 #[test]
995 fn test_add_multiple_edges() {
996 let mut g = EntityGraph::new();
997 g.add_edge("user", "posts")
998 .add_edge("user", "profile")
999 .add_edge("user", "comments");
1000 assert_eq!(g.edge_count(), 3);
1001 }
1002
1003 #[test]
1004 fn test_add_edge_with_sub_graph() {
1005 let mut sub = EntityGraph::new();
1006 sub.add_edge("comments", "author");
1007
1008 let mut g = EntityGraph::new();
1009 g.add_edge_with_graph("user", "posts", sub);
1010
1011 assert_eq!(g.edge_count(), 1);
1012 assert!(g.edges()[0].sub_graph.is_some());
1013 assert_eq!(g.edges()[0].sub_graph.as_ref().unwrap().edge_count(), 1);
1014 }
1015
1016 #[test]
1017 fn test_relations_of() {
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 user_relations = g.relations_of("user");
1024 assert_eq!(user_relations.len(), 2);
1025 assert_eq!(user_relations[0].relation, "posts");
1026 assert_eq!(user_relations[1].relation, "profile");
1027
1028 let post_relations = g.relations_of("post");
1029 assert_eq!(post_relations.len(), 1);
1030
1031 let none = g.relations_of("nonexistent");
1032 assert!(none.is_empty());
1033 }
1034
1035 #[test]
1036 fn test_all_relations() {
1037 let mut g = EntityGraph::new();
1038 g.add_edge("user", "posts")
1039 .add_edge("user", "profile")
1040 .add_edge("post", "comments");
1041
1042 let rels = g.all_relations();
1043 assert_eq!(rels, vec!["comments", "posts", "profile"]);
1044 }
1045
1046 #[test]
1047 fn test_all_parent_fields() {
1048 let mut g = EntityGraph::new();
1049 g.add_edge("user", "posts")
1050 .add_edge("user", "profile")
1051 .add_edge("post", "comments");
1052
1053 let fields = g.all_parent_fields();
1054 assert_eq!(fields, vec!["post", "user"]);
1055 }
1056
1057 #[test]
1058 fn test_all_relations_recursive() {
1059 let mut sub = EntityGraph::new();
1060 sub.add_edge("comments", "author")
1061 .add_edge("comments", "likes");
1062
1063 let mut g = EntityGraph::new();
1064 g.add_edge("user", "posts")
1065 .add_edge_with_graph("user", "comments", sub);
1066
1067 let all = g.all_relations_recursive();
1068 assert!(all.contains(&"posts".to_string()));
1069 assert!(all.contains(&"comments".to_string()));
1070 assert!(all.contains(&"author".to_string()));
1071 assert!(all.contains(&"likes".to_string()));
1072 assert_eq!(all.len(), 4);
1073 }
1074
1075 #[test]
1076 fn test_default_graph_is_empty() {
1077 let g = EntityGraph::default();
1078 assert!(g.is_empty());
1079 }
1080
1081 #[test]
1084 fn test_strategy_name() {
1085 assert_eq!(BatchStrategy::In.name(), "in");
1086 assert_eq!(BatchStrategy::Join.name(), "join");
1087 assert_eq!(BatchStrategy::Subquery.name(), "subquery");
1088 }
1089
1090 #[test]
1091 fn test_strategy_default_is_in() {
1092 assert_eq!(BatchStrategy::default(), BatchStrategy::In);
1093 }
1094
1095 #[test]
1096 fn test_render_in_clause_empty() {
1097 let sql = BatchStrategy::render_in_clause("id", 0);
1098 assert_eq!(sql, "id IN ()");
1099 }
1100
1101 #[test]
1102 fn test_render_in_clause_single() {
1103 let sql = BatchStrategy::render_in_clause("id", 1);
1104 assert_eq!(sql, "id IN (?)");
1105 }
1106
1107 #[test]
1108 fn test_render_in_clause_multiple() {
1109 let sql = BatchStrategy::render_in_clause("user_id", 3);
1110 assert_eq!(sql, "user_id IN (?, ?, ?)");
1111 }
1112
1113 #[test]
1116 fn test_default_config() {
1117 let config = BatchSizeConfig::default();
1118 assert_eq!(config.size, 100);
1119 assert_eq!(config.strategy, BatchStrategy::In);
1120 }
1121
1122 #[test]
1123 fn test_with_size() {
1124 let config = BatchSizeConfig::with_size(50);
1125 assert_eq!(config.size, 50);
1126 assert_eq!(config.strategy, BatchStrategy::In);
1127 }
1128
1129 #[test]
1130 fn test_new_with_strategy() {
1131 let config = BatchSizeConfig::new(200, BatchStrategy::Join);
1132 assert_eq!(config.size, 200);
1133 assert_eq!(config.strategy, BatchStrategy::Join);
1134 }
1135
1136 #[test]
1137 fn test_batch_count_zero() {
1138 let config = BatchSizeConfig::with_size(100);
1139 assert_eq!(config.batch_count(0), 0);
1140 }
1141
1142 #[test]
1143 fn test_batch_count_exact_multiple() {
1144 let config = BatchSizeConfig::with_size(100);
1145 assert_eq!(config.batch_count(100), 1);
1146 assert_eq!(config.batch_count(200), 2);
1147 assert_eq!(config.batch_count(500), 5);
1148 }
1149
1150 #[test]
1151 fn test_batch_count_with_remainder() {
1152 let config = BatchSizeConfig::with_size(100);
1153 assert_eq!(config.batch_count(1), 1);
1154 assert_eq!(config.batch_count(99), 1);
1155 assert_eq!(config.batch_count(101), 2);
1156 assert_eq!(config.batch_count(150), 2);
1157 assert_eq!(config.batch_count(201), 3);
1158 }
1159
1160 #[test]
1161 fn test_batch_range() {
1162 let config = BatchSizeConfig::with_size(100);
1163
1164 assert_eq!(config.batch_range(0, 250), 0..100);
1165 assert_eq!(config.batch_range(1, 250), 100..200);
1166 assert_eq!(config.batch_range(2, 250), 200..250);
1167 }
1168
1169 #[test]
1170 fn test_batch_range_exact() {
1171 let config = BatchSizeConfig::with_size(100);
1172
1173 assert_eq!(config.batch_range(0, 100), 0..100);
1174 assert_eq!(config.batch_range(1, 100), 100..100); }
1176
1177 #[test]
1178 fn test_batch_range_small_batch() {
1179 let config = BatchSizeConfig::with_size(10);
1180
1181 assert_eq!(config.batch_range(0, 25), 0..10);
1182 assert_eq!(config.batch_range(1, 25), 10..20);
1183 assert_eq!(config.batch_range(2, 25), 20..25);
1184 }
1185
1186 fn make_loader() -> BatchLoader<i64, String> {
1189 let loader = Box::new(|ids: &[i64]| -> HashMap<i64, String> {
1190 ids.iter().map(|id| (*id, format!("user_{}", id))).collect()
1191 });
1192 BatchLoader::new(2, loader)
1193 }
1194
1195 #[test]
1196 fn test_batch_loader_load_many_single_batch() {
1197 let loader = make_loader();
1198 let result = loader.load_many(&[1, 2]);
1199 assert_eq!(result.len(), 2);
1200 assert_eq!(result.get(&1), Some(&"user_1".to_string()));
1201 assert_eq!(result.get(&2), Some(&"user_2".to_string()));
1202 }
1203
1204 #[test]
1205 fn test_batch_loader_load_many_multiple_batches() {
1206 let loader = make_loader();
1207 let result = loader.load_many(&[1, 2, 3, 4, 5]);
1209 assert_eq!(result.len(), 5);
1210 for id in 1..=5 {
1211 assert_eq!(
1212 result.get(&id),
1213 Some(&format!("user_{}", id)),
1214 "missing user {}",
1215 id
1216 );
1217 }
1218 }
1219
1220 #[test]
1221 fn test_batch_loader_load_one() {
1222 let loader = make_loader();
1223 let result = loader.load_one(&42);
1224 assert_eq!(result, Some("user_42".to_string()));
1225 }
1226
1227 #[test]
1228 fn test_batch_loader_load_one_missing() {
1229 let loader: BatchLoader<i64, String> =
1231 BatchLoader::new(10, Box::new(|_ids: &[i64]| HashMap::new()));
1232 let result = loader.load_one(&100);
1233 assert_eq!(result, None);
1234 }
1235
1236 #[test]
1237 fn test_batch_loader_caches_results() {
1238 let call_count = std::sync::Arc::new(std::sync::Mutex::new(0u32));
1239 let call_count_clone = call_count.clone();
1240
1241 let loader = Box::new(move |ids: &[i64]| -> HashMap<i64, String> {
1242 *call_count_clone.lock().unwrap() += 1;
1243 ids.iter().map(|id| (*id, format!("user_{}", id))).collect()
1244 });
1245
1246 let batch_loader = BatchLoader::new(100, loader);
1247
1248 batch_loader.load_many(&[1, 2, 3]);
1250 assert_eq!(*call_count.lock().unwrap(), 1);
1251
1252 batch_loader.load_many(&[1, 2, 3]);
1254 assert_eq!(*call_count.lock().unwrap(), 1); batch_loader.load_many(&[4, 5]);
1258 assert_eq!(*call_count.lock().unwrap(), 2);
1259 }
1260
1261 #[test]
1262 fn test_batch_loader_partial_cache_hit() {
1263 let call_count = std::sync::Arc::new(std::sync::Mutex::new(0u32));
1264 let call_count_clone = call_count.clone();
1265
1266 let loader = Box::new(move |ids: &[i64]| -> HashMap<i64, String> {
1267 *call_count_clone.lock().unwrap() += 1;
1268 ids.iter().map(|id| (*id, format!("user_{}", id))).collect()
1269 });
1270
1271 let batch_loader = BatchLoader::new(100, loader);
1272
1273 batch_loader.load_many(&[1, 2, 3]);
1275 assert_eq!(*call_count.lock().unwrap(), 1);
1276
1277 let result = batch_loader.load_many(&[1, 2, 3, 4, 5]);
1279 assert_eq!(result.len(), 5);
1280 assert_eq!(*call_count.lock().unwrap(), 2); assert_eq!(batch_loader.cache_size(), 5);
1284 }
1285
1286 #[test]
1287 fn test_batch_loader_clear_cache() {
1288 let loader = make_loader();
1289 loader.load_many(&[1, 2]);
1290 assert_eq!(loader.cache_size(), 2);
1291
1292 loader.clear_cache();
1293 assert_eq!(loader.cache_size(), 0);
1294 }
1295
1296 #[test]
1297 fn test_batch_loader_empty_input() {
1298 let loader = make_loader();
1299 let result = loader.load_many(&[]);
1300 assert!(result.is_empty());
1301 }
1302
1303 #[test]
1304 fn test_batch_loader_batch_size_attribute() {
1305 let loader = make_loader();
1306 assert_eq!(loader.batch_size(), 2);
1307 }
1308
1309 #[test]
1310 fn test_batch_loader_with_size_1() {
1311 let loader = BatchLoader::new(
1312 1,
1313 Box::new(|ids: &[i64]| ids.iter().map(|id| (*id, *id * 10)).collect()),
1314 );
1315 let result = loader.load_many(&[1, 2, 3]);
1316 assert_eq!(result.len(), 3);
1317 assert_eq!(result.get(&1), Some(&10));
1318 assert_eq!(result.get(&2), Some(&20));
1319 assert_eq!(result.get(&3), Some(&30));
1320 }
1321
1322 #[test]
1325 fn test_batch_loader_with_detector_load_many() {
1326 let detector = Arc::new(N1QueryDetector::with_defaults());
1327 detector.start_window();
1328 let loader = BatchLoader::new(
1329 10,
1330 Box::new(|ids: &[i64]| ids.iter().map(|id| (*id, format!("user_{}", id))).collect()),
1331 )
1332 .with_detector(Arc::clone(&detector), "users");
1333 let result = loader.load_many(&[1, 2, 3]);
1334 assert_eq!(result.len(), 3);
1335 let alerts = detector.end_window();
1336 assert!(alerts.is_empty(), "batch load should not trigger N+1 alert");
1337 }
1338
1339 #[test]
1340 fn test_batch_loader_with_detector_load_one_triggers_n1() {
1341 let detector = Arc::new(N1QueryDetector::with_defaults());
1342 detector.start_window();
1343 let loader = BatchLoader::new(
1344 10,
1345 Box::new(|ids: &[i64]| ids.iter().map(|id| (*id, format!("user_{}", id))).collect()),
1346 )
1347 .with_detector(Arc::clone(&detector), "users");
1348 for i in 0..10 {
1349 let _ = loader.load_one(&i);
1350 }
1351 let alerts = detector.end_window();
1352 assert_eq!(alerts.len(), 1);
1353 assert_eq!(alerts[0].relation, "users");
1354 assert_eq!(alerts[0].query_count, 10);
1355 }
1356
1357 #[test]
1358 fn test_batch_loader_without_detector_no_panic() {
1359 let loader = BatchLoader::new(
1360 10,
1361 Box::new(|ids: &[i64]| ids.iter().map(|id| (*id, format!("user_{}", id))).collect()),
1362 );
1363 let result = loader.load_many(&[1, 2, 3]);
1364 assert_eq!(result.len(), 3);
1365 let val = loader.load_one(&1);
1366 assert!(val.is_some());
1367 }
1368
1369 #[test]
1372 fn test_n1_config_default() {
1373 let cfg = N1DetectionConfig::default();
1374 assert_eq!(cfg.threshold, 5);
1375 assert!(cfg.enabled);
1376 }
1377
1378 #[test]
1379 fn test_n1_config_builder() {
1380 let cfg = N1DetectionConfig::new()
1381 .with_threshold(10)
1382 .with_enabled(false);
1383 assert_eq!(cfg.threshold, 10);
1384 assert!(!cfg.enabled);
1385
1386 let cfg2 = N1DetectionConfig::new().with_threshold(0);
1388 assert_eq!(cfg2.threshold, 1);
1389 }
1390
1391 #[test]
1392 fn test_n1_detector_default() {
1393 let det = N1QueryDetector::default();
1394 assert!(det.is_enabled());
1395 assert_eq!(det.threshold(), 5);
1396 assert!(!det.is_window_active());
1397 assert!(!det.has_n_plus_one());
1398 }
1399
1400 #[test]
1401 fn test_n1_detector_disabled_is_noop() {
1402 let det = N1QueryDetector::new(N1DetectionConfig::new().with_enabled(false));
1403 det.start_window();
1404 for _ in 0..100 {
1405 det.record_single_load("posts");
1406 }
1407 assert_eq!(det.current_count("posts"), 0);
1409 let alerts = det.end_window();
1410 assert!(alerts.is_empty());
1411 }
1412
1413 #[test]
1414 fn test_n1_detector_records_outside_window_ignored() {
1415 let det = N1QueryDetector::with_defaults();
1416 det.record_single_load("posts");
1418 assert_eq!(det.current_count("posts"), 0);
1419 }
1420
1421 #[test]
1422 fn test_n1_detector_below_threshold_no_alert() {
1423 let det = N1QueryDetector::with_defaults(); det.start_window();
1425 for _ in 0..4 {
1426 det.record_single_load("posts");
1427 }
1428 assert_eq!(det.current_count("posts"), 4);
1429 let alerts = det.end_window();
1430 assert!(alerts.is_empty(), "below threshold should not alert");
1431 assert!(!det.has_n_plus_one());
1432 }
1433
1434 #[test]
1435 fn test_n1_detector_at_threshold_triggers_alert() {
1436 let det = N1QueryDetector::with_defaults(); det.start_window();
1438 for _ in 0..5 {
1439 det.record_single_load("posts");
1440 }
1441 let alerts = det.end_window();
1442 assert_eq!(alerts.len(), 1);
1443 assert_eq!(alerts[0].relation, "posts");
1444 assert_eq!(alerts[0].query_count, 5);
1445 assert_eq!(alerts[0].threshold, 5);
1446 assert_eq!(alerts[0].batch_count, 0);
1447 assert!(alerts[0].no_batch_used());
1448 assert!(det.has_n_plus_one());
1449 }
1450
1451 #[test]
1452 fn test_n1_detector_above_threshold_triggers_alert() {
1453 let det = N1QueryDetector::with_defaults();
1454 det.start_window();
1455 for _ in 0..10 {
1456 det.record_single_load("posts");
1457 }
1458 let alerts = det.end_window();
1459 assert_eq!(alerts.len(), 1);
1460 assert_eq!(alerts[0].query_count, 10);
1461 assert!(alerts[0].suggested_batch_size() >= 50);
1463 }
1464
1465 #[test]
1466 fn test_n1_detector_multiple_relations() {
1467 let det = N1QueryDetector::with_defaults();
1468 det.start_window();
1469 for _ in 0..6 {
1470 det.record_single_load("posts");
1471 }
1472 for _ in 0..3 {
1473 det.record_single_load("comments"); }
1475 for _ in 0..8 {
1476 det.record_single_load("tags");
1477 }
1478 let alerts = det.end_window();
1479 assert_eq!(alerts.len(), 2);
1481 assert_eq!(alerts[0].relation, "posts");
1483 assert_eq!(alerts[0].query_count, 6);
1484 assert_eq!(alerts[1].relation, "tags");
1485 assert_eq!(alerts[1].query_count, 8);
1486 }
1487
1488 #[test]
1489 fn test_n1_detector_batch_load_recorded_separately() {
1490 let det = N1QueryDetector::with_defaults();
1491 det.start_window();
1492 for _ in 0..6 {
1494 det.record_single_load("posts");
1495 }
1496 det.record_batch_load("posts", 100);
1498 det.record_batch_load("posts", 50);
1499 let alerts = det.end_window();
1500 assert_eq!(alerts.len(), 1);
1501 assert_eq!(alerts[0].query_count, 6);
1502 assert_eq!(alerts[0].batch_count, 2);
1503 assert!(!alerts[0].no_batch_used());
1505 }
1506
1507 #[test]
1508 fn test_n1_detector_batch_only_does_not_trigger() {
1509 let det = N1QueryDetector::with_defaults();
1511 det.start_window();
1512 for _ in 0..100 {
1513 det.record_batch_load("posts", 50);
1514 }
1515 assert_eq!(det.current_batch_count("posts"), 100);
1516 assert_eq!(det.current_count("posts"), 0);
1517 let alerts = det.end_window();
1518 assert!(alerts.is_empty());
1519 }
1520
1521 #[test]
1522 fn test_n1_detector_start_window_resets() {
1523 let det = N1QueryDetector::with_defaults();
1524 det.start_window();
1525 for _ in 0..10 {
1526 det.record_single_load("posts");
1527 }
1528 let _ = det.end_window();
1529 assert_eq!(det.alerts().len(), 1);
1530
1531 det.start_window();
1533 assert_eq!(det.alerts().len(), 0);
1534 assert_eq!(det.current_count("posts"), 0);
1535 assert!(det.is_window_active());
1536 }
1537
1538 #[test]
1539 fn test_n1_detector_end_window_deactivates() {
1540 let det = N1QueryDetector::with_defaults();
1541 det.start_window();
1542 assert!(det.is_window_active());
1543 det.end_window();
1544 assert!(!det.is_window_active());
1545
1546 det.record_single_load("posts");
1548 assert_eq!(det.current_count("posts"), 0);
1549 }
1550
1551 #[test]
1552 fn test_n1_detector_custom_threshold() {
1553 let det = N1QueryDetector::new(N1DetectionConfig::new().with_threshold(100));
1554 det.start_window();
1555 for _ in 0..50 {
1556 det.record_single_load("posts");
1557 }
1558 let alerts = det.end_window();
1559 assert!(alerts.is_empty(), "below custom threshold should not alert");
1560
1561 det.start_window();
1562 for _ in 0..100 {
1563 det.record_single_load("posts");
1564 }
1565 let alerts = det.end_window();
1566 assert_eq!(alerts.len(), 1);
1567 assert_eq!(alerts[0].threshold, 100);
1568 assert_eq!(alerts[0].query_count, 100);
1569 }
1570
1571 #[test]
1572 fn test_n1_alert_suggested_batch_size() {
1573 let mk = |cnt: u64| N1Alert {
1574 relation: "x".into(),
1575 query_count: cnt,
1576 batch_count: 0,
1577 threshold: 5,
1578 };
1579 assert_eq!(mk(5).suggested_batch_size(), 50);
1580 assert_eq!(mk(50).suggested_batch_size(), 50);
1581 assert_eq!(mk(51).suggested_batch_size(), 100);
1582 assert_eq!(mk(100).suggested_batch_size(), 100);
1583 assert_eq!(mk(101).suggested_batch_size(), 500);
1584 assert_eq!(mk(500).suggested_batch_size(), 500);
1585 assert_eq!(mk(501).suggested_batch_size(), 1000);
1586 assert_eq!(mk(10000).suggested_batch_size(), 1000);
1587 }
1588
1589 #[test]
1590 fn test_n1_detector_real_n_plus_one_scenario() {
1591 let det = N1QueryDetector::with_defaults();
1593 det.start_window();
1594 let user_ids: Vec<i64> = (1..=20).collect();
1595 for _uid in &user_ids {
1596 det.record_single_load("posts");
1598 }
1599 let alerts = det.end_window();
1600 assert_eq!(alerts.len(), 1);
1601 assert_eq!(alerts[0].query_count, 20);
1602 assert!(alerts[0].no_batch_used());
1603
1604 det.start_window();
1606 det.record_batch_load("posts", 20); let alerts2 = det.end_window();
1608 assert!(alerts2.is_empty(), "batch loading should not trigger N+1");
1609 }
1610
1611 #[test]
1614 fn test_workflow_graph_and_batch_loader() {
1615 let mut graph = EntityGraph::new();
1617 graph.add_edge_with_graph("user", "posts", {
1618 let mut sub = EntityGraph::new();
1619 sub.add_edge("posts", "comments");
1620 sub
1621 });
1622 assert_eq!(graph.all_relations_recursive().len(), 2);
1623
1624 let user_loader = BatchLoader::new(
1626 50,
1627 Box::new(|ids: &[i64]| ids.iter().map(|id| (*id, format!("User#{}", id))).collect()),
1628 );
1629
1630 let user_ids: Vec<i64> = (1..=123).collect();
1632 let users = user_loader.load_many(&user_ids);
1633 assert_eq!(users.len(), 123);
1634 assert_eq!(user_loader.cache_size(), 123);
1635 }
1636
1637 #[test]
1638 fn test_n_plus_1_problem_solved() {
1639 let query_count = std::sync::Arc::new(std::sync::Mutex::new(0u32));
1644 let query_count_clone = query_count.clone();
1645
1646 let post_loader = BatchLoader::new(
1647 100,
1648 Box::new(move |user_ids: &[i64]| {
1649 *query_count_clone.lock().unwrap() += 1;
1650 user_ids
1652 .iter()
1653 .map(|uid| (*uid, format!("posts_for_user_{}", uid)))
1654 .collect()
1655 }),
1656 );
1657
1658 let user_ids: Vec<i64> = (1..=250).collect();
1660 let _posts = post_loader.load_many(&user_ids);
1661
1662 assert_eq!(*query_count.lock().unwrap(), 3);
1664 }
1665}
1666
1667#[cfg(all(test, feature = "prod-n1-tuning"))]
1668mod n1_prod_tests {
1669 use super::*;
1670
1671 #[test]
1672 fn test_n1_config_with_window() {
1673 let config = N1DetectionConfig::new().with_window(std::time::Duration::from_secs(5));
1674 assert_eq!(config.window, std::time::Duration::from_secs(5));
1675 }
1676
1677 #[test]
1678 fn test_n1_config_with_block() {
1679 let config = N1DetectionConfig::new().with_block(true);
1680 assert!(config.block);
1681 }
1682
1683 #[test]
1684 fn test_n1_config_default_window_block() {
1685 let config = N1DetectionConfig::default();
1686 assert_eq!(config.window, std::time::Duration::from_secs(1));
1687 assert!(!config.block);
1688 }
1689
1690 #[test]
1691 fn test_n1_detector_stats_initial() {
1692 let detector = N1QueryDetector::new(N1DetectionConfig::default());
1693 let stats = detector.stats();
1694 assert_eq!(stats.trigger_count, 0);
1695 assert_eq!(stats.block_count, 0);
1696 }
1697
1698 #[test]
1699 fn test_n1_config_backward_compatible() {
1700 let config = N1DetectionConfig::new()
1701 .with_threshold(10)
1702 .with_enabled(true);
1703 assert_eq!(config.threshold, 10);
1704 assert!(config.enabled);
1705 assert_eq!(config.window, std::time::Duration::from_secs(1));
1706 assert!(!config.block);
1707 }
1708}