refactor(datafusion-flight-sql-server): remove tokio-stream listener …

datafusion-flight-sql-server/Cargo.toml

@@ -26,7 +26,6 @@ once_cell = "1.21"
 prost.workspace = true
 tonic.workspace = true
 async-trait.workspace = true
-tokio-stream = "0.1.17"
 tokio = { version = "1.47", features = ["net"], default-features = false }
 
 [dev-dependencies]

datafusion-flight-sql-server/src/service.rs

@@ -53,7 +53,11 @@ use log::info;
 use once_cell::sync::Lazy;
 use prost::bytes::Bytes;
 use prost::Message;
-use tonic::transport::Server;
+use tokio::io::{AsyncRead, AsyncWrite};
+use tonic::transport::{
+    server::{Connected, TcpIncoming},
+    Server,
+};
 use tonic::{Request, Response, Status, Streaming};
 
 use super::config::FlightSqlServiceConfig;
@@ -120,12 +124,25 @@ impl FlightSqlService {
     ) -> Result<(), Box<dyn std::error::Error>> {
         info!("Listening on {}", listener.local_addr()?);
 
-        let svc = FlightServiceServer::new(self);
         let listener = tokio::net::TcpListener::from_std(listener)?;
+        let incoming = TcpIncoming::from(listener).with_nodelay(Some(true));
+
+        self.serve_with_incoming(incoming).await
+    }
 
+    pub async fn serve_with_incoming<I, IO, IE>(
+        self,
+        incoming: I,
+    ) -> Result<(), Box<dyn std::error::Error>>
+    where
+        I: Stream<Item = std::result::Result<IO, IE>>,
+        IO: AsyncRead + AsyncWrite + Connected + Unpin + Send + 'static,
+        IE: Into<Box<dyn std::error::Error + Send + Sync>>,
+    {
+        let svc = FlightServiceServer::new(self);
         Ok(Server::builder()
             .add_service(svc)
-            .serve_with_incoming(tokio_stream::wrappers::TcpListenerStream::new(listener))
+            .serve_with_incoming(incoming)
             .await?)
     }