g3-daemon: use vectored IO in register

lib/g3-daemon/src/register/task.rs

@@ -14,12 +14,13 @@
  * limitations under the License.
  */
 
+use std::io::IoSlice;
 use std::sync::Arc;
 
 use anyhow::anyhow;
 use http::{Method, StatusCode};
 use serde_json::{Map, Value};
-use tokio::io::BufStream;
+use tokio::io::{AsyncWriteExt, BufStream};
 use tokio::net::TcpStream;
 
 use g3_http::HttpBodyReader;
@@ -58,24 +59,25 @@ impl RegisterTask {
         }
 
         let body = Value::Object(content).to_string();
-        let data = format!(
+        let header = format!(
             "POST {} HTTP/1.1\r\n\
              Host: {}\r\n\
              Content-Type: application/json\r\n\
              Content-Length: {}\r\n\
              Connection: Keep-Alive\r\n\
-             \r\n{body}",
+             \r\n",
             self.config.register_path,
             self.config.upstream.host(),
             body.len()
         );
 
-        self.write_request(data.as_bytes()).await?;
+        self.write_request(header.as_bytes(), Some(body.as_bytes()))
+            .await?;
         self.check_response(Method::POST).await
     }
 
     pub async fn ping_until_end(&mut self) -> anyhow::Result<()> {
-        let data = format!(
+        let header = format!(
             "GET {} HTTP/1.1\r\n\
              Host: {}\r\n\
              Content-Length: 0\r\n\
@@ -89,7 +91,7 @@ impl RegisterTask {
         loop {
             tokio::select! {
                 _ = interval.tick() => {
-                    self.write_request(data.as_bytes()).await?;
+                    self.write_request(header.as_bytes(), None).await?;
                     self.check_response(Method::GET).await?;
                 }
                 _ = self.stream.fill_wait_data() => {
@@ -99,11 +101,23 @@ impl RegisterTask {
         }
     }
 
-    async fn write_request(&mut self, data: &[u8]) -> anyhow::Result<()> {
-        self.stream
-            .write_all_flush(data)
-            .await
-            .map_err(|e| anyhow!("failed to write data: {e:?}"))
+    async fn write_request(&mut self, header: &[u8], body: Option<&[u8]>) -> anyhow::Result<()> {
+        if let Some(body) = body {
+            let slices = [IoSlice::new(header), IoSlice::new(body)];
+            self.stream
+                .write_all_vectored(slices)
+                .await
+                .map_err(|e| anyhow!("failed to write request: {e:?}"))?;
+            self.stream
+                .flush()
+                .await
+                .map_err(|e| anyhow!("failed to write request: {e:?}"))
+        } else {
+            self.stream
+                .write_all_flush(header)
+                .await
+                .map_err(|e| anyhow!("failed to write request: {e:?}"))
+        }
     }
 
     async fn check_response(&mut self, method: Method) -> anyhow::Result<()> {