forge
ExploreLog in

Tahinli/forge

public
0 watching
★ 0
13.6 KiB
1//! Job log shipping: merge the job's output, mask secrets, batch it, and POST it with byte offsets.
2
3use std::collections::VecDeque;
4use std::time::Duration;
5
6use forge_ci::mask::MaskStream;
7use tokio::io::{AsyncBufReadExt, AsyncRead, BufReader};
8use tokio::sync::mpsc;
9use tokio::time::{Instant, MissedTickBehavior};
10
11use crate::client::{ClientError, ForgeClient, LogAck};
12
13/// A batch never grows past this many bytes. The forge accepts chunks up to 1 MiB.
14pub const MAX_BATCH: usize = 64 * 1024;
15/// A line longer than this is sent in pieces, so one line cannot grow without bound.
16const MAX_LINE: usize = 16 * 1024;
17/// How many resyncs one batch may take before the pump gives up.
18const MAX_RESYNCS: u32 = 3;
19
20/// Errors from shipping logs.
21#[derive(Debug, thiserror::Error)]
22pub enum LogError {
23 #[error(transparent)]
24 Client(#[from] ClientError),
25 #[error("forge keeps disagreeing about the log size")]
26 Resync,
27}
28
29/// Timing and retry settings for the pump.
30#[derive(Clone, Copy, Debug)]
31pub struct PumpConfig {
32 pub flush_every: Duration,
33 pub retries: u32,
34 pub retry_delay: Duration,
35}
36
37impl Default for PumpConfig {
38 fn default() -> Self {
39 Self {
40 flush_every: Duration::from_secs(1),
41 retries: 5,
42 retry_delay: Duration::from_secs(2),
43 }
44 }
45}
46
47/// Bytes kept for resends. Acknowledged bytes are trimmed past this, oldest first. Unacknowledged
48/// bytes are never trimmed (they are at most one batch).
49pub const RETAIN_CAP: usize = 8 * 1024 * 1024;
50/// Written in place of bytes the forge lost and the runner no longer holds.
51pub const GAP_MARKER: &[u8] = b"[log gap]\n";
52
53/// The job's log as the runner holds it, in forge coordinates (byte positions in the log the
54/// forge stores).
55///
56/// - `offset` is the forge's confirmed size. Bytes below it were acknowledged with a 200.
57/// - `start..end` is the retained window. It holds every byte from `start` up to `end`: the
58/// acknowledged tail kept for resends, and the unacknowledged batch after `offset`.
59/// - A 409 that reports size `S` inside the window resends from `S`. A size outside it leaves a
60/// gap, which is marked with `GAP_MARKER`.
61///
62/// Pure state, so it can be tested without a network.
63#[derive(Debug)]
64pub struct LogBuffer {
65 start: u64,
66 offset: u64,
67 data: VecDeque<u8>,
68 cap: usize,
69}
70
71impl Default for LogBuffer {
72 fn default() -> Self {
73 Self::new()
74 }
75}
76
77impl LogBuffer {
78 pub fn new() -> Self {
79 Self::with_cap(RETAIN_CAP)
80 }
81
82 /// A buffer that keeps at most `cap` bytes of acknowledged output. `cap` is at least one batch.
83 pub fn with_cap(cap: usize) -> Self {
84 Self {
85 start: 0,
86 offset: 0,
87 data: VecDeque::new(),
88 cap: cap.max(MAX_BATCH),
89 }
90 }
91
92 /// The forge's confirmed size.
93 pub fn offset(&self) -> u64 {
94 self.offset
95 }
96
97 /// The forge position of the oldest retained byte.
98 pub fn retained_start(&self) -> u64 {
99 self.start
100 }
101
102 /// The position just past the last byte the runner holds.
103 pub fn end(&self) -> u64 {
104 self.start + self.data.len() as u64
105 }
106
107 /// Bytes the forge has not acknowledged.
108 pub fn pending(&self) -> usize {
109 (self.end() - self.offset) as usize
110 }
111
112 /// The next batch to send: up to `MAX_BATCH` unacknowledged bytes from `offset`.
113 pub fn batch(&self) -> Vec<u8> {
114 let from = (self.offset - self.start) as usize;
115 let n = self.pending().min(MAX_BATCH);
116 self.data.range(from..from + n).copied().collect()
117 }
118
119 /// Room left before the pending bytes reach one batch.
120 pub fn room(&self) -> usize {
121 MAX_BATCH.saturating_sub(self.pending())
122 }
123
124 /// Appends as much of `bytes` as fits, and returns how many bytes were taken.
125 pub fn take(&mut self, bytes: &[u8]) -> usize {
126 let n = self.room().min(bytes.len());
127 self.data.extend(&bytes[..n]);
128 n
129 }
130
131 /// The batch returned by `batch` was stored. Advances `offset` and trims the acknowledged
132 /// tail to the cap.
133 pub fn acked(&mut self) {
134 self.offset += self.pending().min(MAX_BATCH) as u64;
135 self.trim();
136 }
137
138 fn trim(&mut self) {
139 let excess = self.data.len().saturating_sub(self.cap);
140 let acked = (self.offset - self.start) as usize;
141 let drop = excess.min(acked);
142 if drop > 0 {
143 self.data.drain(..drop);
144 self.start += drop as u64;
145 tracing::debug!(
146 dropped = drop,
147 "trimmed acknowledged log bytes past the retain cap"
148 );
149 }
150 }
151
152 /// The forge reports that it holds `held` bytes. Returns true when nothing is pending.
153 ///
154 /// - If `held` is inside the retained window, the runner resends from `held`.
155 /// - Otherwise the bytes between `held` and the window are unavailable. The runner writes
156 /// `GAP_MARKER` at `held`, then continues with what it still holds.
157 pub fn resync(&mut self, held: u64) -> bool {
158 if held >= self.start && held <= self.end() {
159 self.offset = held;
160 } else {
161 tracing::warn!(
162 held,
163 retained_start = self.start,
164 end = self.end(),
165 "forge log size is outside the retained window, marking a log gap"
166 );
167 let kept: VecDeque<u8> = if held < self.start {
168 // Keep the retained bytes after the marker.
169 GAP_MARKER
170 .iter()
171 .copied()
172 .chain(self.data.drain(..))
173 .collect()
174 } else {
175 GAP_MARKER.iter().copied().collect()
176 };
177 self.data = kept;
178 self.start = held;
179 self.offset = held;
180 }
181 self.pending() == 0
182 }
183}
184
185/// Ships `rx` to the forge: mask each chunk, batch it, and POST it with its offset.
186///
187/// Returns the number of bytes the forge holds for this job once `rx` is closed and the last
188/// batch is stored.
189pub async fn pump(
190 client: &ForgeClient,
191 job_id: &str,
192 job_token: &str,
193 mut stream: MaskStream,
194 mut rx: mpsc::Receiver<Vec<u8>>,
195 cfg: PumpConfig,
196) -> Result<u64, LogError> {
197 let mut buf = LogBuffer::new();
198 let mut ticker = tokio::time::interval_at(Instant::now() + cfg.flush_every, cfg.flush_every);
199 ticker.set_missed_tick_behavior(MissedTickBehavior::Delay);
200
201 loop {
202 tokio::select! {
203 msg = rx.recv() => match msg {
204 Some(bytes) => {
205 let masked = stream.push(&bytes);
206 append(&mut buf, &masked, client, job_id, job_token, &cfg).await?;
207 }
208 None => break,
209 },
210 _ = ticker.tick() => {
211 send(&mut buf, client, job_id, job_token, &cfg).await?;
212 }
213 }
214 }
215
216 let tail = stream.finish();
217 append(&mut buf, &tail, client, job_id, job_token, &cfg).await?;
218 send(&mut buf, client, job_id, job_token, &cfg).await?;
219 Ok(buf.offset())
220}
221
222async fn append(
223 buf: &mut LogBuffer,
224 mut bytes: &[u8],
225 client: &ForgeClient,
226 job_id: &str,
227 job_token: &str,
228 cfg: &PumpConfig,
229) -> Result<(), LogError> {
230 while !bytes.is_empty() {
231 if buf.room() == 0 {
232 send(buf, client, job_id, job_token, cfg).await?;
233 }
234 let n = buf.take(bytes);
235 bytes = &bytes[n..];
236 }
237 Ok(())
238}
239
240/// Sends the pending batch until the forge has it, retrying transient errors.
241async fn send(
242 buf: &mut LogBuffer,
243 client: &ForgeClient,
244 job_id: &str,
245 job_token: &str,
246 cfg: &PumpConfig,
247) -> Result<(), LogError> {
248 let mut failures = 0u32;
249 let mut resyncs = 0u32;
250 while buf.pending() > 0 {
251 let batch = buf.batch();
252 match client
253 .post_logs(job_id, job_token, buf.offset(), batch)
254 .await
255 {
256 Ok(LogAck::Accepted) => buf.acked(),
257 Ok(LogAck::Resync(held)) => {
258 resyncs += 1;
259 if resyncs > MAX_RESYNCS {
260 return Err(LogError::Resync);
261 }
262 tracing::warn!(
263 held,
264 offset = buf.offset(),
265 "log offset mismatch, resyncing"
266 );
267 buf.resync(held);
268 }
269 Err(e) => {
270 failures += 1;
271 if failures > cfg.retries {
272 return Err(e.into());
273 }
274 tracing::warn!(error = %e, failures, "log upload failed, retrying");
275 tokio::time::sleep(cfg.retry_delay * failures).await;
276 }
277 }
278 }
279 Ok(())
280}
281
282/// Reads `reader` and sends its output to `tx`, one line per message. A line longer than
283/// `MAX_LINE` is split. If the receiver goes away, the reader keeps draining so the child does
284/// not block on a full pipe.
285pub async fn forward_lines<R>(reader: R, tx: mpsc::Sender<Vec<u8>>)
286where
287 R: AsyncRead + Unpin,
288{
289 let mut reader = BufReader::new(reader);
290 let mut line: Vec<u8> = Vec::new();
291 let mut alive = true;
292 loop {
293 let (take, done) = {
294 let available = match reader.fill_buf().await {
295 Ok(b) if !b.is_empty() => b,
296 _ => break,
297 };
298 let (take, done) = match available.iter().position(|&b| b == b'\n') {
299 Some(i) => (i + 1, true),
300 None => (available.len(), false),
301 };
302 line.extend_from_slice(&available[..take]);
303 (take, done)
304 };
305 reader.consume(take);
306 if (done || line.len() >= MAX_LINE) && alive {
307 alive = tx.send(std::mem::take(&mut line)).await.is_ok();
308 }
309 }
310 if !line.is_empty() && alive {
311 let _ = tx.send(line).await;
312 }
313}
314
315#[cfg(test)]
316mod tests {
317 use super::*;
318
319 #[test]
320 fn batches_never_exceed_the_cap() {
321 let mut buf = LogBuffer::new();
322 let data = vec![b'x'; MAX_BATCH + 10];
323 let n = buf.take(&data);
324 assert_eq!(n, MAX_BATCH);
325 assert_eq!(buf.room(), 0);
326 assert_eq!(buf.batch().len(), MAX_BATCH);
327 buf.acked();
328 assert_eq!(buf.offset(), MAX_BATCH as u64);
329 assert_eq!(buf.take(&data[n..]), 10);
330 }
331
332 #[test]
333 fn resync_drops_bytes_the_forge_already_has() {
334 // We sent "abcdef" at offset 0 and the response was lost. The forge stored it.
335 let mut buf = LogBuffer::new();
336 buf.take(b"abcdef");
337 assert!(buf.resync(6));
338 assert_eq!(buf.offset(), 6);
339 assert!(buf.batch().is_empty());
340 }
341
342 #[test]
343 fn resync_resends_the_uncovered_tail() {
344 let mut buf = LogBuffer::new();
345 buf.take(b"abcdef");
346 assert!(!buf.resync(2));
347 assert_eq!(buf.offset(), 2);
348 assert_eq!(buf.batch(), b"cdef");
349 }
350
351 #[test]
352 fn resync_inside_the_window_resends_acknowledged_bytes() {
353 // "abc" was acknowledged, "def" is pending. The forge then reports size 1.
354 let mut buf = LogBuffer::new();
355 buf.take(b"abc");
356 buf.acked();
357 buf.take(b"def");
358 assert!(!buf.resync(1));
359 assert_eq!(buf.offset(), 1);
360 assert_eq!(buf.batch(), b"bcdef");
361 assert_eq!(buf.end(), 6);
362 }
363
364 #[test]
365 fn resync_ahead_of_everything_we_hold_marks_a_gap() {
366 let mut buf = LogBuffer::new();
367 buf.take(b"ab");
368 assert!(!buf.resync(10));
369 assert_eq!(buf.offset(), 10);
370 assert_eq!(buf.batch(), GAP_MARKER);
371 }
372
373 #[test]
374 fn resync_below_the_window_marks_a_gap_and_keeps_the_rest() {
375 // Two full batches are acknowledged. A one-batch cap trims the first.
376 let mut small = LogBuffer::with_cap(MAX_BATCH);
377 small.take(&vec![b'a'; MAX_BATCH]);
378 small.acked();
379 small.take(&vec![b'b'; MAX_BATCH]);
380 small.acked();
381 assert_eq!(small.retained_start(), MAX_BATCH as u64);
382 // The forge says it holds 1 byte, which is below the window. The bytes 1..start are gone.
383 assert!(!small.resync(1));
384 assert_eq!(small.offset(), 1);
385 let batch = small.batch();
386 assert_eq!(&batch[..GAP_MARKER.len()], GAP_MARKER);
387 assert!(batch[GAP_MARKER.len()..].iter().all(|&b| b == b'b'));
388 assert_eq!(small.end(), 1 + (GAP_MARKER.len() + MAX_BATCH) as u64);
389 }
390
391 #[test]
392 fn trim_drops_only_acknowledged_bytes() {
393 let mut buf = LogBuffer::with_cap(MAX_BATCH);
394 buf.take(&vec![b'a'; MAX_BATCH]);
395 buf.acked();
396 buf.take(b"tail");
397 // The cap is one batch, so the acked batch is trimmed once it is acked.
398 buf.acked();
399 assert_eq!(buf.offset(), (MAX_BATCH + 4) as u64);
400 assert!(buf.retained_start() > 0);
401 assert!(buf.retained_start() <= buf.offset());
402 assert_eq!(buf.end(), buf.offset());
403 }
404
405 #[tokio::test]
406 async fn forward_lines_splits_on_newlines() {
407 let (tx, mut rx) = mpsc::channel(16);
408 let input: &[u8] = b"one\ntwo\nthree";
409 forward_lines(input, tx).await;
410 let mut got = Vec::new();
411 while let Some(line) = rx.recv().await {
412 got.push(line);
413 }
414 assert_eq!(
415 got,
416 vec![b"one\n".to_vec(), b"two\n".to_vec(), b"three".to_vec()]
417 );
418 }
419
420 #[tokio::test]
421 async fn forward_lines_caps_long_lines() {
422 let (tx, mut rx) = mpsc::channel(64);
423 let input = vec![b'z'; MAX_LINE * 2 + 5];
424 forward_lines(&input[..], tx).await;
425 let mut total = 0;
426 while let Some(piece) = rx.recv().await {
427 assert!(piece.len() <= MAX_LINE);
428 total += piece.len();
429 }
430 assert_eq!(total, input.len());
431 }
432}