relay_knowledge/code/index/plan/
mod.rs

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

relay_knowledge/code/index/plan/
mod.rs
