relay_knowledge/code/index/
plan.rs

1use std::{collections::BTreeMap, path::PathBuf, thread};
2
3use crate::domain::{
4    CodeIndexBatch, CodeIndexResourceBudget, CodeIndexSession, CodeMonorepoWorkspace,
5    CodeRepositoryRegistration, CodeRepositorySelector, CodeWorkspaceDetectionConfig,
6    code_snapshot_scope_id,
7};
8
9use super::{
10    CodeIndexError,
11    changes::GitTreeEntry,
12    identity, parse_indexed_file,
13    scope::scoped_source_snapshot,
14    snapshot::{SnapshotBuild, SnapshotScopeFilters, detect_workspaces_for_source_snapshot},
15    source::{
16        RepositorySourceKind, ensure_filesystem_blobs_match_content_hashes,
17        ensure_filesystem_paths_match_content_hashes, filesystem_content_hashes_for_paths,
18        filesystem_tree_hash_from_path_hashes, source_snapshot_batch_bytes,
19    },
20};
21
22const GIT_BLOB_FETCH_GROUP: usize = CodeIndexResourceBudget::DEFAULT_MAX_FILES_PER_BATCH;
23const MIN_PARALLEL_PARSE_FILES: usize = 12;
24const MIN_PARALLEL_PARSE_BYTES: usize = 256 * 1024;
25const TARGET_PARSE_FILES_PER_WORKER: usize = 16;
26const TARGET_PARSE_BYTES_PER_WORKER: usize = 512 * 1024;
27
28/// Blocking plan for a checkpointed full repository index.
29#[derive(Debug, Clone)]
30pub struct CodeIndexPlan {
31    registration: CodeRepositoryRegistration,
32    root: PathBuf,
33    commit: String,
34    tree_hash: String,
35    source_scope: String,
36    path_filters: Vec<String>,
37    language_filters: Vec<String>,
38    source_kind: RepositorySourceKind,
39    filesystem_path_hashes: BTreeMap<String, String>,
40    paths: Vec<GitTreeEntry>,
41    workspaces: Vec<CodeMonorepoWorkspace>,
42    cursor: usize,
43    next_batch_index: usize,
44    resource_budget: CodeIndexResourceBudget,
45}
46
47impl CodeIndexPlan {
48    /// Returns the durable session metadata that storage checkpoints.
49    pub fn session(&self) -> CodeIndexSession {
50        CodeIndexSession {
51            repository_id: self.registration.repository_id.clone(),
52            source_scope: self.source_scope.clone(),
53            base_resolved_commit_sha: None,
54            resolved_commit_sha: self.commit.clone(),
55            tree_hash: self.tree_hash.clone(),
56            path_filters: self.path_filters.clone(),
57            language_filters: self.language_filters.clone(),
58            full_replace: true,
59            total_path_count: self.paths.len(),
60            changed_path_count: self.paths.len(),
61            skipped_unchanged_count: 0,
62            deleted_paths: Vec::new(),
63            tombstones: Vec::new(),
64            workspaces: self.workspaces.clone(),
65            resource_budget: self.resource_budget,
66        }
67    }
68
69    /// Parses the next bounded file batch without retaining prior batches.
70    pub fn parse_next_batch(mut self) -> Result<(Self, Option<CodeIndexBatch>), CodeIndexError> {
71        if self.cursor >= self.paths.len() {
72            return Ok((self, None));
73        }
74
75        let mut build = SnapshotBuild::new_with_scope_filters(
76            &self.registration,
77            self.commit.clone(),
78            self.tree_hash.clone(),
79            SnapshotScopeFilters {
80                path_filters: self.path_filters.clone(),
81                language_filters: self.language_filters.clone(),
82            },
83            true,
84            self.paths.len(),
85            0,
86        );
87        let mut parsed_bytes = 0usize;
88        while self.cursor < self.paths.len() {
89            let fetch_end = next_fetch_end(&self, build.files.len(), parsed_bytes);
90            if fetch_end == self.cursor {
91                break;
92            }
93            let fetched_paths = self.paths[self.cursor..fetch_end]
94                .iter()
95                .map(|entry| entry.path.clone())
96                .collect::<Vec<_>>();
97            ensure_filesystem_paths_match_content_hashes(
98                &self.root,
99                &self.commit,
100                &fetched_paths,
101                &self.filesystem_path_hashes,
102            )?;
103            let blobs = source_snapshot_batch_bytes(
104                &self.root,
105                self.source_kind,
106                &self.commit,
107                &fetched_paths,
108            )?;
109            ensure_filesystem_blobs_match_content_hashes(
110                &self.commit,
111                &fetched_paths,
112                &blobs,
113                &self.filesystem_path_hashes,
114            )?;
115            let parsed_files = parse_fetched_files(&self, &fetched_paths, &blobs)?;
116            for (bytes, parsed_file) in blobs.iter().zip(parsed_files) {
117                parsed_bytes = parsed_bytes.saturating_add(bytes.len());
118                build.append_file_records(parsed_file);
119                self.cursor += 1;
120
121                if !build.files.is_empty()
122                    && (build.files.len() >= self.resource_budget.max_files_per_batch
123                        || parsed_bytes >= self.resource_budget.max_bytes_per_batch
124                        || batch_row_count(&build) >= self.resource_budget.max_rows_per_batch)
125                {
126                    break;
127                }
128            }
129            if !build.files.is_empty()
130                && (build.files.len() >= self.resource_budget.max_files_per_batch
131                    || parsed_bytes >= self.resource_budget.max_bytes_per_batch
132                    || batch_row_count(&build) >= self.resource_budget.max_rows_per_batch)
133            {
134                break;
135            }
136        }
137        identity::enrich_symbol_identities(&build.repository_id, &mut build.symbols);
138
139        let batch = CodeIndexBatch {
140            repository_id: build.repository_id,
141            source_scope: build.source_scope,
142            batch_index: self.next_batch_index,
143            parsed_byte_count: parsed_bytes,
144            files: build.files,
145            symbols: build.symbols,
146            references: build.references,
147            imports: build.imports,
148            dependencies: build.dependencies,
149            feature_flags: build.feature_flags,
150            routes: build.routes,
151            chunks: build.chunks,
152            diagnostics: build.diagnostics,
153        };
154        self.next_batch_index += 1;
155
156        Ok((self, Some(batch)))
157    }
158}
159
160fn parse_fetched_files(
161    plan: &CodeIndexPlan,
162    paths: &[String],
163    blobs: &[Vec<u8>],
164) -> Result<Vec<SnapshotBuild>, CodeIndexError> {
165    let worker_count = worker_count(paths.len(), total_blob_bytes(blobs));
166    if paths.len() <= 1 || worker_count <= 1 {
167        return paths
168            .iter()
169            .zip(blobs.iter())
170            .map(|(path, bytes)| parse_one_file(plan, path, bytes))
171            .collect();
172    }
173
174    let mut parsed = thread::scope(|scope| {
175        let handles = (0..worker_count)
176            .map(|worker_index| {
177                scope.spawn(move || {
178                    parse_worker_stride(plan, paths, blobs, worker_index, worker_count)
179                })
180            })
181            .collect::<Vec<_>>();
182        let mut parsed = Vec::with_capacity(paths.len());
183        for handle in handles {
184            let worker_output = handle.join().map_err(|_| {
185                CodeIndexError::InvalidInput("code parser worker panicked".to_owned())
186            })??;
187            parsed.extend(worker_output);
188        }
189
190        Ok::<_, CodeIndexError>(parsed)
191    })?;
192    parsed.sort_by_key(|(index, _)| *index);
193
194    Ok(parsed.into_iter().map(|(_, build)| build).collect())
195}
196
197fn parse_one_file(
198    plan: &CodeIndexPlan,
199    path: &str,
200    bytes: &[u8],
201) -> Result<SnapshotBuild, CodeIndexError> {
202    let mut build = SnapshotBuild::new_with_scope_filters(
203        &plan.registration,
204        plan.commit.clone(),
205        plan.tree_hash.clone(),
206        SnapshotScopeFilters {
207            path_filters: plan.path_filters.clone(),
208            language_filters: plan.language_filters.clone(),
209        },
210        true,
211        plan.paths.len(),
212        0,
213    );
214    parse_indexed_file(&mut build, path, bytes)?;
215
216    Ok(build)
217}
218
219fn parse_worker_stride(
220    plan: &CodeIndexPlan,
221    paths: &[String],
222    blobs: &[Vec<u8>],
223    worker_index: usize,
224    worker_count: usize,
225) -> Result<Vec<(usize, SnapshotBuild)>, CodeIndexError> {
226    let mut parsed = Vec::new();
227    let mut index = worker_index;
228    while index < paths.len() {
229        parsed.push((index, parse_one_file(plan, &paths[index], &blobs[index])?));
230        index += worker_count;
231    }
232
233    Ok(parsed)
234}
235
236fn total_blob_bytes(blobs: &[Vec<u8>]) -> usize {
237    blobs
238        .iter()
239        .fold(0usize, |total, blob| total.saturating_add(blob.len()))
240}
241
242fn worker_count(item_count: usize, total_bytes: usize) -> usize {
243    if item_count == 0 {
244        return 0;
245    }
246    if item_count < MIN_PARALLEL_PARSE_FILES && total_bytes < MIN_PARALLEL_PARSE_BYTES {
247        return 1;
248    }
249    let desired_workers = item_count
250        .div_ceil(TARGET_PARSE_FILES_PER_WORKER)
251        .max(total_bytes.div_ceil(TARGET_PARSE_BYTES_PER_WORKER))
252        .max(1);
253
254    thread::available_parallelism()
255        .map(usize::from)
256        .unwrap_or(1)
257        .min(item_count)
258        .min(desired_workers)
259}
260
261/// Prepares a full repository index as a bounded, checkpointable batch plan.
262pub fn prepare_full_index_plan(
263    registration: CodeRepositoryRegistration,
264    selector: CodeRepositorySelector,
265    resource_budget: CodeIndexResourceBudget,
266) -> Result<CodeIndexPlan, CodeIndexError> {
267    prepare_full_index_plan_with_workspace_detection(
268        registration,
269        selector,
270        resource_budget,
271        &CodeWorkspaceDetectionConfig::default(),
272    )
273}
274
275/// Prepares a full repository index plan with caller-controlled workspace
276/// detection metadata for finalization.
277pub fn prepare_full_index_plan_with_workspace_detection(
278    registration: CodeRepositoryRegistration,
279    selector: CodeRepositorySelector,
280    resource_budget: CodeIndexResourceBudget,
281    workspace_detection: &CodeWorkspaceDetectionConfig,
282) -> Result<CodeIndexPlan, CodeIndexError> {
283    let root = PathBuf::from(&registration.root_path);
284    let snapshot = scoped_source_snapshot(&registration, &selector, &root, &selector.ref_selector)?;
285    let filesystem_path_hashes = filesystem_plan_path_hashes(&snapshot)?;
286    let source_scope = code_snapshot_scope_id(
287        &registration.repository_id,
288        &snapshot.tree_hash,
289        &snapshot.path_filters,
290        &snapshot.language_filters,
291    );
292    let workspaces = detect_workspaces_for_source_snapshot(
293        &snapshot.root,
294        snapshot.kind,
295        &snapshot.resolved_commit_sha,
296        &snapshot.entries,
297        &snapshot.path_filters,
298        workspace_detection,
299    );
300
301    Ok(CodeIndexPlan {
302        registration,
303        root: snapshot.root,
304        commit: snapshot.resolved_commit_sha,
305        tree_hash: snapshot.tree_hash,
306        source_scope,
307        path_filters: snapshot.path_filters,
308        language_filters: snapshot.language_filters,
309        source_kind: snapshot.kind,
310        filesystem_path_hashes,
311        paths: snapshot.entries,
312        workspaces,
313        cursor: 0,
314        next_batch_index: 1,
315        resource_budget,
316    })
317}
318
319fn filesystem_plan_path_hashes(
320    snapshot: &super::scope::ScopedSourceSnapshot,
321) -> Result<BTreeMap<String, String>, CodeIndexError> {
322    if !snapshot.kind.is_filesystem() {
323        return Ok(BTreeMap::new());
324    }
325    let paths = snapshot
326        .entries
327        .iter()
328        .map(|entry| entry.path.clone())
329        .collect::<Vec<_>>();
330    let path_hashes = filesystem_content_hashes_for_paths(&snapshot.root, &paths)?;
331    let tree_hash = filesystem_tree_hash_from_path_hashes(&path_hashes);
332    if tree_hash != snapshot.tree_hash {
333        return Err(CodeIndexError::InvalidInput(format!(
334            "filesystem source snapshot {} no longer matches planned filesystem content {tree_hash}",
335            snapshot.tree_hash
336        )));
337    }
338
339    Ok(path_hashes)
340}
341
342fn next_fetch_end(plan: &CodeIndexPlan, batch_file_count: usize, parsed_bytes: usize) -> usize {
343    let remaining_files = plan
344        .resource_budget
345        .max_files_per_batch
346        .saturating_sub(batch_file_count)
347        .max(1);
348    let file_limited_end = plan.paths.len().min(
349        plan.cursor
350            .saturating_add(GIT_BLOB_FETCH_GROUP.min(remaining_files)),
351    );
352    let remaining_bytes = plan
353        .resource_budget
354        .max_bytes_per_batch
355        .saturating_sub(parsed_bytes);
356    let mut byte_count = 0usize;
357    let mut end = plan.cursor;
358    while end < file_limited_end {
359        let entry_bytes = plan.paths[end].byte_count;
360        if end > plan.cursor && byte_count.saturating_add(entry_bytes) > remaining_bytes {
361            break;
362        }
363        byte_count = byte_count.saturating_add(entry_bytes);
364        end += 1;
365    }
366
367    if end == plan.cursor && batch_file_count == 0 {
368        return (plan.cursor + 1).min(plan.paths.len());
369    }
370
371    end
372}
373
374fn batch_row_count(build: &SnapshotBuild) -> usize {
375    build
376        .files
377        .len()
378        .saturating_add(build.symbols.len())
379        .saturating_add(build.references.len())
380        .saturating_add(build.imports.len())
381        .saturating_add(build.dependencies.len())
382        .saturating_add(build.feature_flags.len())
383        .saturating_add(build.routes.len())
384        .saturating_add(build.chunks.len())
385        .saturating_add(build.diagnostics.len())
386}
387
388#[cfg(test)]
389mod tests {
390    use super::*;
391
392    #[test]
393    fn parser_worker_count_keeps_tiny_batches_serial() {
394        assert_eq!(worker_count(7, 32 * 1024), 1);
395    }
396
397    #[test]
398    fn parser_worker_count_scales_with_bounded_batch_work() {
399        let available = thread::available_parallelism()
400            .map(usize::from)
401            .unwrap_or(1);
402        let workers = worker_count(96, 4 * 1024 * 1024);
403
404        assert_eq!(workers, available.min(8).min(96));
405        assert!(workers >= 1);
406    }
407
408    #[test]
409    fn parser_worker_count_caps_thread_fanout_for_small_byte_batches() {
410        let available = thread::available_parallelism()
411            .map(usize::from)
412            .unwrap_or(1);
413        let workers = worker_count(40, 128 * 1024);
414
415        assert_eq!(workers, available.min(3).min(40));
416    }
417
418    #[test]
419    fn batch_row_count_includes_feature_flags() {
420        let registration =
421            CodeRepositoryRegistration::new("repo", "fixture", "/tmp/repo", Vec::new(), Vec::new())
422                .expect("registration should validate");
423        let mut build = SnapshotBuild::new(
424            &registration,
425            "commit".to_owned(),
426            "tree".to_owned(),
427            true,
428            1,
429            0,
430        );
431        build.feature_flags = crate::code::feature_flags::extract_feature_flags(
432            crate::code::feature_flags::FeatureFlagFileInput {
433                repository_id: &build.repository_id,
434                source_scope: &build.source_scope,
435                file_id: "file",
436                path: "src/lib.rs",
437                language_id: "rust",
438                content: "if env::var(\"CHECKOUT_V2\").is_ok() && env::var(\"PAYMENTS_V2\").is_ok() {}",
439                config_facts: &[],
440            },
441        )
442        .expect("feature flags should extract");
443
444        assert_eq!(batch_row_count(&build), 2);
445    }
446}
relay_knowledge/code/index/plan.rs

relay_knowledge/code/index/
plan.rs
