fix: bound and validate requests to the Comet Native Executors cluster

native/core/src/execution/native_executors/fragment.rs

@@ -22,7 +22,7 @@ use std::sync::Arc;
 use arrow::array::{Array, AsArray, BinaryArray, RecordBatch, RecordBatchOptions, UInt32Array};
 use arrow::datatypes::{DataType, Field, Schema, SchemaRef, UInt32Type};
 use datafusion::common::tree_node::TreeNodeRecursion;
-use datafusion::common::{internal_datafusion_err, DataFusionError, Result};
+use datafusion::common::{internal_datafusion_err, plan_datafusion_err, DataFusionError, Result};
 use datafusion::datasource::memory::MemorySourceConfig;
 use datafusion::execution::disk_manager::DiskManager;
 use datafusion::execution::memory_pool::{MemoryPool, UnboundedMemoryPool};
@@ -52,6 +52,14 @@ use crate::execution::spark_config::{
 
 const DEFAULT_BATCH_SIZE: usize = 8192;
 
+/// The most partitions a stage or an exchange may have. A request names the counts, and the
+/// scheduler allocates for each partition, so a request can't ask for more than this. Spark's
+/// own default is 200 shuffle partitions, and large queries use a few thousand.
+pub const MAX_PARTITIONS: usize = 1 << 16;
+
+/// The largest `spark.comet.batchSize` a request may set. Each operator sizes its buffers by it.
+const MAX_BATCH_SIZE: usize = 1 << 20;
+
 fn descriptor_schema() -> SchemaRef {
     Arc::new(Schema::new(vec![
         Field::new("partition", DataType::UInt32, false),
@@ -71,6 +79,11 @@ pub fn partition_descriptor(
     num_partitions: usize,
     plans: &[Vec<u8>],
 ) -> Result<Arc<dyn ExecutionPlan>> {
+    if num_partitions == 0 || num_partitions > MAX_PARTITIONS {
+        return Err(plan_datafusion_err!(
+            "A stage must have between 1 and {MAX_PARTITIONS} partitions, not {num_partitions}"
+        ));
+    }
     let schema = descriptor_schema();
     let partitions = (0..num_partitions)
         .map(|index| {
@@ -350,14 +363,14 @@ fn fragment_session(
 ) -> Result<Arc<SessionContext>> {
     let mut runtime = comet_runtime_env(memory_pool);
     runtime.disk_manager = disk_manager;
-    let session = comet_session_context(
-        runtime,
-        conf.get_usize(COMET_BATCH_SIZE, DEFAULT_BATCH_SIZE),
-        1,
-        conf,
-        spark_plan,
-    )
-    .map_err(|e| DataFusionError::External(Box::new(e)))?;
+    let batch_size = conf.get_usize(COMET_BATCH_SIZE, DEFAULT_BATCH_SIZE);
+    if batch_size == 0 || batch_size > MAX_BATCH_SIZE {
+        return Err(plan_datafusion_err!(
+            "{COMET_BATCH_SIZE} must be between 1 and {MAX_BATCH_SIZE}, not {batch_size}"
+        ));
+    }
+    let session = comet_session_context(runtime, batch_size, 1, conf, spark_plan)
+        .map_err(|e| DataFusionError::External(Box::new(e)))?;
     Ok(Arc::new(session))
 }
 

native/core/src/execution/native_executors/mod.rs

@@ -26,7 +26,7 @@
 
 mod fragment;
 
-pub use fragment::{block_schema, partition_descriptor, CometFragmentExec};
+pub use fragment::{block_schema, partition_descriptor, CometFragmentExec, MAX_PARTITIONS};
 
 /// The version of this Comet build. A Comet Native Executor cluster requires a driver's to match
 /// its own, together with [`datafusion_comet_proto::PROTO_FINGERPRINT`].

native/native-executors/README.md

@@ -29,6 +29,9 @@ This directory holds the cluster. The Spark driver side is in progress in
 [#7289](https://github.com/apache/datafusion-comet/pull/7289). It is its own Cargo workspace, so the
 main Comet build never depends on Ballista.
 
+**Security.** The scheduler and the executors have no authentication and no TLS, and a query names
+the files and object stores the executors read. Run a cluster only on a trusted network.
+
 ## The scheduler's Flight service
 
 A client runs a query through the scheduler's Arrow Flight service:
@@ -39,12 +42,17 @@ A client runs a query through the scheduler's Arrow Flight service:
    from another Comet build and plans each block with Comet's planner to check it. That builds no
    object store for the block's files, so the scheduler uses no storage credentials. It then
    builds a Ballista physical plan (`src/assembly.rs`), runs it on the Ballista scheduler it
-   wraps, and answers with one endpoint per result partition.
+   wraps, and answers with one endpoint per result partition. It refuses a request with more
+   partitions or a larger batch size than a cluster plausibly needs, a stage whose block reads a
+   different number of inputs than the stage declares, and a block that writes files.
 2. `DoGet` with each endpoint's ticket, at the endpoint's location: the executor that holds the
    partition. A query's results stay there for `COMET_SCHEDULER_RESULT_RETENTION_SECONDS` after
    the query ends (five minutes by default), whether or not a client has read them, and each
    endpoint's expiration time says until when.
 
+The scheduler cancels a job whose client disconnects, or that runs longer than
+`COMET_SCHEDULER_QUERY_TIMEOUT_SECONDS`.
+
 A stage that several other stages read runs once for each of them, since Ballista's planner
 doesn't share a stage between consumers yet.
 
@@ -60,8 +68,11 @@ for each partition. Ballista moves the data between stages:
 
 A stage that fails to plan is refused, not run wrongly or loaded with code it names:
 
-- A stage can't call into the JVM: scalar subqueries and JVM UDFs fail to plan, and a scan whose
-  object store needs a JVM credential provider fails when it runs.
+- A stage can't call into the JVM: JVM UDFs and a scalar subquery the driver has not replaced with
+  its value fail to plan, and a scan whose object store needs a JVM credential provider fails when
+  it runs.
+- A stage can't write files: a Parquet writer, an Iceberg write or a shuffle writer in a block is
+  refused. The cluster moves data between stages itself.
 - A stage can't load a native UDF library. These processes serve every client, and a plan names
   the library path, so a client would choose what the processes load.
 - A stage that reads files can't also read a hash exchange. Spark places the rows of a bucketed
@@ -77,7 +88,13 @@ native/target/debug/comet-executor &
 
 Sharing `native/target` with the main workspace saves building the common dependencies twice. Both
 binaries read their settings from environment variables, listed at the top of
-`src/bin/comet-scheduler.rs` and `src/bin/comet-executor.rs`.
+`src/bin/comet-scheduler.rs` and `src/bin/comet-executor.rs`. The scheduler's include:
+
+| Variable                                   | Meaning                                                              |
+| ------------------------------------------ | -------------------------------------------------------------------- |
+| `COMET_SCHEDULER_RESULT_RETENTION_SECONDS` | How long a query's results stay readable after it ends (default 300) |
+| `COMET_SCHEDULER_FLIGHT_MAX_REQUEST_BYTES` | The largest query the Flight service accepts (default 512 MiB)       |
+| `COMET_SCHEDULER_QUERY_TIMEOUT_SECONDS`    | The longest a query may take (default 0: no limit)                   |
 
 ## Testing
 

native/native-executors/src/assembly.rs

@@ -21,7 +21,7 @@ use std::collections::HashMap;
 use std::sync::Arc;
 
 use comet::execution::native_executors::{
-    block_schema, partition_descriptor, CometFragmentExec, COMET_VERSION,
+    block_schema, partition_descriptor, CometFragmentExec, COMET_VERSION, MAX_PARTITIONS,
 };
 use datafusion::common::{plan_datafusion_err, plan_err, Result};
 use datafusion::physical_expr::expressions::Column;
@@ -32,6 +32,8 @@ use datafusion::physical_plan::ExecutionPlanProperties;
 use datafusion::physical_plan::{ExecutionPlan, Partitioning};
 use datafusion_comet_proto::spark_expression::expr::ExprStruct;
 use datafusion_comet_proto::spark_native_executors::{Fragment, QueryRequest, Stage, StageInput};
+use datafusion_comet_proto::spark_operator::operator::OpStruct;
+use datafusion_comet_proto::spark_operator::Operator;
 use datafusion_comet_proto::spark_partitioning::partitioning::PartitioningStruct;
 use datafusion_comet_proto::PROTO_FINGERPRINT;
 use prost::Message;
@@ -74,11 +76,20 @@ fn build_stage(
     producers: &[Arc<dyn ExecutionPlan>],
     conf: &HashMap<String, String>,
 ) -> Result<Arc<dyn ExecutionPlan>> {
-    let block = stage
+    let plan = stage
         .plan
         .as_ref()
-        .ok_or_else(|| plan_datafusion_err!("The stage has no plan"))?
-        .encode_to_vec();
+        .ok_or_else(|| plan_datafusion_err!("The stage has no plan"))?;
+    // The block's input leaves are fed from the stage's inputs, in order, so the two must agree.
+    check_block(plan, stage.inputs.len())?;
+    for (index, partition_plan) in stage.partition_plans.iter().enumerate() {
+        let partition_plan = Operator::decode(partition_plan.as_slice()).map_err(|e| {
+            plan_datafusion_err!("Cannot decode the plan of partition {index}: {e}")
+        })?;
+        check_block(&partition_plan, stage.inputs.len())
+            .map_err(|e| e.context(format!("Plan of partition {index}")))?;
+    }
+    let block = plan.encode_to_vec();
     // Spark places the rows of a bucketed scan by Murmur3 and shuffles only the other side of a
     // join to match. The hash exchange below would place the other side by DataFusion's hash, so
     // partition `p` of one side would no longer meet the rows of bucket `p` of the other, and the
@@ -122,6 +133,11 @@ fn build_stage(
     if num_partitions == 0 {
         return plan_err!("The stage has no partitions");
     }
+    if num_partitions > MAX_PARTITIONS {
+        return plan_err!(
+            "The stage has {num_partitions} partitions, more than the {MAX_PARTITIONS} allowed"
+        );
+    }
     if !stage.partition_plans.is_empty() && stage.partition_plans.len() != num_partitions {
         return plan_err!(
             "The stage has {num_partitions} partitions but {} partition plans",
@@ -161,6 +177,40 @@ fn build_stage(
     )?))
 }
 
+/// Checks a stage's block, `block`, before it is planned: it must read as many inputs as the
+/// stage declares, and it must not write files. Like a native UDF library or a call into the JVM,
+/// writing is something a driver's stage never does on the cluster, so a request that asks for it
+/// is refused rather than run by processes that serve every client.
+fn check_block(block: &Operator, inputs: usize) -> Result<()> {
+    let mut leaves = 0;
+    let mut pending = vec![block];
+    while let Some(op) = pending.pop() {
+        let writer = match op.op_struct {
+            Some(OpStruct::ParquetWriter(_)) => Some("ParquetWriter"),
+            Some(OpStruct::IcebergWrite(_)) => Some("IcebergWrite"),
+            Some(OpStruct::ShuffleWriter(_)) => Some("ShuffleWriter"),
+            Some(OpStruct::Scan(_)) | Some(OpStruct::ShuffleScan(_)) => {
+                leaves += 1;
+                None
+            }
+            _ => None,
+        };
+        if let Some(writer) = writer {
+            return plan_err!(
+                "A stage can't contain a {writer} operator: Comet Native Executors don't write \
+                 files, and the cluster moves data between stages itself"
+            );
+        }
+        pending.extend(&op.children);
+    }
+    if leaves != inputs {
+        return plan_err!(
+            "The stage's block reads {leaves} inputs but the stage declares {inputs}"
+        );
+    }
+    Ok(())
+}
+
 /// The exchange that carries `producer`'s output to a stage, as a node Ballista's planner turns
 /// into a stage boundary.
 ///
@@ -181,6 +231,16 @@ fn exchange(
         .and_then(|p| p.partitioning_struct.as_ref());
     match partitioning {
         Some(PartitioningStruct::HashPartition(hash)) => {
+            let num_partitions = usize::try_from(hash.num_partitions)
+                .ok()
+                .filter(|count| (1..=MAX_PARTITIONS).contains(count))
+                .ok_or_else(|| {
+                    plan_datafusion_err!(
+                        "A hash exchange must have between 1 and {MAX_PARTITIONS} partitions, \
+                         not {}",
+                        hash.num_partitions
+                    )
+                })?;
             let schema = producer.schema();
             let keys = hash
                 .hash_expression
@@ -198,7 +258,7 @@ fn exchange(
                 .collect::<Result<Vec<_>>>()?;
             Ok(Arc::new(RepartitionExec::try_new(
                 producer,
-                Partitioning::Hash(keys, hash.num_partitions as usize),
+                Partitioning::Hash(keys, num_partitions),
             )?))
         }
         // Ballista can't plan a range shuffle for a physical plan yet

native/native-executors/src/bin/comet-executor.rs

@@ -38,6 +38,10 @@
 //!   memory, or no limit when neither can be found. Memory outside the pools, such as batches in
 //!   flight between stages, comes on top of it.
 //! - `COMET_EXECUTOR_WORK_DIR`: where shuffle and spill files go (default: a temporary directory)
+//!
+//! The executor's Flight service and its gRPC connection with the scheduler don't authenticate
+//! their callers or use TLS, and a query names the files and object stores the executor reads.
+//! Run a cluster only on a trusted network.
 
 use std::sync::Arc;
 

native/native-executors/src/bin/comet-scheduler.rs

@@ -29,9 +29,14 @@
 //! - `COMET_GRPC_MAX_MESSAGE_SIZE`: the largest message, in bytes, of Ballista's gRPC service
 //!   between the scheduler and its executors, such as a task with many partitions of a scan
 //!   (default: Ballista's 16 MiB). Set it on the executors as well, for their side of the same
-//!   messages. The Flight service for drivers has no limit. Shuffle reads between executors keep
-//!   the limit of the query's `ballista.client.grpc_max_message_size` (16 MiB by default), which
-//!   this does not change.
+//!   messages. Shuffle reads between executors keep the limit of the query's
+//!   `ballista.client.grpc_max_message_size` (16 MiB by default), which this does not change.
+//! - `COMET_SCHEDULER_FLIGHT_MAX_REQUEST_BYTES`: the largest query, in bytes, the Flight service
+//!   accepts from a driver. A query carries every stage's plan and the files of every scan
+//!   partition, so a large one is big (default `536870912`, 512 MiB).
+//! - `COMET_SCHEDULER_QUERY_TIMEOUT_SECONDS`: the longest a query may take. The scheduler fails
+//!   and cancels a query that runs longer. A query is also cancelled when the driver that
+//!   submitted it disconnects. (default `0`: no limit)
 //! - `COMET_SCHEDULER_EVENT_LOG_DIR`: a directory to write one event log per query to, as
 //!   Ballista's scheduler does (default: no event logs). The scheduler creates it if needed and
 //!   doesn't start if it can't write there. A log has the query's plan and a record for each
@@ -46,17 +51,22 @@
 //!   results, not to logs.
 //!
 //! Ballista's REST API, which lists and cancels jobs without authentication, is not served.
+//!
+//! Neither the Flight service nor Ballista's gRPC service authenticates its callers or uses TLS,
+//! and a query names the files and object stores executors read. Run a cluster only on a trusted
+//! network.
 
 use std::net::SocketAddr;
 use std::path::Path;
 use std::sync::Arc;
+use std::time::Duration;
 
 use arrow_flight::flight_service_server::FlightServiceServer;
 use ballista_scheduler::cluster::BallistaCluster;
 use ballista_scheduler::config::SchedulerConfig;
 use ballista_scheduler::scheduler_process::{create_scheduler, start_grpc_service};
 use comet_native_executors::codec::CometPhysicalCodec;
-use comet_native_executors::frontend::CometFlightService;
+use comet_native_executors::frontend::{CometFlightService, Limits};
 use comet_native_executors::{env_var, init_logging};
 use datafusion_proto::protobuf::{LogicalPlanNode, PhysicalPlanNode};
 use log::info;
@@ -98,6 +108,12 @@ async fn main() -> Result<(), Box<dyn std::error::Error>> {
         })?;
         config.event_log_dir = Some(dir);
     }
+    let max_request_bytes: usize =
+        env_var("COMET_SCHEDULER_FLIGHT_MAX_REQUEST_BYTES")?.unwrap_or(512 << 20);
+    let mut limits = Limits::default();
+    if let Some(seconds) = env_var::<u64>("COMET_SCHEDULER_QUERY_TIMEOUT_SECONDS")? {
+        limits.query_timeout = (seconds > 0).then(|| Duration::from_secs(seconds));
+    }
     let flight_port: u16 = env_var("COMET_SCHEDULER_FLIGHT_PORT")?.unwrap_or(50060);
     let grpc_addr: SocketAddr = format!("{}:{}", config.bind_host, config.bind_port).parse()?;
     let flight_addr = SocketAddr::new(grpc_addr.ip(), flight_port);
@@ -112,8 +128,8 @@ async fn main() -> Result<(), Box<dyn std::error::Error>> {
     let grpc = tokio::spawn(start_grpc_service(grpc_addr, scheduler.clone()));
     let flight = tonic::transport::Server::builder()
         .add_service(
-            FlightServiceServer::new(CometFlightService::new(scheduler))
-                .max_decoding_message_size(usize::MAX)
+            FlightServiceServer::new(CometFlightService::new(scheduler).with_limits(limits))
+                .max_decoding_message_size(max_request_bytes)
                 .max_encoding_message_size(usize::MAX),
         )
         .serve(flight_addr);