relay_knowledge/code/
pipeline.rs

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

relay_knowledge/code/
pipeline.rs
