Skip to content

Commit 966da3e

Browse files
helqhelq
committed
fetcher: implement non-blocking in-memory delay retry queue
Implement an elegant In-Memory Delay Vector queue inside FetchAgent to handle transient fetch/extraction failures (such as Gerrit replication delays) with natural backoff and max 3 retries. Failed requests are parked with exponential backoff (20s, 40s, 80s) and promoted back to the active queue on every 10-second tick. Includes full duplicate-purging, complete state hygiene on completion/exhaustion, and a virtual-time compliant unit test using sub-second intervals. TAG=agy Signed-off-by: Elkin Cruz <elkin@google.com>
1 parent 5c1ce97 commit 966da3e

2 files changed

Lines changed: 229 additions & 47 deletions

File tree

‎Cargo.toml‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -40,7 +40,7 @@ tempfile = "3.24.0"
4040
termcolor = "1.4.1"
4141
thiserror = "2.0.17"
4242
tiktoken-rs = "0.6.0"
43-
tokio = { version = "1.48.0", features = ["full"] }
43+
tokio = { version = "1.48.0", features = ["full", "test-util"] }
4444
tokio-util = { version = "0.7.17", features = ["codec", "io"] }
4545
toml = "0.8"
4646
tower-http = { version = "0.6.8", features = ["cors", "fs", "trace"] }

‎src/fetcher.rs‎

Lines changed: 228 additions & 46 deletions
Original file line numberDiff line numberDiff line change
@@ -29,10 +29,19 @@ pub struct FetchRequest {
2929
pub commit_hash: String,
3030
}
3131

32+
#[derive(Clone)]
33+
struct DelayedRequest {
34+
repo_url: Option<String>,
35+
commit_hash: String,
36+
run_at: std::time::Instant,
37+
}
38+
3239
pub struct FetchAgent {
3340
repo_path: PathBuf,
3441
rx: mpsc::Receiver<FetchRequest>,
3542
main_tx: mpsc::Sender<Event>,
43+
tick_interval: Duration,
44+
backoff_base: Duration,
3645
}
3746

3847
impl FetchAgent {
@@ -46,15 +55,25 @@ impl FetchAgent {
4655
repo_path,
4756
rx,
4857
main_tx,
58+
tick_interval: Duration::from_secs(10),
59+
backoff_base: Duration::from_secs(10),
4960
},
5061
tx,
5162
)
5263
}
5364

65+
pub fn with_intervals(mut self, tick_interval: Duration, backoff_base: Duration) -> Self {
66+
self.tick_interval = tick_interval;
67+
self.backoff_base = backoff_base;
68+
self
69+
}
70+
5471
pub async fn run(mut self) {
5572
info!("FetchAgent started");
5673
let mut queue: HashMap<Option<String>, HashSet<String>> = HashMap::new();
57-
let mut ticker = interval(Duration::from_secs(10));
74+
let mut delayed_retries: Vec<DelayedRequest> = Vec::new(); // In-memory delay vector
75+
let mut attempts_tracker: HashMap<String, u32> = HashMap::new(); // Local attempt tracker
76+
let mut ticker = interval(self.tick_interval); // Configurable batch timer
5877

5978
loop {
6079
tokio::select! {
@@ -64,15 +83,41 @@ impl FetchAgent {
6483
.insert(req.commit_hash);
6584
}
6685
_ = ticker.tick() => {
86+
let now = std::time::Instant::now();
87+
let mut expired_retries = Vec::new();
88+
89+
// Filter out expired retries
90+
delayed_retries.retain(|item| {
91+
if now >= item.run_at {
92+
expired_retries.push(item.clone());
93+
false // Remove from delayed list (promoted)
94+
} else {
95+
true // Keep in delayed list
96+
}
97+
});
98+
99+
// Promote expired retries into the active batch queue
100+
for retry in expired_retries {
101+
queue.entry(retry.repo_url)
102+
.or_default()
103+
.insert(retry.commit_hash);
104+
}
105+
106+
// Process active batch queue
67107
if !queue.is_empty() {
68-
self.process_queue(&mut queue).await;
108+
self.process_queue(&mut queue, &mut delayed_retries, &mut attempts_tracker).await;
69109
}
70110
}
71111
}
72112
}
73113
}
74114

