Address more review comments

sadhansood · sadhansood · commit a3adba86ffdb · 2025-12-19T15:36:03.000-08:00
diff --git a/crates/walrus-sdk/src/client.rs b/crates/walrus-sdk/src/client.rs
@@ -206,6 +206,13 @@ struct SendBlobOptions<'a> {
     initial_completed_weight: Option<&'a HashMap<BlobId, usize>>,
 }
 
+/// Pending upload state carried through a single reserve+store execution.
+#[derive(Debug, Default, Clone)]
+struct PendingUploadContext {
+    initial_upload_weight: Option<HashMap<BlobId, usize>>,
+    initial_failed_nodes: Option<HashMap<BlobId, Vec<NodeIndex>>>,
+}
+
 struct PendingUploadHandle<'a> {
     cancel: CancellationToken,
     future: PendingUploadFuture<'a>,
@@ -1019,7 +1026,7 @@ impl WalrusNodeClient<SuiContractClient> {
             return Ok(vec![]);
         }
 
-        let mut store_args = self.pending_upload_store_args(store_args);
+        let store_args = self.pending_upload_store_args(store_args);
 
         tracing::info!(
             "writing {blobs_count} blob{} to Walrus",
@@ -1069,12 +1076,8 @@ impl WalrusNodeClient<SuiContractClient> {
         tracing::debug!(?registered_blobs);
         store_args.maybe_observe_store_operation(store_op_duration);
 
-        store_args = self.apply_pending_upload_outcome(
-            store_args,
-            pending_upload_result,
-            &pending_blobs,
-            &committees,
-        );
+        let pending_context =
+            self.apply_pending_upload_outcome(pending_upload_result, &pending_blobs, &committees);
 
         let (mut final_result, blobs_awaiting_upload, mut blobs_pending_certify_and_extend) =
             Self::partition_registered_blobs(registered_blobs, blobs_count);
@@ -1097,9 +1100,11 @@ impl WalrusNodeClient<SuiContractClient> {
             // has changed in the meantime.
             // This operation can be safely interrupted as it does not require a wallet.
             blobs_with_certificates = self
-                .await_while_checking_notification(
-                    self.get_all_blob_certificates(blobs_awaiting_upload, &store_args),
-                )
+                .await_while_checking_notification(self.get_all_blob_certificates(
+                    blobs_awaiting_upload,
+                    &store_args,
+                    &pending_context,
+                ))
                 .await?;
 
             debug_assert_eq!(blobs_with_certificates.len(), num_to_be_certified);
@@ -1141,7 +1146,7 @@ impl WalrusNodeClient<SuiContractClient> {
         encoded_blobs_with_status: &[WalrusStoreBlobMaybeFinished<BlobWithStatus>],
         store_args: &StoreArgs,
     ) -> Vec<(VerifiedBlobMetadataWithId, Arc<Vec<SliverPair>>)> {
-        let optimistic_max_blob_bytes = self.optimistic_upload_max_blob_bytes();
+        let pending_upload_max_blob_bytes = self.pending_upload_max_blob_bytes();
         let pending_uploads_enabled = self.config.communication_config.pending_uploads_enabled;
 
         let pending: Vec<_> = encoded_blobs_with_status
@@ -1150,16 +1155,16 @@ impl WalrusNodeClient<SuiContractClient> {
                 blob.pending_upload_payload(
                     pending_uploads_enabled,
                     store_args,
-                    optimistic_max_blob_bytes,
+                    pending_upload_max_blob_bytes,
                 )
             })
             .collect();
 
         tracing::debug!(
             pending_candidates = pending.len(),
             pending_enabled = pending_uploads_enabled,
-            optimistic_uploads = store_args.store_optimizations.optimistic_uploads_enabled(),
-            max_blob_bytes = optimistic_max_blob_bytes,
+            pending_uploads = store_args.store_optimizations.pending_uploads_enabled(),
+            max_blob_bytes = pending_upload_max_blob_bytes,
             "computed pending upload candidates",
         );
 
@@ -1250,11 +1255,11 @@ impl WalrusNodeClient<SuiContractClient> {
 
     fn apply_pending_upload_outcome(
         &self,
-        mut store_args: StoreArgs,
         mut pending_upload_result: Option<RunOutput<Vec<BlobId>, StoreError>>,
         pending_blobs: &[(VerifiedBlobMetadataWithId, Arc<Vec<SliverPair>>)],
         committees: &ActiveCommittees,
-    ) -> StoreArgs {
+    ) -> PendingUploadContext {
+        let mut context = PendingUploadContext::default();
         // Cancel any in-flight pending tail. Tail handle is returned during the upload with
         // immediate intent later.
         if let Some(ref mut pending_output) = pending_upload_result
@@ -1264,9 +1269,9 @@ impl WalrusNodeClient<SuiContractClient> {
         }
 
         if let Some(pending_output) = pending_upload_result {
-            let weight_map = Self::success_weight_by_blob(&pending_output.results, &HashSet::new());
+            let weight_map = Self::success_weight_by_blob(&pending_output.results);
             if !weight_map.is_empty() {
-                store_args = store_args.with_initial_upload_weight(weight_map);
+                context.initial_upload_weight = Some(weight_map);
             }
             let mut failed_nodes = failed_node_indices(&pending_output.results);
             // Add nodes that were not seen in the pending upload results to the failed nodes.
@@ -1286,11 +1291,11 @@ impl WalrusNodeClient<SuiContractClient> {
                         .iter()
                         .map(|(meta, _)| (*meta.blob_id(), failed_nodes.clone())),
                 );
-                store_args = store_args.with_initial_failed_nodes(failed_map);
+                context.initial_failed_nodes = Some(failed_map);
             }
         }
 
-        store_args
+        context
     }
 
     fn partition_registered_blobs(
@@ -1365,6 +1370,7 @@ impl WalrusNodeClient<SuiContractClient> {
         &self,
         blobs_to_be_certified: Vec<WalrusStoreBlobUnfinished<BlobAwaitingUpload>>,
         store_args: &StoreArgs,
+        pending_context: &PendingUploadContext,
     ) -> ClientResult<Vec<WalrusStoreBlobMaybeFinished<BlobPendingCertifyAndExtend>>> {
         if blobs_to_be_certified.is_empty() {
             return Ok(vec![]);
@@ -1377,8 +1383,13 @@ impl WalrusNodeClient<SuiContractClient> {
             |blob_to_be_certified| {
                 let multi_pb = Arc::clone(&multi_pb);
                 async move {
-                    self.get_certificate(blob_to_be_certified, multi_pb.as_ref(), store_args)
-                        .await
+                    self.get_certificate(
+                        blob_to_be_certified,
+                        multi_pb.as_ref(),
+                        store_args,
+                        pending_context,
+                    )
+                    .await
                 }
             },
         ))
@@ -1401,6 +1412,7 @@ impl WalrusNodeClient<SuiContractClient> {
         blob_to_be_certified: WalrusStoreBlobUnfinished<BlobAwaitingUpload>,
         multi_pb: &MultiProgress,
         store_args: &StoreArgs,
+        pending_context: &PendingUploadContext,
     ) -> ClientResult<WalrusStoreBlobMaybeFinished<BlobPendingCertifyAndExtend>> {
         let committees = self.get_committees().await?;
 
@@ -1447,6 +1459,7 @@ impl WalrusNodeClient<SuiContractClient> {
                             &blob_object.blob_persistence_type(),
                             Some(multi_pb),
                             store_args,
+                            pending_context,
                         )
                         .await
                     }
@@ -1672,7 +1685,7 @@ impl<T> WalrusNodeClient<T> {
         }
     }
 
-    fn optimistic_upload_max_blob_bytes(&self) -> u64 {
+    fn pending_upload_max_blob_bytes(&self) -> u64 {
         self.config
             .communication_config
             .optimistic_upload_max_blob_bytes
@@ -2028,9 +2041,10 @@ impl<T> WalrusNodeClient<T> {
         blob_persistence_type: &BlobPersistenceType,
         multi_pb: Option<&MultiProgress>,
         store_args: &StoreArgs,
+        pending_context: &PendingUploadContext,
     ) -> ClientResult<ConfirmationCertificate> {
         let blobs = vec![(metadata.clone(), pairs.clone())];
-        let target_nodes = if let Some(failed) = store_args
+        let target_nodes = if let Some(failed) = pending_context
             .initial_failed_nodes
             .as_ref()
             .and_then(|m| m.get(metadata.blob_id()))
@@ -2048,7 +2062,7 @@ impl<T> WalrusNodeClient<T> {
             quorum_forwarder: store_args.quorum_event_tx.clone(),
             tail_handle_collector: store_args.tail_handle_collector.clone(),
             target_nodes,
-            initial_completed_weight: store_args.initial_upload_weight.as_ref(),
+            initial_completed_weight: pending_context.initial_upload_weight.as_ref(),
         };
 
         let (confirmation_results, upload_result) = self
@@ -2340,15 +2354,11 @@ impl<T> WalrusNodeClient<T> {
 
     fn success_weight_by_blob(
         results: &[NodeResult<Vec<BlobId>, StoreError>],
-        excluded_nodes: &HashSet<NodeIndex>,
     ) -> HashMap<BlobId, usize> {
         let mut totals: HashMap<BlobId, usize> = HashMap::new();
         let mut seen: HashMap<BlobId, HashSet<NodeIndex>> = HashMap::new();
 
         for res in results {
-            if excluded_nodes.contains(&res.node) {
-                continue;
-            }
             if let Ok(blob_ids) = &res.result {
                 for blob_id in blob_ids {
                     let counted_nodes = seen.entry(*blob_id).or_default();
diff --git a/crates/walrus-sdk/src/client/client_types.rs b/crates/walrus-sdk/src/client/client_types.rs
@@ -698,19 +698,19 @@ impl WalrusStoreBlobMaybeFinished<BlobWithStatus> {
         Ok(self)
     }
 
-    /// Returns the metadata and sliver pairs needed for an optimistic upload if eligible.
+    /// Returns the metadata and sliver pairs needed for a pending upload if eligible.
     pub fn pending_upload_payload(
         &self,
         pending_uploads_enabled: bool,
         store_args: &StoreArgs,
-        optimistic_upload_max_blob_bytes: u64,
+        pending_upload_max_blob_bytes: u64,
     ) -> Option<(VerifiedBlobMetadataWithId, Arc<Vec<SliverPair>>)> {
         let WalrusStoreBlobState::Unfinished(blob_with_status) = &self.state else {
             return None;
         };
         if blob_with_status.status.is_registered()
             || !pending_uploads_enabled
-            || !store_args.store_optimizations.optimistic_uploads_enabled()
+            || !store_args.store_optimizations.pending_uploads_enabled()
         {
             return None;
         }
@@ -719,7 +719,7 @@ impl WalrusStoreBlobMaybeFinished<BlobWithStatus> {
             return None;
         };
         let unencoded_len = encoded.metadata.metadata().unencoded_length();
-        if unencoded_len > optimistic_upload_max_blob_bytes {
+        if unencoded_len > pending_upload_max_blob_bytes {
             return None;
         }
 
diff --git a/crates/walrus-sdk/src/client/store_args.rs b/crates/walrus-sdk/src/client/store_args.rs
@@ -15,11 +15,7 @@ use tokio::{
 use walrus_core::{DEFAULT_ENCODING, EncodingType, EpochCount};
 use walrus_sui::client::{BlobPersistence, PostStoreAction};
 
-use super::{
-    communication::node::NodeIndex,
-    metrics::ClientMetrics,
-    upload_relay_client::UploadRelayClient,
-};
+use super::{metrics::ClientMetrics, upload_relay_client::UploadRelayClient};
 use crate::{
     client::upload_relay_client::UploadRelayClientError,
     store_optimizations::StoreOptimizations,
@@ -76,11 +72,6 @@ pub struct StoreArgs {
     pub tail_handle_collector: Option<Arc<Mutex<Vec<JoinHandle<()>>>>>,
     /// Optional channel to forward encoding progress events to.
     pub encoding_event_tx: Option<UnboundedSender<EncodingProgressEvent>>,
-    /// Initial upload weight accumulated before certification (e.g., from pending uploads).
-    pub initial_upload_weight: Option<std::collections::HashMap<walrus_core::BlobId, usize>>,
-    /// Nodes that failed during pending uploads, keyed by blob ID, to scope immediate retries.
-    pub initial_failed_nodes:
-        Option<std::collections::HashMap<walrus_core::BlobId, Vec<NodeIndex>>>,
 }
 
 impl StoreArgs {
@@ -120,8 +111,6 @@ impl StoreArgs {
             quorum_event_tx: None,
             tail_handle_collector: None,
             encoding_event_tx: None,
-            initial_upload_weight: None,
-            initial_failed_nodes: None,
         }
     }
 
@@ -159,24 +148,6 @@ impl StoreArgs {
         self
     }
 
-    /// Marks blob IDs that already had optimistic uploads scheduled earlier.
-    pub fn with_initial_upload_weight(
-        mut self,
-        weight: std::collections::HashMap<walrus_core::BlobId, usize>,
-    ) -> Self {
-        self.initial_upload_weight = Some(weight);
-        self
-    }
-
-    /// Sets the initial failed nodes map.
-    pub fn with_initial_failed_nodes(
-        mut self,
-        failed: std::collections::HashMap<walrus_core::BlobId, Vec<NodeIndex>>,
-    ) -> Self {
-        self.initial_failed_nodes = Some(failed);
-        self
-    }
-
     /// Returns a reference to the upload relay client if present.
     pub fn upload_relay_client_ref(&self) -> Option<&UploadRelayClient> {
         self.upload_relay_client
diff --git a/crates/walrus-sdk/src/error.rs b/crates/walrus-sdk/src/error.rs
@@ -52,18 +52,6 @@ pub struct SliverStoreError {
     pub error: NodeError,
 }
 
-impl SliverStoreError {
-    /// Returns true if the underlying node rejected the write because its pending cache is full.
-    pub fn is_cache_saturated(&self) -> bool {
-        self.error.is_cache_saturated()
-    }
-
-    /// Returns true if the underlying node reports the blob is not yet registered.
-    pub fn is_not_registered(&self) -> bool {
-        self.error.is_not_registered()
-    }
-}
-
 /// A helper type for the client to handle errors.
 pub type ClientResult<T> = Result<T, ClientError>;
 
diff --git a/crates/walrus-sdk/src/store_optimizations.rs b/crates/walrus-sdk/src/store_optimizations.rs
@@ -79,8 +79,13 @@ impl StoreOptimizations {
         self.reuse_resources
     }
 
+    /// Returns true if pending (optimistic) uploads are enabled.
+    pub fn pending_uploads_enabled(&self) -> bool {
+        self.optimistic_uploads
+    }
+
     /// Returns true if optimistic buffering is enabled.
     pub fn optimistic_uploads_enabled(&self) -> bool {
-        self.optimistic_uploads
+        self.pending_uploads_enabled()
     }
 }
diff --git a/crates/walrus-storage-node-client/src/error.rs b/crates/walrus-storage-node-client/src/error.rs
@@ -66,16 +66,6 @@ impl NodeError {
             .unwrap_or(false)
     }
 
-    /// Returns true if the error was caused by the pending upload cache being full.
-    pub fn is_cache_saturated(&self) -> bool {
-        self.has_reason("CACHE_SATURATED")
-    }
-
-    /// Returns true if the error indicates the blob is not yet registered.
-    pub fn is_not_registered(&self) -> bool {
-        self.has_reason("NOT_REGISTERED")
-    }
-
     /// Returns true if the HTTP error status code associated with number 499
     pub fn is_expired(&self) -> bool {
         Some(StatusCode::from_u16(499).expect("status code is in a valid range"))
@@ -139,12 +129,6 @@ impl NodeError {
             None
         }
     }
-
-    fn has_reason(&self, reason: &str) -> bool {
-        self.status()
-            .map(|status| status.is_for_reason(reason, STORAGE_NODE_ERROR_DOMAIN))
-            .unwrap_or(false)
-    }
 }
 
 /// Errors returned during the communication with a storage node.

Original file line number	Diff line number	Diff line change
`@@ -79,8 +79,13 @@ impl StoreOptimizations {`
`79`	`79`	`self.reuse_resources`
`80`	`80`	`}`
`81`	`81`
	`82`	`+ /// Returns true if pending (optimistic) uploads are enabled.`
	`83`	`+ pub fn pending_uploads_enabled(&self) -> bool {`
	`84`	`+ self.optimistic_uploads`
	`85`	`+ }`
	`86`	`+`
`82`	`87`	`/// Returns true if optimistic buffering is enabled.`
`83`	`88`	`pub fn optimistic_uploads_enabled(&self) -> bool {`
`84`		`- self.optimistic_uploads`
	`89`	`+ self.pending_uploads_enabled()`
`85`	`90`	`}`
`86`	`91`	`}`