From 1d5d59341c4b15067a80745fd6ee8be04d1d581e Mon Sep 17 00:00:00 2001
From: Daniel Szoke <daniel.szoke@sentry.io>
Date: Wed, 4 Dec 2024 17:42:05 +0100
Subject: [PATCH] ref(dif): Genericize `poll_assemble`

This requires adding `should_wait` and `max_wait` functions to the `ChunkOptions` trait.
---
 src/utils/chunks/upload.rs |  8 ++++++++
 src/utils/dif_upload.rs    | 20 ++++++++++++++------
 2 files changed, 22 insertions(+), 6 deletions(-)

diff --git a/src/utils/chunks/upload.rs b/src/utils/chunks/upload.rs
index 663a7044ee..2b1718c872 100644
--- a/src/utils/chunks/upload.rs
+++ b/src/utils/chunks/upload.rs
@@ -1,3 +1,5 @@
+use std::time::Duration;
+
 /// A trait representing options for chunk uploads.
 pub trait ChunkOptions {
     /// Determines whether we need to strip debug_ids from the requests.
@@ -11,4 +13,10 @@ pub trait ChunkOptions {
 
     /// Returns the project that we are uploading to.
     fn project(&self) -> &str;
+
+    /// Returns whether we should wait for assembling to complete.
+    fn should_wait(&self) -> bool;
+
+    /// Returns the maximum wait time for the upload to complete.
+    fn max_wait(&self) -> Duration;
 }
diff --git a/src/utils/dif_upload.rs b/src/utils/dif_upload.rs
index a8c70c5aaf..df8f0c91ea 100644
--- a/src/utils/dif_upload.rs
+++ b/src/utils/dif_upload.rs
@@ -1374,7 +1374,7 @@ fn render_detail(detail: &Option<String>, fallback: Option<&str>) {
 /// missing chunks in the assemble response, this likely indicates a bug in the server.
 fn poll_assemble<T>(
     chunked_objects: &[&Chunked<T>],
-    options: &DifUpload,
+    options: &impl ChunkOptions,
 ) -> Result<(Vec<DebugInfoFile>, bool)>
 where
     T: Display + Assemblable,
@@ -1392,14 +1392,14 @@ where
     let assemble_start = Instant::now();
 
     let mut request: AssembleDifsRequest<'_> = chunked_objects.iter().copied().collect();
-    if !options.pdbs_allowed {
+    if options.should_strip_debug_ids() {
         request.strip_debug_ids();
     }
 
     let response = loop {
         let response =
             api.authenticated()?
-                .assemble_difs(&options.org, &options.project, &request)?;
+                .assemble_difs(options.org(), options.project(), &request)?;
 
         let chunks_missing = response
             .values()
@@ -1415,11 +1415,11 @@ where
         // Poll until there is a response, unless the user has specified to skip polling. In
         // that case, we return the potentially partial response from the server. This might
         // still contain a cached error.
-        if !options.wait {
+        if !options.should_wait() {
             break response;
         }
 
-        if assemble_start.elapsed() > options.max_wait {
+        if assemble_start.elapsed() > options.max_wait() {
             break response;
         }
 
@@ -1446,7 +1446,7 @@ where
 
     let (errors, mut successes): (Vec<_>, _) = response
         .into_iter()
-        .partition(|(_, r)| r.state.is_err() || options.wait && r.state.is_pending());
+        .partition(|(_, r)| r.state.is_err() || options.should_wait() && r.state.is_pending());
 
     // Print a summary of all successes first, so that errors show up at the
     // bottom for the user
@@ -2108,4 +2108,12 @@ impl ChunkOptions for DifUpload {
     fn project(&self) -> &str {
         &self.project
     }
+
+    fn should_wait(&self) -> bool {
+        self.wait
+    }
+
+    fn max_wait(&self) -> Duration {
+        self.max_wait
+    }
 }