mago_analyzer/external/
scan.rs1use std::sync::Arc;
2use std::time::Instant;
3
4use mago_database::GlobSettings;
5use mago_database::file::File;
6use mago_database::file::FileId;
7use mago_database::file::FileType;
8use mago_database::matcher::ExclusionMatcher;
9use mago_extension::PayloadWriter;
10use mago_extension::source::SourceSnapshot;
11use mago_names::ResolvedNames;
12use mago_syntax::cst::Program;
13
14use crate::external::AnalyzerTransport;
15use crate::external::Backend;
16use crate::external::ExternalAnalyzerError;
17use crate::external::error::protocol;
18use crate::external::protocol;
19
20const BOOTSTRAP_GROUP: u64 = 0x434F_4445_5343_414E;
21const MESSAGE_OVERHEAD: usize = 12 + 1 + 1 + 4 + 4;
22
23#[derive(Debug, Clone)]
24struct HookMatcher {
25 index: u16,
26 paths: ExclusionMatcher<String>,
27}
28
29#[derive(Debug, Clone)]
30struct BackendPlan {
31 backend: u16,
32 hooks: Box<[HookMatcher]>,
33}
34
35#[derive(Debug, Clone)]
37pub struct CodebaseScanPlan {
38 backends: Arc<[BackendPlan]>,
39}
40
41impl CodebaseScanPlan {
42 pub(super) fn compile<T>(backends: &[Backend<T>]) -> Result<Option<Self>, ExternalAnalyzerError> {
43 let mut plans = Vec::new();
44 let mut hook_count = 0usize;
45 let mut target_count = 0usize;
46 for (backend, registered) in backends.iter().enumerate() {
47 let mut hooks = Vec::with_capacity(registered.registration.codebase_scan_hooks.len());
48 for hook in ®istered.registration.codebase_scan_hooks {
49 hook_count += 1;
50 target_count += hook.targets.len();
51 let paths = ExclusionMatcher::compile(hook.targets.iter().cloned(), GlobSettings::default()).map_err(
52 |error| {
53 protocol(format!(
54 "codebase-scan hook {} contains an invalid source-file target: {error}",
55 hook.index
56 ))
57 },
58 )?;
59 hooks.push(HookMatcher { index: hook.index, paths });
60 }
61
62 if !hooks.is_empty() {
63 let backend = u16::try_from(backend)
64 .map_err(|_| protocol("more than 65,536 external analyzer backends were configured"))?;
65 plans.push(BackendPlan { backend, hooks: hooks.into_boxed_slice() });
66 }
67 }
68
69 tracing::trace!(
70 backends = plans.len(),
71 hooks = hook_count,
72 targets = target_count,
73 "Compiled external codebase-scan selectors."
74 );
75 Ok((!plans.is_empty()).then(|| Self { backends: plans.into() }))
76 }
77
78 pub fn capture(
87 &self,
88 file: &Arc<File>,
89 program: &Program<'_>,
90 resolved_names: &ResolvedNames<'_>,
91 ) -> Result<Option<CodebaseScanFile>, ExternalAnalyzerError> {
92 if file.file_type != FileType::Host {
93 return Ok(None);
94 }
95
96 let Ok(path) = std::str::from_utf8(&file.name) else {
97 return Ok(None);
98 };
99 let mut routes = Vec::new();
100 for backend in self.backends.iter() {
101 let hooks = backend
102 .hooks
103 .iter()
104 .filter_map(|hook| hook.paths.is_match(path).then_some(hook.index))
105 .collect::<Vec<_>>();
106 if !hooks.is_empty() {
107 routes.push(CodebaseScanRoute { backend: backend.backend, hooks: hooks.into_boxed_slice() });
108 }
109 }
110
111 if routes.is_empty() {
112 return Ok(None);
113 }
114
115 let snapshot = SourceSnapshot::complete_with_literals(program, resolved_names)?;
116 let mut writer = PayloadWriter::with_capacity(snapshot.encoded_len_with_literals());
117 snapshot.write_to_with_literals(&mut writer)?;
118 Ok(Some(CodebaseScanFile {
119 file: Arc::clone(file),
120 snapshot: writer.finish().into(),
121 routes: routes.into_boxed_slice(),
122 }))
123 }
124}
125
126#[derive(Debug, Clone)]
127struct CodebaseScanRoute {
128 backend: u16,
129 hooks: Box<[u16]>,
130}
131
132#[derive(Debug, Clone)]
134pub struct CodebaseScanFile {
135 file: Arc<File>,
136 snapshot: Arc<[u8]>,
137 routes: Box<[CodebaseScanRoute]>,
138}
139
140impl CodebaseScanFile {
141 #[inline]
142 #[must_use]
143 pub fn file_id(&self) -> FileId {
144 self.file.id
145 }
146
147 fn route(&self, backend: u16) -> Option<&[u16]> {
148 self.routes.iter().find(|route| route.backend == backend).map(|route| route.hooks.as_ref())
149 }
150
151 fn encoded_len(&self, hooks: &[u16]) -> usize {
152 4usize
153 .saturating_add(hooks.len().saturating_mul(2))
154 .saturating_add(4)
155 .saturating_add(self.file.name.len())
156 .saturating_add(4)
157 .saturating_add(self.file.contents.len())
158 .saturating_add(self.snapshot.len())
159 }
160}
161
162pub(super) fn dispatch<T>(
163 backends: &[Backend<T>],
164 mut files: Vec<CodebaseScanFile>,
165) -> Result<(), ExternalAnalyzerError>
166where
167 T: AnalyzerTransport,
168{
169 files.sort_unstable_by(|left, right| left.file.name.cmp(&right.file.name));
170 for (backend_index, backend) in backends.iter().enumerate() {
171 if backend.registration.codebase_scan_hooks.is_empty() {
172 continue;
173 }
174
175 let backend_index = u16::try_from(backend_index)
176 .map_err(|_| protocol("more than 65,536 external analyzer backends were configured"))?;
177 let selected =
178 files.iter().filter_map(|file| file.route(backend_index).map(|hooks| (file, hooks))).collect::<Vec<_>>();
179 let active_hooks = backend.registration.codebase_scan_hooks.iter().map(|hook| hook.index).collect::<Vec<_>>();
180 let maximum = backend.transport.maximum_payload_size();
181 let batches = encode_batches(&active_hooks, &selected, maximum)?;
182 let request_bytes = batches.iter().map(Vec::len).sum::<usize>();
183 let started_at = tracing::enabled!(tracing::Level::TRACE).then(Instant::now);
184 tracing::trace!(
185 backend = backend_index,
186 files = selected.len(),
187 batches = batches.len(),
188 request_bytes,
189 "Broadcasting filtered codebase-scan snapshots."
190 );
191 let responses = backend.transport.broadcast_sequence(BOOTSTRAP_GROUP, &batches)?;
192 for batch in responses {
193 for response in batch {
194 protocol::decode_codebase_scan_response(&response)?;
195 }
196 }
197 if let Some(started_at) = started_at {
198 tracing::trace!(
199 backend = backend_index,
200 files = selected.len(),
201 batches = batches.len(),
202 request_bytes,
203 elapsed = ?started_at.elapsed(),
204 "Filtered codebase-scan broadcast completed."
205 );
206 }
207 }
208
209 Ok(())
210}
211
212fn encode_batches(
213 active_hooks: &[u16],
214 files: &[(&CodebaseScanFile, &[u16])],
215 maximum: usize,
216) -> Result<Vec<Vec<u8>>, ExternalAnalyzerError> {
217 let message_overhead = MESSAGE_OVERHEAD.saturating_add(active_hooks.len().saturating_mul(2));
218 let mut ranges = Vec::new();
219 let mut start = 0;
220 let mut length = message_overhead;
221 for (index, (file, hooks)) in files.iter().enumerate() {
222 let record = file.encoded_len(hooks);
223 if message_overhead.saturating_add(record) > maximum {
224 return Err(protocol(format!(
225 "codebase-scan snapshot for `{}` requires {} bytes, exceeding the worker payload limit of {maximum}",
226 mago_bytes::BytesDisplay(&file.file.name),
227 message_overhead.saturating_add(record),
228 )));
229 }
230 if index > start && length.saturating_add(record) > maximum {
231 ranges.push(start..index);
232 start = index;
233 length = message_overhead;
234 }
235 length = length.saturating_add(record);
236 }
237 ranges.push(start..files.len());
238
239 let range_count = ranges.len();
240 ranges
241 .into_iter()
242 .enumerate()
243 .map(|(index, range)| {
244 let mut writer = protocol::message_writer_with_capacity(
245 protocol::CODEBASE_SCAN_REQUEST,
246 files[range.clone()]
247 .iter()
248 .fold(message_overhead, |size, (file, hooks)| size.saturating_add(file.encoded_len(hooks))),
249 );
250 writer.write_bool(index == 0);
251 writer.write_bool(index + 1 == range_count);
252 writer.write_length(active_hooks.len())?;
253 for hook in active_hooks {
254 writer.write_u16(*hook);
255 }
256 writer.write_length(range.len())?;
257 for (file, hooks) in &files[range] {
258 writer.write_length(hooks.len())?;
259 for hook in *hooks {
260 writer.write_u16(*hook);
261 }
262 writer.write_bytes(&file.file.name)?;
263 writer.write_bytes(&file.file.contents)?;
264 writer.write_raw(&file.snapshot);
265 }
266 Ok(writer.finish())
267 })
268 .collect()
269}