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

comindly_engine/
indexer.rs
