Fix sending large messages with NATS client

autonomys · Jul 11, 2024 · a1f9bb0 · a1f9bb0
1 parent 1d36819
commit a1f9bb0
Showing 1 changed file with 17 additions and 3 deletions.
diff --git a/crates/subspace-farmer/src/cluster/nats_client.rs b/crates/subspace-farmer/src/cluster/nats_client.rs
@@ -385,6 +385,7 @@ struct Inner {
     client: Client,
     request_retry_backoff_policy: ExponentialBackoff,
     approximate_max_message_size: usize,
+    max_message_size: usize,
 }
 
 /// NATS client wrapper that can be used to interact with other Subspace-specific clients
@@ -438,6 +439,8 @@ impl NatsClient {
             request_retry_backoff_policy,
             // Allow up to 90%, the rest will be wrapper data structures, etc.
             approximate_max_message_size: max_payload * 9 / 10,
+            // Allow up to 90%, the rest will be wrapper data structures, etc.
+            max_message_size: max_payload,
         };
 
         Ok(Self {
@@ -728,8 +731,9 @@ impl NatsClient {
                 return;
             }
         };
-        let approximate_max_message_size = self.approximate_max_message_size();
-        let max_responses_per_message = approximate_max_message_size / first_element.encoded_size();
+        let max_message_size = self.inner.max_message_size;
+        let max_responses_per_message =
+            self.approximate_max_message_size() / first_element.encoded_size();
 
         let ack_subject = format!("stream-response-ack.{}", Ulid::new());
         let mut ack_subscription = match self.subscribe(ack_subject.clone()).await {
@@ -788,11 +792,21 @@ impl NatsClient {
                     }
                 };
                 let encoded_response = response.encode();
+                let encoded_response_len = encoded_response.len();
                 // When encoded response is too large, remove one of the responses from it and try
                 // again
-                if encoded_response.len() > approximate_max_message_size {
+                if encoded_response_len > max_message_size {
                     buffer = response.into();
                     if let Some(element) = buffer.pop_back() {
+                        if buffer.is_empty() {
+                            error!(
+                                ?element,
+                                encoded_response_len,
+                                max_message_size,
+                                "Element was too large to fit into NATS message, this is an \
+                                implementation bug"
+                            );
+                        }
                         overflow_buffer.push_front(element);
                         continue;
                     } else {