forge
ExploreLog in

Tahinli/forge

public
0 watching
★ 0
14.7 KiB
1//! The runner loop: claim a job, run it, ship its logs and artifacts, and report the result.
2
3use std::sync::Arc;
4use std::sync::atomic::{AtomicBool, Ordering};
5use std::time::Duration;
6
7use forge_ci::mask::{MaskStream, Masker};
8use forge_ci::proto::{FinishRequest, JobAssignment, Outcome};
9use tokio::sync::{Semaphore, mpsc};
10use tokio::task::{JoinHandle, JoinSet};
11use tokio::time::Instant;
12use tokio_util::sync::CancellationToken;
13
14use crate::artifacts;
15use crate::client::{ClientError, ForgeClient};
16use crate::config::RunnerConfig;
17use crate::executor::{ExecRequest, ExecResult, Executor};
18use crate::logs::{self, PumpConfig};
19use crate::update::{self, Installed, UpdateState};
20use crate::workspace;
21
22/// Log lines queued between the executor and the log shipper.
23const LOG_CHANNEL: usize = 256;
24/// How many times the runner tries to report a finished job.
25const FINISH_ATTEMPTS: u32 = 5;
26/// Largest single artifact the runner will upload. The forge enforces its own total cap.
27const MAX_ARTIFACT_BYTES: u64 = 512 * 1024 * 1024;
28
29/// Timing knobs. Production uses the defaults. Tests shorten them.
30#[derive(Clone, Copy, Debug)]
31pub struct RunOptions {
32 pub heartbeat_every: Duration,
33 pub claim_backoff: Duration,
34 pub pump: PumpConfig,
35 /// When set, replaces `timeout_minutes` from the job.
36 pub timeout_override: Option<Duration>,
37}
38
39impl Default for RunOptions {
40 fn default() -> Self {
41 Self {
42 heartbeat_every: Duration::from_secs(10),
43 claim_backoff: Duration::from_secs(5),
44 pump: PumpConfig::default(),
45 timeout_override: None,
46 }
47 }
48}
49
50/// How [`Runner::run`] ended.
51#[derive(Clone, Copy, Debug, PartialEq, Eq)]
52pub enum RunEnd {
53 /// Shutdown was requested.
54 Stopped,
55 /// The binary was replaced. The caller exits with [`update::RESTART_EXIT_CODE`].
56 Updated,
57}
58
59/// A runner bound to one forge and one executor.
60#[derive(Clone)]
61pub struct Runner {
62 client: ForgeClient,
63 config: Arc<RunnerConfig>,
64 executor: Arc<dyn Executor>,
65 options: RunOptions,
66}
67
68impl Runner {
69 pub fn new(
70 config: RunnerConfig,
71 executor: Arc<dyn Executor>,
72 options: RunOptions,
73 ) -> Result<Self, ClientError> {
74 let client = ForgeClient::new(&config.url)?;
75 Ok(Self {
76 client,
77 config: Arc::new(config),
78 executor,
79 options,
80 })
81 }
82
83 pub fn client(&self) -> &ForgeClient {
84 &self.client
85 }
86
87 /// Claims and runs jobs, up to `concurrency` at a time, until `shutdown` is cancelled. Jobs
88 /// still running at shutdown are cancelled and reported as cancelled.
89 ///
90 /// When the forge advertises a different binary and `auto_update` is on, no new jobs are
91 /// claimed. Once the running jobs finish, the binary is replaced and the loop returns
92 /// [`RunEnd::Updated`].
93 pub async fn run(&self, shutdown: CancellationToken) -> RunEnd {
94 let semaphore = Arc::new(Semaphore::new(self.config.concurrency));
95 let mut jobs: JoinSet<()> = JoinSet::new();
96 // Hashed once at startup. Without it, or with auto_update off, updates stay disabled.
97 let exe = if self.config.auto_update {
98 match update::locate_self() {
99 Ok(found) => Some(found),
100 Err(e) => {
101 tracing::warn!(error = %e, "auto update disabled, cannot read own binary");
102 None
103 }
104 }
105 } else {
106 None
107 };
108 let mut updates = UpdateState::new(
109 self.config.auto_update,
110 exe.as_ref().map(|(_, sha)| sha.clone()),
111 );
112 tracing::info!(
113 name = %self.config.name,
114 concurrency = self.config.concurrency,
115 auto_update = updates.enabled(),
116 "runner started"
117 );
118 let end = loop {
119 if let Some((exe_path, _)) = &exe
120 && let Some(sha) = updates.pending().map(str::to_owned)
121 {
122 if !jobs.is_empty() {
123 // Claiming is paused until the running jobs finish.
124 tokio::select! {
125 _ = shutdown.cancelled() => break RunEnd::Stopped,
126 Some(joined) = jobs.join_next() => {
127 if let Err(e) = joined {
128 tracing::error!(error = %e, "job task panicked");
129 }
130 }
131 }
132 continue;
133 }
134 match update::install(&self.client, &self.config.runner_secret, exe_path, &sha)
135 .await
136 {
137 Ok(Installed::Replaced) => {
138 tracing::info!(sha = %sha, "binary updated, restarting");
139 break RunEnd::Updated;
140 }
141 Ok(Installed::Listed) => {
142 tracing::warn!(sha = %sha, "forge binary is on the bad list, not updating");
143 updates.give_up(&sha);
144 }
145 Err(e) => {
146 tracing::warn!(sha = %sha, error = %e, "update failed, keeping this binary");
147 updates.give_up(&sha);
148 }
149 }
150 continue;
151 }
152
153 let permit = tokio::select! {
154 _ = shutdown.cancelled() => break RunEnd::Stopped,
155 permit = semaphore.clone().acquire_owned() => match permit {
156 Ok(permit) => permit,
157 Err(_) => break RunEnd::Stopped,
158 },
159 };
160 let claimed = tokio::select! {
161 _ = shutdown.cancelled() => {
162 drop(permit);
163 break RunEnd::Stopped;
164 }
165 r = self.client.claim(&self.config.runner_secret, &self.config.labels) => r,
166 };
167 match claimed {
168 Ok(reply) => {
169 if updates.observe(reply.binary_sha256.as_deref())
170 && let Some(sha) = updates.pending()
171 {
172 tracing::info!(sha = %sha, "forge advertises a new binary, updating when idle");
173 }
174 match reply.job {
175 Some(job) => {
176 tracing::info!(
177 job_id = %job.job_id,
178 job_name = %job.job_name,
179 repo = %job.repo,
180 "claimed job"
181 );
182 let runner = self.clone();
183 let cancel = shutdown.child_token();
184 jobs.spawn(async move {
185 let _permit = permit;
186 runner.process(job, cancel).await;
187 });
188 }
189 None => drop(permit),
190 }
191 }
192 Err(e) => {
193 drop(permit);
194 tracing::warn!(error = %e, "claim failed, backing off");
195 tokio::select! {
196 _ = shutdown.cancelled() => break RunEnd::Stopped,
197 _ = tokio::time::sleep(self.options.claim_backoff) => {}
198 }
199 }
200 }
201 while let Some(joined) = jobs.try_join_next() {
202 if let Err(e) = joined {
203 tracing::error!(error = %e, "job task panicked");
204 }
205 }
206 };
207 while let Some(joined) = jobs.join_next().await {
208 if let Err(e) = joined {
209 tracing::error!(error = %e, "job task panicked");
210 }
211 }
212 tracing::info!("runner stopped");
213 end
214 }
215
216 /// Runs one claimed job end to end and reports its result. Returns the final outcome.
217 pub async fn process(&self, job: JobAssignment, shutdown: CancellationToken) -> Outcome {
218 // Mask the job's secret values and its own token, both raw and base64 (done by Masker).
219 let mut secret_values: Vec<String> = job.secrets.values().cloned().collect();
220 secret_values.push(job.job_token.clone());
221 let masker = Masker::new(secret_values);
222
223 let (tx, rx) = mpsc::channel::<Vec<u8>>(LOG_CHANNEL);
224 let pump = {
225 let client = self.client.clone();
226 let job_id = job.job_id.clone();
227 let token = job.job_token.clone();
228 let cfg = self.options.pump;
229 tokio::spawn(async move {
230 logs::pump(&client, &job_id, &token, MaskStream::new(masker), rx, cfg).await
231 })
232 };
233
234 send_line(
235 &tx,
236 format!(
237 "forge-runner {}: job {} on runner {}, image {}, commit {}",
238 env!("CARGO_PKG_VERSION"),
239 job.job_name,
240 self.config.name,
241 job.image,
242 job.commit
243 ),
244 )
245 .await;
246
247 let timeout = self
248 .options
249 .timeout_override
250 .unwrap_or_else(|| Duration::from_secs(60 * u64::from(job.timeout_minutes)));
251 let job_cancel = shutdown.child_token();
252 let timed_out = Arc::new(AtomicBool::new(false));
253 let heartbeat = self.spawn_heartbeat(&job, job_cancel.clone());
254 let timer = {
255 let cancel = job_cancel.clone();
256 let flag = timed_out.clone();
257 tokio::spawn(async move {
258 tokio::time::sleep(timeout).await;
259 flag.store(true, Ordering::SeqCst);
260 cancel.cancel();
261 })
262 };
263
264 let mut workspace_dir = None;
265 let prepared = workspace::prepare(&self.config.work_dir, &job, &self.config.url).await;
266 let mut result = match prepared {
267 Ok(dir) => {
268 let result = self
269 .executor
270 .run(ExecRequest {
271 job: &job,
272 workspace: &dir,
273 logs: tx.clone(),
274 cancel: job_cancel.clone(),
275 })
276 .await;
277 workspace_dir = Some(dir);
278 result
279 }
280 Err(e) => ExecResult {
281 outcome: Outcome::Error,
282 exit_code: None,
283 message: Some(format!("workspace: {e}")),
284 },
285 };
286
287 // Artifacts only come from a successful run.
288 let upload_error = match workspace_dir.as_deref() {
289 Some(dir) if result.outcome == Outcome::Success && !job.artifacts.is_empty() => {
290 self.upload_artifacts(&job, dir).await.err()
291 }
292 _ => None,
293 };
294 if let Some(message) = upload_error {
295 result = ExecResult {
296 outcome: Outcome::Error,
297 exit_code: result.exit_code,
298 message: Some(message),
299 };
300 }
301
302 heartbeat.abort();
303 timer.abort();
304 if timed_out.load(Ordering::SeqCst) && result.outcome == Outcome::Cancelled {
305 result = ExecResult {
306 outcome: Outcome::Failure,
307 exit_code: None,
308 message: Some(format!("timed out after {} minutes", job.timeout_minutes)),
309 };
310 }
311
312 if let Some(message) = &result.message {
313 send_line(&tx, format!("runner: {message}")).await;
314 }
315 drop(tx);
316 match pump.await {
317 Ok(Ok(bytes)) => tracing::debug!(job_id = %job.job_id, bytes, "logs shipped"),
318 Ok(Err(e)) => tracing::warn!(job_id = %job.job_id, error = %e, "log shipping failed"),
319 Err(e) => tracing::error!(job_id = %job.job_id, error = %e, "log task panicked"),
320 }
321
322 self.finish(&job, &result).await;
323 if let Some(dir) = workspace_dir
324 && let Err(e) = workspace::cleanup(&dir).await
325 {
326 tracing::warn!(job_id = %job.job_id, error = %e, "workspace cleanup failed");
327 }
328 tracing::info!(job_id = %job.job_id, outcome = ?result.outcome, "job finished");
329 result.outcome
330 }
331
332 fn spawn_heartbeat(&self, job: &JobAssignment, cancel: CancellationToken) -> JoinHandle<()> {
333 let client = self.client.clone();
334 let job_id = job.job_id.clone();
335 let token = job.job_token.clone();
336 let every = self.options.heartbeat_every;
337 tokio::spawn(async move {
338 let mut ticker = tokio::time::interval_at(Instant::now() + every, every);
339 loop {
340 ticker.tick().await;
341 match client.heartbeat(&job_id, &token).await {
342 Ok(beat) if beat.cancel => {
343 tracing::info!(job_id = %job_id, "forge requested cancel");
344 cancel.cancel();
345 return;
346 }
347 Ok(_) => {}
348 Err(e) => tracing::warn!(job_id = %job_id, error = %e, "heartbeat failed"),
349 }
350 }
351 })
352 }
353
354 /// Collects the job's artifacts and uploads them. The error is a message for the job record.
355 async fn upload_artifacts(
356 &self,
357 job: &JobAssignment,
358 dir: &std::path::Path,
359 ) -> Result<(), String> {
360 let found = artifacts::collect(dir, &job.artifacts, MAX_ARTIFACT_BYTES)
361 .map_err(|e| format!("artifacts: {e}"))?;
362 for artifact in &found {
363 self.client
364 .upload_artifact(&job.job_id, &job.job_token, &artifact.name, &artifact.path)
365 .await
366 .map_err(|e| format!("artifact {}: {e}", artifact.name))?;
367 }
368 tracing::info!(job_id = %job.job_id, count = found.len(), "artifacts uploaded");
369 Ok(())
370 }
371
372 async fn finish(&self, job: &JobAssignment, result: &ExecResult) {
373 let req = FinishRequest {
374 outcome: result.outcome,
375 exit_code: result.exit_code,
376 message: result.message.clone(),
377 };
378 for attempt in 1..=FINISH_ATTEMPTS {
379 match self.client.finish(&job.job_id, &job.job_token, &req).await {
380 Ok(()) => return,
381 Err(e) => {
382 tracing::warn!(job_id = %job.job_id, attempt, error = %e, "finish failed");
383 tokio::time::sleep(Duration::from_secs(2 * u64::from(attempt))).await;
384 }
385 }
386 }
387 tracing::error!(
388 job_id = %job.job_id,
389 "could not report the job result; the forge will reap it after heartbeats stop"
390 );
391 }
392}
393
394async fn send_line(tx: &mpsc::Sender<Vec<u8>>, line: String) {
395 let mut bytes = line.into_bytes();
396 bytes.push(b'\n');
397 let _ = tx.send(bytes).await;
398}