v0.34.0_2

src/proxy/body.rs

@@ -301,6 +301,7 @@ pub struct MeteredBody {
     /// For a response, backend selection is final before this wrapper is made.
     /// Keep the selected lease alive and avoid locking Exchange.backend per frame.
     response_backend: Option<Arc<crate::balance::Lease>>,
+    request_backend: Option<Arc<crate::balance::Lease>>,
     direction: Direction,
     timer: Option<Pin<Box<Sleep>>>,
     deadline: tokio::time::Instant,
@@ -317,6 +318,7 @@ impl MeteredBody {
             inner,
             exchange: Some(exchange),
             response_backend,
+            request_backend: None,
             direction,
             // Ready/empty bodies never need a timer allocation. Keep the
             // original deadline so a delayed first poll still times out.
@@ -330,6 +332,7 @@ impl MeteredBody {
         self.finished = true;
         self.timer.take();
         self.response_backend.take();
+        self.request_backend.take();
         self.exchange.take();
     }
 }
@@ -350,7 +353,17 @@ impl HttpBody for MeteredBody {
                 if let (Some(data), Some(exchange)) = (frame.data_ref(), this.exchange.as_ref()) {
                     let count = data.len() as u64;
                     match this.direction {
-                        Direction::Request => exchange.backend_bytes(count, true),
+                        Direction::Request => {
+                            // Backend is selected before a non-empty upstream
+                            // upload is polled. No retry of a streamed write.
+                            if this.request_backend.is_none() {
+                                this.request_backend = exchange.backend
+                                    .lock().unwrap_or_else(|e| e.into_inner()).clone();
+                            }
+                            if let Some(lease) = &this.request_backend {
+                                lease.bytes(count, true);
+                            }
+                        }
                         Direction::Response => {
                             if let Some(lease) = &this.response_backend {
                                 lease.bytes(count, false);