75-
async fn process_queue(&self, queue: &mut HashMap<Option<String>, HashSet<String>>) {
115+
async fn process_queue(
116+
&self,
117+
queue: &mut HashMap<Option<String>, HashSet<String>>,
118+
delayed_retries: &mut Vec<DelayedRequest>,
119+
attempts_tracker: &mut HashMap<String, u32>,
120+
) {
76121
info!("Processing fetch queue with {} repos", queue.len());
77122

78123
for (url_opt, commits) in queue.drain() {
@@ -102,9 +147,9 @@ impl FetchAgent {
102147
"All commits present locally, skipping fetch for {}",
103148
url_display
104149
);
105-
} else if let Some(url) = url_opt {
150+
} else if let Some(ref url) = url_opt {
106151
// Remote fetch logic
107-
let remote_name = self.get_remote_name(&url);
152+
let remote_name = self.get_remote_name(url);
108153

109154
// Check if repo is local (same as self.repo_path)
110155
let is_local = {
@@ -126,17 +171,17 @@ impl FetchAgent {
126171
);
127172
// Do not continue here; let it fall through to Step 3 where it will fail individually
128173
} else {
129-
if let Err(e) = self.ensure_remote(&remote_name, &url).await {
174+
if let Err(e) = self.ensure_remote(&remote_name, url).await {
130175
error!("Failed to ensure remote {}: {}", url, e);
131176
for commit in &missing_commits {
132-
let _ = self
133-
.main_tx
134-
.send(Event::IngestionFailed {
135-
article_id: commit.clone(),
136-
error: format!("Failed to set up remote {}: {}", url, e),
137-
source: MessageSource::GitFetch,
138-
})
139-
.await;
177+
self.handle_fetch_failure(
178+
&url_opt,
179+
commit,
180+
&format!("Failed to set up remote {}: {}", url, e),
181+
delayed_retries,
182+
attempts_tracker,
183+
)
184+
.await;
140185
}
141186
continue;
142187
}
@@ -151,14 +196,14 @@ impl FetchAgent {
151196
if let Err(e) = self.fetch_all(&remote_name).await {
152197
error!("Full fetch failed for {}: {}", url, e);
153198
for commit in &missing_commits {
154-
let _ = self
155-
.main_tx
156-
.send(Event::IngestionFailed {
157-
article_id: commit.clone(),
158-
error: format!("Failed to fetch from {}: {}", url, e),
159-
source: MessageSource::GitFetch,
160-
})
161-
.await;
199+
self.handle_fetch_failure(
200+
&url_opt,
201+
commit,
202+
&format!("Failed to fetch from {}: {}", url, e),
203+
delayed_retries,
204+
attempts_tracker,
205+
)
206+
.await;
162207
}
163208
continue;
164209
}
@@ -182,14 +227,14 @@ impl FetchAgent {
182227
{
183228
Ok(shas) => shas,
184229
Err(e) => {
185-
let _ = self
186-
.main_tx
187-
.send(Event::IngestionFailed {
188-
article_id: range.clone(),
189-
error: format!("Failed to resolve git range: {}", e),
190-
source: MessageSource::GitFetch,
191-
})
192-
.await;
230+
self.handle_fetch_failure(
231+
&url_opt,
232+
&commit_or_range,
233+
&format!("Failed to resolve git range: {}", e),
234+
delayed_retries,
235+
attempts_tracker,
236+
)
237+
.await;
193238
continue;
194239
}
195240
};
@@ -217,26 +262,33 @@ impl FetchAgent {
217262
}
218263
}
219264
}
265+
// Range Success! Clean up attempts tracker and delay queue for the range key
266+
attempts_tracker.remove(&commit_or_range);
267+
delayed_retries.retain(|item| item.commit_hash != commit_or_range);
220268
info!("Successfully submitted remote range {}", range);
221269
} else {
222270
// Single commit
223271
let full_sha = match self.resolve_sha(&commit_or_range).await {
224272
Ok(sha) => sha,
225273
Err(e) => {
226-
let _ = self
227-
.main_tx
228-
.send(Event::IngestionFailed {
229-
article_id: commit_or_range.clone(),
230-
error: format!("Failed to resolve SHA: {}", e),
231-
source: MessageSource::GitFetch,
232-
})
233-
.await;
274+
self.handle_fetch_failure(
275+
&url_opt,
276+
&commit_or_range,
277+
&format!("Failed to resolve SHA: {}", e),
278+
delayed_retries,
279+
attempts_tracker,
280+
)
281+
.await;
234282
continue;
235283
}
236284
};
237285

238286
match self.extract_patch(&full_sha, &commit_or_range, 1, 1).await {
239287
Ok(mut event) => {
288+
// Success! Clean up attempts tracker and the parked delay queue
289+
attempts_tracker.remove(&commit_or_range);
290+
delayed_retries.retain(|item| item.commit_hash != commit_or_range);
291+
240292
if let Event::PatchSubmitted {
241293
ref mut message_id, ..
242294
} = event
@@ -251,14 +303,14 @@ impl FetchAgent {
251303
}
252304
Err(e) => {
253305
error!("Failed to extract patch {}: {}", commit_or_range, e);
254-
let _ = self
255-
.main_tx
256-
.send(Event::IngestionFailed {
257-
article_id: commit_or_range,
258-
error: format!("Failed to extract patch: {}", e),
259-
source: MessageSource::GitFetch,
260-
})
261-
.await;
306+
self.handle_fetch_failure(
307+
&url_opt,
308+
&commit_or_range,
309+
&format!("Failed to extract patch: {}", e),
310+
delayed_retries,
311+
attempts_tracker,
312+
)
313+
.await;
262314
}
263315
}
264316
}
@@ -439,6 +491,61 @@ impl FetchAgent {
439491
total,
440492
})
441493
}
494+
495+
async fn handle_fetch_failure(
496+
&self,
497+
url_opt: &Option<String>,
498+
commit: &str,
499+
error_msg: &str,
500+
delayed_retries: &mut Vec<DelayedRequest>,
501+
attempts_tracker: &mut HashMap<String, u32>,
502+
) {
503+
let attempts = attempts_tracker.get(commit).cloned().unwrap_or(0);
504+
let max_retries = 3;
505+
506+
if attempts < max_retries {
507+
let next_attempt = attempts + 1;
508+
attempts_tracker.insert(commit.to_string(), next_attempt);
509+
510+
// Exponential backoff using configurable base: base * (2 ^ attempts)
511+
let delay = self.backoff_base * (2u32.pow(next_attempt));
512+
let run_at = std::time::Instant::now() + delay;
513+
514+
warn!(
515+
"Fetch/extract failed for {} (attempts: {}/{}). Parking in delay queue for {}s.",
516+
commit,
517+
next_attempt,
518+
max_retries,
519+
delay.as_secs()
520+
);
521+
522+
// Deduplicate inside delayed_retries first to prevent duplicate parked entries!
523+
delayed_retries.retain(|item| item.commit_hash != commit);
524+
525+
delayed_retries.push(DelayedRequest {
526+
repo_url: url_opt.clone(),
527+
commit_hash: commit.to_string(),
528+
run_at,
529+
});
530+
} else {
531+
// Retries exhausted, clean up attempts tracker and the parked delay queue
532+
attempts_tracker.remove(commit);
533+
delayed_retries.retain(|item| item.commit_hash != commit);
534+
error!(
535+
"Fetch/extract failed for {} after {} retries. Marking as Failed. Error: {}",
536+
commit, max_retries, error_msg
537+
);
538+
539+
let _ = self
540+
.main_tx
541+
.send(Event::IngestionFailed {
542+
article_id: commit.to_string(),
543+
error: error_msg.to_string(),
544+
source: MessageSource::GitFetch,
545+
})
546+
.await;
547+
}
548+
}
442549
}
443550

444551
#[cfg(test)]
@@ -590,4 +697,79 @@ mod tests {
590697

591698
Ok(())
592699
}
700+
701+
#[tokio::test]
702+
async fn test_fetch_agent_delay_queue_retry() -> Result<()> {
703+
let temp_dir = tempfile::tempdir()?;
704+
let repo_path = temp_dir.path().to_path_buf();
705+
706+
// Setup empty dummy repo
707+
Command::new("git")
708+
.current_dir(&repo_path)
709+
.arg("init")
710+
.output()
711+
.await?;
712+
713+
let (event_tx, mut event_rx) = mpsc::channel(10);
714+
715+
// Create agent with tiny sub-second intervals for fast real-time testing!
716+
let (agent, fetch_tx) = {
717+
let (a, tx) = FetchAgent::new(repo_path.clone(), event_tx);
718+
(
719+
a.with_intervals(Duration::from_millis(50), Duration::from_millis(50)),
720+
tx,
721+
)
722+
};
723+
724+
// Spawn the agent in the background
725+
let agent_handle = tokio::spawn(agent.run());
726+
727+
// Enqueue a FetchRequest for a missing commit (which will fail resolve_sha!)
728+
let req = FetchRequest {
729+
repo_url: None,
730+
commit_hash: "0123456789abcdef0123456789abcdef01234567".to_string(),
731+
};
732+
fetch_tx.send(req).await?;
733+
734+
// Let the first tick (50ms) process. Wait 100ms to be safe.
735+
tokio::time::sleep(Duration::from_millis(100)).await;
736+
737+
// The first tick ran process_queue, failed, and parked it in delayed_retries (Attempts = 1, delay = 100ms)
738+
// Verify no event is received yet.
739+
tokio::select! {
740+
Some(event) = event_rx.recv() => {
741+
panic!("Received unexpected event: {:?}", event);
742+
}
743+
_ = tokio::time::sleep(Duration::from_millis(50)) => {
744+
// Expected: no event sent!
745+
}
746+
}
747+
748+
// Wait for retries to exhaust:
749+
// Attempt 1: delay 100ms.
750+
// Attempt 2: delay 200ms.
751+
// Attempt 3: delay 400ms.
752+
// Total delay is 700ms. Plus ticks processing time, we sleep 1.2 seconds to be completely safe.
753+
tokio::time::sleep(Duration::from_millis(1200)).await;
754+
755+
// Verify that Event::IngestionFailed is finally received!
756+
tokio::select! {
757+
Some(Event::IngestionFailed { article_id, error, .. }) = event_rx.recv() => {
758+
assert_eq!(article_id, "0123456789abcdef0123456789abcdef01234567");
759+
assert!(
760+
error.contains("Failed to resolve SHA") || error.contains("Failed to extract patch"),
761+
"Actual error was: {}",
762+
error
763+
);
764+
}
765+
_ = tokio::time::sleep(Duration::from_millis(500)) => {
766+
panic!("Timed out waiting for IngestionFailed event after retries!");
767+
}
768+
}
769+
770+
// Clean up the agent
771+
agent_handle.abort();
772+
773+
Ok(())
774+
}
593775
}

0 commit comments

Comments
 (0)