allow configuring message sizes

Author: adriangb

examples/in_memory_cluster.rs

@@ -76,6 +76,13 @@ async fn main() -> Result<(), Box<dyn Error>> {
 
 const DUMMY_URL: &str = "http://localhost:50051";
 
+/// Maximum message size for FlightData chunks in ArrowFlightEndpoint.
+const ENDPOINT_MESSAGE_SIZE: usize = 128 * 1024 * 1024; // 128 MB
+
+/// Maximum message size for gRPC server encoding and decoding.
+/// This should be 2x the ArrowFlightEndpoint max_message_size to allow for overhead.
+const MAX_MESSAGE_SIZE: usize = 256 * 1024 * 1024; // 256 MB
+
 /// [ChannelResolver] implementation that returns gRPC clients baked by an in-memory
 /// tokio duplex rather than a TCP connection.
 #[derive(Clone)]
@@ -113,11 +120,16 @@ impl InMemoryChannelResolver {
                     Ok(builder.build())
                 }
             })
-            .unwrap();
+            .unwrap()
+            .with_max_message_size(ENDPOINT_MESSAGE_SIZE);
 
         tokio::spawn(async move {
             Server::builder()
-                .add_service(FlightServiceServer::new(endpoint))
+                .add_service(
+                    FlightServiceServer::new(endpoint)
+                        .max_decoding_message_size(MAX_MESSAGE_SIZE)
+                        .max_encoding_message_size(MAX_MESSAGE_SIZE),
+                )
                 .serve_with_incoming(tokio_stream::once(Ok::<_, std::io::Error>(server)))
                 .await
         });

examples/localhost_worker.rs

@@ -25,6 +25,13 @@ struct Args {
     cluster_ports: Vec<u16>,
 }
 
+/// Maximum message size for FlightData chunks in ArrowFlightEndpoint.
+const ENDPOINT_MESSAGE_SIZE: usize = 128 * 1024 * 1024; // 128 MB
+
+/// Maximum message size for gRPC server encoding and decoding.
+/// This should be 2x the ArrowFlightEndpoint max_message_size to allow for overhead.
+const MAX_MESSAGE_SIZE: usize = 256 * 1024 * 1024; // 256 MB
+
 #[tokio::main]
 async fn main() -> Result<(), Box<dyn Error>> {
     let args = Args::from_args();
@@ -43,10 +50,15 @@ async fn main() -> Result<(), Box<dyn Error>> {
                 .with_default_features()
                 .build())
         }
-    })?;
+    })?
+    .with_max_message_size(ENDPOINT_MESSAGE_SIZE);
 
     Server::builder()
-        .add_service(FlightServiceServer::new(endpoint))
+        .add_service(
+            FlightServiceServer::new(endpoint)
+                .max_decoding_message_size(MAX_MESSAGE_SIZE)
+                .max_encoding_message_size(MAX_MESSAGE_SIZE),
+        )
         .serve(SocketAddr::new(IpAddr::V4(Ipv4Addr::LOCALHOST), args.port))
         .await?;
 

src/flight_service/do_get.rs

@@ -152,6 +152,7 @@ impl ArrowFlightEndpoint {
             // Note that we do garbage collection of unused dictionary values above, so we are not sending
             // unused dictionary values over the wire.
             .with_dictionary_handling(DictionaryHandling::Resend)
+            .with_max_flight_data_size(usize::MAX)
             .build(stream.map_err(|err| {
                 FlightError::Tonic(Box::new(datafusion_error_to_tonic_status(&err)))
             }));

src/flight_service/service.rs

@@ -28,8 +28,13 @@ pub struct ArrowFlightEndpoint {
     pub(super) task_data_entries: Arc<TTLMap<StageKey, Arc<OnceCell<TaskData>>>>,
     pub(super) session_builder: Arc<dyn DistributedSessionBuilder + Send + Sync>,
     pub(super) hooks: ArrowFlightEndpointHooks,
+    max_message_size: usize,
 }
 
+/// Default maximum message size for FlightData chunks in ArrowFlightEndpoint.
+/// This is the size used for chunking FlightData within the endpoint.
+const DEFAULT_MESSAGE_SIZE: usize = 2 * 1024 * 1024; // 2 MB
+
 impl ArrowFlightEndpoint {
     pub fn try_new(
         session_builder: impl DistributedSessionBuilder + Send + Sync + 'static,
@@ -40,6 +45,7 @@ impl ArrowFlightEndpoint {
             task_data_entries: Arc::new(ttl_map),
             session_builder: Arc::new(session_builder),
             hooks: ArrowFlightEndpointHooks::default(),
+            max_message_size: DEFAULT_MESSAGE_SIZE,
         })
     }
 
@@ -54,6 +60,18 @@ impl ArrowFlightEndpoint {
     ) {
         self.hooks.on_plan.push(Arc::new(hook));
     }
+
+    /// Set the maximum message size for FlightData chunks.
+    /// Defaults to 2 MB.
+    /// If you change this, ensure you configure the server's max_encoding_message_size and
+    /// max_decoding_message_size to at least 2x this value to allow for overhead.
+    /// If your service communication is purely internal and there is no risk of DOS attacks,
+    /// you may want to set this to a considerably larger value to minimize the overhead of chunking
+    /// larger datasets.
+    pub fn with_max_message_size(mut self, size: usize) -> Self {
+        self.max_message_size = size;
+        self
+    }
 }
 
 #[async_trait]

