forge
ExploreLog in

Tahinli/forge

public
0 watching
★ 0
11.5 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::workspace;
20
21/// Log lines queued between the executor and the log shipper.
22const LOG_CHANNEL: usize = 256;
23/// How many times the runner tries to report a finished job.
24const FINISH_ATTEMPTS: u32 = 5;
25/// Largest single artifact the runner will upload. The forge enforces its own total cap.
26const MAX_ARTIFACT_BYTES: u64 = 512 * 1024 * 1024;
27
28/// Timing knobs. Production uses the defaults. Tests shorten them.
29#[derive(Clone, Copy, Debug)]
30pub struct RunOptions {
31 pub heartbeat_every: Duration,
32 pub claim_backoff: Duration,
33 pub pump: PumpConfig,
34 /// When set, replaces `timeout_minutes` from the job.
35 pub timeout_override: Option<Duration>,
36}
37
38impl Default for RunOptions {
39 fn default() -> Self {
40 Self {
41 heartbeat_every: Duration::from_secs(10),
42 claim_backoff: Duration::from_secs(5),
43 pump: PumpConfig::default(),
44 timeout_override: None,
45 }
46 }
47}
48
49/// A runner bound to one forge and one executor.
50#[derive(Clone)]
51pub struct Runner {
52 client: ForgeClient,
53 config: Arc<RunnerConfig>,
54 executor: Arc<dyn Executor>,
55 options: RunOptions,
56}
57
58impl Runner {
59 pub fn new(
60 config: RunnerConfig,
61 executor: Arc<dyn Executor>,
62 options: RunOptions,
63 ) -> Result<Self, ClientError> {
64 let client = ForgeClient::new(&config.url)?;
65 Ok(Self {
66 client,
67 config: Arc::new(config),
68 executor,
69 options,
70 })
71 }
72
73 pub fn client(&self) -> &ForgeClient {
74 &self.client
75 }
76
77 /// Claims and runs jobs, up to `concurrency` at a time, until `shutdown` is cancelled. Jobs
78 /// still running at shutdown are cancelled and reported as cancelled.
79 pub async fn run(&self, shutdown: CancellationToken) {
80 let semaphore = Arc::new(Semaphore::new(self.config.concurrency));
81 let mut jobs: JoinSet<()> = JoinSet::new();
82 tracing::info!(
83 name = %self.config.name,
84 concurrency = self.config.concurrency,
85 "runner started"
86 );
87 loop {
88 let permit = tokio::select! {
89 _ = shutdown.cancelled() => break,
90 permit = semaphore.clone().acquire_owned() => match permit {
91 Ok(permit) => permit,
92 Err(_) => break,
93 },
94 };
95 let claimed = tokio::select! {
96 _ = shutdown.cancelled() => {
97 drop(permit);
98 break;
99 }
100 r = self.client.claim(&self.config.runner_secret, &self.config.labels) => r,
101 };
102 match claimed {
103 Ok(Some(job)) => {
104 tracing::info!(
105 job_id = %job.job_id,
106 job_name = %job.job_name,
107 repo = %job.repo,
108 "claimed job"
109 );
110 let runner = self.clone();
111 let cancel = shutdown.child_token();
112 jobs.spawn(async move {
113 let _permit = permit;
114 runner.process(job, cancel).await;
115 });
116 }
117 Ok(None) => drop(permit),
118 Err(e) => {
119 drop(permit);
120 tracing::warn!(error = %e, "claim failed, backing off");
121 tokio::select! {
122 _ = shutdown.cancelled() => break,
123 _ = tokio::time::sleep(self.options.claim_backoff) => {}
124 }
125 }
126 }
127 while let Some(joined) = jobs.try_join_next() {
128 if let Err(e) = joined {
129 tracing::error!(error = %e, "job task panicked");
130 }
131 }
132 }
133 while let Some(joined) = jobs.join_next().await {
134 if let Err(e) = joined {
135 tracing::error!(error = %e, "job task panicked");
136 }
137 }
138 tracing::info!("runner stopped");
139 }
140
141 /// Runs one claimed job end to end and reports its result. Returns the final outcome.
142 pub async fn process(&self, job: JobAssignment, shutdown: CancellationToken) -> Outcome {
143 // Mask the job's secret values and its own token, both raw and base64 (done by Masker).
144 let mut secret_values: Vec<String> = job.secrets.values().cloned().collect();
145 secret_values.push(job.job_token.clone());
146 let masker = Masker::new(secret_values);
147
148 let (tx, rx) = mpsc::channel::<Vec<u8>>(LOG_CHANNEL);
149 let pump = {
150 let client = self.client.clone();
151 let job_id = job.job_id.clone();
152 let token = job.job_token.clone();
153 let cfg = self.options.pump;
154 tokio::spawn(async move {
155 logs::pump(&client, &job_id, &token, MaskStream::new(masker), rx, cfg).await
156 })
157 };
158
159 send_line(
160 &tx,
161 format!(
162 "forge-runner {}: job {} on runner {}, image {}, commit {}",
163 env!("CARGO_PKG_VERSION"),
164 job.job_name,
165 self.config.name,
166 job.image,
167 job.commit
168 ),
169 )
170 .await;
171
172 let timeout = self
173 .options
174 .timeout_override
175 .unwrap_or_else(|| Duration::from_secs(60 * u64::from(job.timeout_minutes)));
176 let job_cancel = shutdown.child_token();
177 let timed_out = Arc::new(AtomicBool::new(false));
178 let heartbeat = self.spawn_heartbeat(&job, job_cancel.clone());
179 let timer = {
180 let cancel = job_cancel.clone();
181 let flag = timed_out.clone();
182 tokio::spawn(async move {
183 tokio::time::sleep(timeout).await;
184 flag.store(true, Ordering::SeqCst);
185 cancel.cancel();
186 })
187 };
188
189 let mut workspace_dir = None;
190 let prepared = workspace::prepare(&self.config.work_dir, &job, &self.config.url).await;
191 let mut result = match prepared {
192 Ok(dir) => {
193 let result = self
194 .executor
195 .run(ExecRequest {
196 job: &job,
197 workspace: &dir,
198 logs: tx.clone(),
199 cancel: job_cancel.clone(),
200 })
201 .await;
202 workspace_dir = Some(dir);
203 result
204 }
205 Err(e) => ExecResult {
206 outcome: Outcome::Error,
207 exit_code: None,
208 message: Some(format!("workspace: {e}")),
209 },
210 };
211
212 // Artifacts only come from a successful run.
213 let upload_error = match workspace_dir.as_deref() {
214 Some(dir) if result.outcome == Outcome::Success && !job.artifacts.is_empty() => {
215 self.upload_artifacts(&job, dir).await.err()
216 }
217 _ => None,
218 };
219 if let Some(message) = upload_error {
220 result = ExecResult {
221 outcome: Outcome::Error,
222 exit_code: result.exit_code,
223 message: Some(message),
224 };
225 }
226
227 heartbeat.abort();
228 timer.abort();
229 if timed_out.load(Ordering::SeqCst) && result.outcome == Outcome::Cancelled {
230 result = ExecResult {
231 outcome: Outcome::Failure,
232 exit_code: None,
233 message: Some(format!("timed out after {} minutes", job.timeout_minutes)),
234 };
235 }
236
237 if let Some(message) = &result.message {
238 send_line(&tx, format!("runner: {message}")).await;
239 }
240 drop(tx);
241 match pump.await {
242 Ok(Ok(bytes)) => tracing::debug!(job_id = %job.job_id, bytes, "logs shipped"),
243 Ok(Err(e)) => tracing::warn!(job_id = %job.job_id, error = %e, "log shipping failed"),
244 Err(e) => tracing::error!(job_id = %job.job_id, error = %e, "log task panicked"),
245 }
246
247 self.finish(&job, &result).await;
248 if let Some(dir) = workspace_dir
249 && let Err(e) = workspace::cleanup(&dir).await
250 {
251 tracing::warn!(job_id = %job.job_id, error = %e, "workspace cleanup failed");
252 }
253 tracing::info!(job_id = %job.job_id, outcome = ?result.outcome, "job finished");
254 result.outcome
255 }
256
257 fn spawn_heartbeat(&self, job: &JobAssignment, cancel: CancellationToken) -> JoinHandle<()> {
258 let client = self.client.clone();
259 let job_id = job.job_id.clone();
260 let token = job.job_token.clone();
261 let every = self.options.heartbeat_every;
262 tokio::spawn(async move {
263 let mut ticker = tokio::time::interval_at(Instant::now() + every, every);
264 loop {
265 ticker.tick().await;
266 match client.heartbeat(&job_id, &token).await {
267 Ok(beat) if beat.cancel => {
268 tracing::info!(job_id = %job_id, "forge requested cancel");
269 cancel.cancel();
270 return;
271 }
272 Ok(_) => {}
273 Err(e) => tracing::warn!(job_id = %job_id, error = %e, "heartbeat failed"),
274 }
275 }
276 })
277 }
278
279 /// Collects the job's artifacts and uploads them. The error is a message for the job record.
280 async fn upload_artifacts(
281 &self,
282 job: &JobAssignment,
283 dir: &std::path::Path,
284 ) -> Result<(), String> {
285 let found = artifacts::collect(dir, &job.artifacts, MAX_ARTIFACT_BYTES)
286 .map_err(|e| format!("artifacts: {e}"))?;
287 for artifact in &found {
288 self.client
289 .upload_artifact(&job.job_id, &job.job_token, &artifact.name, &artifact.path)
290 .await
291 .map_err(|e| format!("artifact {}: {e}", artifact.name))?;
292 }
293 tracing::info!(job_id = %job.job_id, count = found.len(), "artifacts uploaded");
294 Ok(())
295 }
296
297 async fn finish(&self, job: &JobAssignment, result: &ExecResult) {
298 let req = FinishRequest {
299 outcome: result.outcome,
300 exit_code: result.exit_code,
301 message: result.message.clone(),
302 };
303 for attempt in 1..=FINISH_ATTEMPTS {
304 match self.client.finish(&job.job_id, &job.job_token, &req).await {
305 Ok(()) => return,
306 Err(e) => {
307 tracing::warn!(job_id = %job.job_id, attempt, error = %e, "finish failed");
308 tokio::time::sleep(Duration::from_secs(2 * u64::from(attempt))).await;
309 }
310 }
311 }
312 tracing::error!(
313 job_id = %job.job_id,
314 "could not report the job result; the forge will reap it after heartbeats stop"
315 );
316 }
317}
318
319async fn send_line(tx: &mpsc::Sender<Vec<u8>>, line: String) {
320 let mut bytes = line.into_bytes();
321 bytes.push(b'\n');
322 let _ = tx.send(bytes).await;
323}