comindly_engine/
indexer.rs

1//! Index jobs run in a child `comindly index-job`, so the memory a run takes goes back when it ends.
2
3use std::collections::BTreeMap;
4use std::process::Stdio;
5use std::sync::Arc;
6
7use anyhow::{Context, Result};
8use serde::{Deserialize, Serialize};
9use tokio::io::{AsyncBufReadExt, AsyncReadExt, AsyncWriteExt, BufReader};
10
11use crate::context::{Cancelled, IndexContext, Settings};
12use crate::events::{Event, EventSink, Kind, Level};
13
14/// The hidden subcommand a server starts for each index job.
15pub const COMMAND: &str = "index-job";
16
17/// What the child indexes, sent on its stdin so no secret shows in its arguments or environment.
18#[derive(Serialize, Deserialize)]
19pub struct Spec {
20    pub root: String,
21    pub repos: Vec<crate::config::RepoEntry>,
22    pub llm: bool,
23    pub llm_key: Option<String>,
24    /// Stored git tokens by host.
25    pub git: BTreeMap<String, String>,
26    /// LLM tasks to write again from scratch.
27    #[serde(default)]
28    pub rewrite: Vec<crate::prompts::Task>,
29    /// Folders to index as they are, instead of `repos` at their default branches.
30    #[serde(default)]
31    pub paths: Vec<String>,
32    /// The LLM steps to run; none: the team settings decide.
33    #[serde(default)]
34    pub steps: Vec<crate::prompts::Task>,
35}
36
37/// Run `spec` in a child process, its events going to `job`; cancelling `job` stops it.
38pub async fn run(spec: &Spec, job: Arc<crate::jobs::Running>) -> Result<()> {
39    let mut cmd = tokio::process::Command::new(std::env::current_exe()?);
40    cmd.arg(COMMAND)
41        .stdin(Stdio::piped())
42        .stdout(Stdio::piped())
43        .stderr(Stdio::piped())
44        .kill_on_drop(true);
45    #[cfg(unix)]
46    cmd.process_group(0);
47    let mut child = cmd.spawn().context("start the indexer")?;
48    let pid = child.id();
49    // Kept open until the end: the child stops when it closes.
50    let mut stdin = child.stdin.take().context("indexer stdin")?;
51    stdin.write_all(&serde_json::to_vec(spec)?).await?;
52    stdin.write_all(b"\n").await?;
53    let out = child.stdout.take().context("indexer stdout")?;
54    let err = child.stderr.take().context("indexer stderr")?;
55    let events = async {
56        let mut failure = None;
57        let mut lines = BufReader::new(out).lines();
58        while let Ok(Some(l)) = lines.next_line().await {
59            match serde_json::from_str::<Event>(&l) {
60                Ok(e) => {
61                    if e.level == Level::Error {
62                        failure = Some(e.message.clone());
63                    }
64                    job.emit(e);
65                }
66                Err(_) if !l.trim().is_empty() => job.emit(Event::new(Level::Info, Kind::Note, l)),
67                Err(_) => {}
68            }
69        }
70        failure
71    };
72    let logs = async {
73        let mut lines = BufReader::new(err).lines();
74        while let Ok(Some(l)) = lines.next_line().await {
75            crate::telemetry::write(&l);
76        }
77    };
78    let exit = async {
79        tokio::select! {
80            s = child.wait() => s,
81            _ = job.cancel.cancelled() => {
82                signal(pid, Signal::Term);
83                match tokio::time::timeout(std::time::Duration::from_secs(60), child.wait()).await {
84                    Ok(s) => s,
85                    Err(_) => {
86                        signal(pid, Signal::Kill);
87                        child.wait().await
88                    }
89                }
90            }
91        }
92    };
93    let (failure, (), status) = tokio::join!(events, logs, exit);
94    drop(stdin);
95    let status = status.context("wait for the indexer")?;
96    if status.success() {
97        return Ok(());
98    }
99    if job.cancel.is_cancelled() {
100        return Err(Cancelled.into());
101    }
102    Err(anyhow::anyhow!(failure.unwrap_or_else(|| stopped(&status))))
103}
104
105fn stopped(status: &std::process::ExitStatus) -> String {
106    #[cfg(unix)]
107    if std::os::unix::process::ExitStatusExt::signal(status) == Some(9) {
108        return "the indexer was killed, most likely for running out of memory".into();
109    }
110    format!("the indexer stopped ({status})")
111}
112
113enum Signal {
114    Term,
115    Kill,
116}
117
118/// Ask the child to stop between steps, or kill its whole group, git and SCIP indexers included.
119fn signal(pid: Option<u32>, sig: Signal) {
120    #[cfg(unix)]
121    if let Some(pid) = pid.and_then(|p| i32::try_from(p).ok()) {
122        let (to, sig) = match sig {
123            Signal::Term => (pid, libc::SIGTERM),
124            Signal::Kill => (-pid, libc::SIGKILL),
125        };
126        // SAFETY: kill(2) only sends a signal; a negative pid addresses the group.
127        unsafe {
128            libc::kill(to, sig);
129        }
130    }
131    #[cfg(not(unix))]
132    let _ = (pid, sig);
133}
134
135/// Events as JSON lines on stdout, for the server that started this process.
136struct Stdout;
137
138impl EventSink for Stdout {
139    fn emit(&self, e: Event) {
140        use std::io::Write;
141        if let Ok(line) = serde_json::to_string(&e) {
142            let mut out = std::io::stdout().lock();
143            let _ = writeln!(out, "{line}");
144            let _ = out.flush();
145        }
146    }
147}
148
149/// The child's side: read the job from stdin, index, and report on stdout.
150pub fn child() -> Result<()> {
151    crate::rt::shared().block_on(async {
152        let mut stdin = BufReader::new(tokio::io::stdin());
153        let mut line = String::new();
154        stdin.read_line(&mut line).await?;
155        let spec: Spec = serde_json::from_str(&line).context("read the index job")?;
156        let team = crate::team::load(&spec.root)
157            .await
158            .map(|t| t.settings)
159            .unwrap_or_default();
160        let mut settings = Settings::resolve(&team, crate::config::env);
161        settings.rewrite = spec.rewrite.iter().copied().collect();
162        if !spec.llm {
163            settings.enrich = false;
164            settings.flows = false;
165            settings.guide = false;
166        } else if !spec.steps.is_empty() {
167            use crate::prompts::Task;
168            settings.enrich = spec.steps.contains(&Task::Summary);
169            settings.flows = spec.steps.contains(&Task::Flow);
170            settings.guide = spec.steps.contains(&Task::Guide);
171        }
172        let client = if settings.enrich || settings.flows || settings.guide {
173            crate::llm::LlmClient::new(
174                settings.llm_base_url.as_deref(),
175                &settings.llm_model,
176                spec.llm_key.as_deref(),
177            )
178            .ok()
179            .map(|c| c.with_models(settings.task_models.clone()))
180        } else {
181            None
182        };
183        let scip = settings.scip;
184        let git = spec.git;
185        let creds = move |host: &str| {
186            git.get(&host.to_lowercase())
187                .cloned()
188                .or_else(|| crate::config::env_token(host))
189        };
190        let ctx = IndexContext::new(
191            spec.root,
192            settings,
193            Arc::new(Stdout),
194            crate::config::cache_dir(),
195        )
196        .with_llm(client)
197        .with_git(Arc::new(creds));
198        let cancel = ctx.cancel.clone();
199        tokio::spawn(async move {
200            let mut rest = Vec::new();
201            #[cfg(unix)]
202            {
203                use tokio::signal::unix::{signal, SignalKind};
204                let mut term = signal(SignalKind::terminate()).ok();
205                tokio::select! {
206                    _ = async { match term.as_mut() { Some(t) => { t.recv().await; } None => std::future::pending().await } } => {}
207                    _ = stdin.read_to_end(&mut rest) => {}
208                }
209            }
210            #[cfg(not(unix))]
211            let _ = stdin.read_to_end(&mut rest).await;
212            cancel.cancel();
213        });
214        let run = async {
215            let sources = if spec.paths.is_empty() {
216                crate::pipeline::checkout(&ctx, &crate::config::expanded(&spec.repos)).await?
217            } else {
218                spec.paths.iter().map(crate::pipeline::Source::at).collect()
219            };
220            crate::pipeline::link(&ctx, sources, true, scip).await
221        };
222        match run.await {
223            Ok(_) => Ok(()),
224            Err(e) if e.is::<Cancelled>() => {
225                ctx.report.warn("Cancelled; what was already saved stays.");
226                Err(e)
227            }
228            Err(e) => {
229                ctx.report.error(crate::storage::explain(&format!("{e:#}")));
230                Err(e)
231            }
232        }
233    })
234}
comindly_engine/indexer.rs

comindly_engine/
indexer.rs
