refactor: InsertFormatted::send_inner() can just send the data first

Author: abonander

src/insert_formatted.rs

@@ -249,15 +249,23 @@ impl InsertFormatted {
     }
 
     async fn send_inner(&mut self, mut data: Bytes) -> Result<()> {
+        if self.state.is_not_started() {
+            self.init_request()?;
+        }
+
         std::future::poll_fn(move |cx| {
             loop {
-                ready!(self.poll_ready(cx))?;
-
                 // Potentially cheaper than cloning `data` which touches the refcount
                 match self.try_send(mem::take(&mut data)) {
-                    ControlFlow::Break(res) => return Poll::Ready(res),
+                    ControlFlow::Break(Ok(())) => return Poll::Ready(Ok(())),
+                    ControlFlow::Break(Err(_)) => {
+                        // If the channel is closed, we should return the actual error
+                        return self.poll_wait_handle(cx);
+                    }
                     ControlFlow::Continue(unsent) => {
                         data = unsent;
+                        // Shorter code-path if we just try to send the data first
+                        ready!(self.poll_ready(cx))?;
                     }
                 }
             }

tests/it/insert_formatted.rs

@@ -48,6 +48,29 @@ async fn insert() {
     verify_insert(&client).await;
 }
 
+#[tokio::test]
+async fn insert_small_chunks() {
+    let client = prepare_database!()
+        // Separate test for compression
+        .with_compression(Compression::None);
+
+    create_table(&client).await;
+
+    let mut bytes = Bytes::copy_from_slice(TAXI_DATA_TSV);
+
+    let mut insert =
+        client.insert_formatted_with("INSERT INTO nyc_taxi_trips_small FORMAT TabSeparated");
+
+    while !bytes.is_empty() {
+        let chunk = bytes.split_to(cmp::min(16, bytes.len()));
+        insert.send(chunk).await.unwrap();
+    }
+
+    insert.end().await.unwrap();
+
+    verify_insert(&client).await;
+}
+
 #[tokio::test]
 #[cfg(feature = "lz4")]
 async fn insert_compressed() {