src/test_utils/in_memory_channel_resolver.rs

@@ -12,6 +12,13 @@ use tonic::transport::{Endpoint, Server};
 
 const DUMMY_URL: &str = "http://localhost:50051";
 
+/// Maximum message size for FlightData chunks in ArrowFlightEndpoint.
+const ENDPOINT_MESSAGE_SIZE: usize = 128 * 1024 * 1024; // 128 MB
+
+/// Maximum message size for gRPC server encoding and decoding.
+/// This should be 2x the ArrowFlightEndpoint max_message_size to allow for overhead.
+const MAX_MESSAGE_SIZE: usize = 256 * 1024 * 1024; // 256 MB
+
 /// [ChannelResolver] implementation that returns gRPC clients backed by an in-memory
 /// tokio duplex rather than a TCP connection.
 #[derive(Clone)]
@@ -55,11 +62,16 @@ impl InMemoryChannelResolver {
                     Ok(builder.build())
                 }
             })
-            .unwrap();
+            .unwrap()
+            .with_max_message_size(ENDPOINT_MESSAGE_SIZE);
 
         tokio::spawn(async move {
             Server::builder()
-                .add_service(FlightServiceServer::new(endpoint))
+                .add_service(
+                    FlightServiceServer::new(endpoint)
+                        .max_decoding_message_size(MAX_MESSAGE_SIZE)
+                        .max_encoding_message_size(MAX_MESSAGE_SIZE),
+                )
                 .serve_with_incoming(tokio_stream::once(Ok::<_, std::io::Error>(server)))
                 .await
         });

src/test_utils/localhost.rs

@@ -18,6 +18,13 @@ use tokio::net::TcpListener;
 use tonic::transport::{Channel, Server};
 use url::Url;
 
+/// Maximum message size for FlightData chunks in ArrowFlightEndpoint.
+const ENDPOINT_MESSAGE_SIZE: usize = 128 * 1024 * 1024; // 128 MB
+
+/// Maximum message size for gRPC server encoding and decoding.
+/// This should be 2x the ArrowFlightEndpoint max_message_size to allow for overhead.
+const MAX_MESSAGE_SIZE: usize = 256 * 1024 * 1024; // 256 MB
+
 pub async fn start_localhost_context<B>(
     num_workers: usize,
     session_builder: B,
@@ -113,12 +120,17 @@ pub async fn spawn_flight_service(
     session_builder: impl DistributedSessionBuilder + Send + Sync + 'static,
     incoming: TcpListener,
 ) -> Result<(), Box<dyn Error + Send + Sync>> {
-    let endpoint = ArrowFlightEndpoint::try_new(session_builder)?;
+    let endpoint =
+        ArrowFlightEndpoint::try_new(session_builder)?.with_max_message_size(ENDPOINT_MESSAGE_SIZE);
 
     let incoming = tokio_stream::wrappers::TcpListenerStream::new(incoming);
 
     Ok(Server::builder()
-        .add_service(FlightServiceServer::new(endpoint))
+        .add_service(
+            FlightServiceServer::new(endpoint)
+                .max_decoding_message_size(MAX_MESSAGE_SIZE)
+                .max_encoding_message_size(MAX_MESSAGE_SIZE),
+        )
         .serve_with_incoming(incoming)
         .await?)
 }