do_get: use TTLMap in ArrowFlightEndpoint instead of DashMap

Author: jayshrivastava

src/common/ttl_map.rs

@@ -56,7 +56,7 @@ where
     /// new creates a new Bucket
     fn new<V>(data: Arc<DashMap<K, V>>) -> Self
     where
-        V: Default + Clone + Send + Sync + 'static,
+        V: Send + Sync + 'static,
     {
         // TODO: To avoid unbounded growth, consider making this bounded. Alternatively, we can
         // introduce a dynamic GC period to ensure that GC can keep up.
@@ -84,7 +84,7 @@ where
     /// task is responsible for managing a subset of keys in the map.
     async fn task<V>(mut rx: UnboundedReceiver<BucketOp<K>>, data: Arc<DashMap<K, V>>)
     where
-        V: Default + Clone + Send + Sync + 'static,
+        V: Send + Sync + 'static,
     {
         let mut shard = HashSet::new();
         while let Some(op) = rx.recv().await {
@@ -207,7 +207,7 @@ where
 
     /// get_or_default executes the provided closure with a reference to the map entry for the given key.
     /// If the key does not exist, it inserts a new entry with the default value.
-    pub fn get_or_init<F>(&self, key: K, f: F) -> V
+    pub fn get_or_init<F>(&self, key: K, init: F) -> V
     where
         F: FnOnce() -> V,
     {
@@ -218,7 +218,7 @@ where
 
         let value = match self.data.entry(key.clone()) {
             Entry::Vacant(entry) => {
-                let value = f();
+                let value = init();
                 entry.insert(value.clone());
                 new_entry = true;
                 value
@@ -257,6 +257,18 @@ where
         self.data.remove(&key);
     }
 
+    /// Returns the number of entries currently stored in the map
+    #[cfg(test)]
+    pub fn len(&self) -> usize {
+        self.data.len()
+    }
+
+    /// Returns an iterator over the keys currently stored in the map
+    #[cfg(test)]
+    pub fn keys(&self) -> impl Iterator<Item = K> + '_ {
+        self.data.iter().map(|entry| entry.key().clone())
+    }
+
     /// run_gc_loop will continuously clear expired entries from the map, checking every `period`. The
     /// function terminates if `shutdown` is signalled.
     async fn run_gc_loop(time: Arc<AtomicU64>, period: Duration, buckets: &[Bucket<K>]) {

src/flight_service/do_get.rs

@@ -14,7 +14,9 @@ use arrow_flight::Ticket;
 use datafusion::execution::SessionState;
 use futures::TryStreamExt;
 use prost::Message;
+use std::sync::atomic::{AtomicUsize, Ordering};
 use std::sync::Arc;
+use tokio::sync::OnceCell;
 use tonic::metadata::MetadataMap;
 use tonic::{Request, Response, Status};
 
@@ -37,6 +39,13 @@ pub struct DoGet {
     pub stage_key: Option<StageKey>,
 }
 
+#[derive(Clone, Debug)]
+pub struct TaskData {
+    pub(super) state: SessionState,
+    pub(super) stage: Arc<ExecutionStage>,
+    partition_count: Arc<AtomicUsize>,
+}
+
 impl ArrowFlightEndpoint {
     pub(super) async fn get(
         &self,
@@ -50,7 +59,10 @@ impl ArrowFlightEndpoint {
 
         let partition = doget.partition as usize;
         let task_number = doget.task_number as usize;
-        let (mut state, stage) = self.get_state_and_stage(doget, metadata).await?;
+        let task_data = self.get_state_and_stage(doget, metadata).await?;
+
+        let stage = task_data.stage;
+        let mut state = task_data.state;
 
         // find out which partition group we are executing
         let task = stage
@@ -90,16 +102,15 @@ impl ArrowFlightEndpoint {
         &self,
         doget: DoGet,
         metadata_map: MetadataMap,
-    ) -> Result<(SessionState, Arc<ExecutionStage>), Status> {
+    ) -> Result<TaskData, Status> {
         let key = doget
             .stage_key
             .ok_or(Status::invalid_argument("DoGet is missing the stage key"))?;
-        let once_stage = {
-            let entry = self.stages.entry(key).or_default();
-            Arc::clone(&entry)
-        };
+        let once_stage = self
+            .stages
+            .get_or_init(key.clone(), || Arc::new(OnceCell::<TaskData>::new()));
 
-        let (state, stage) = once_stage
+        let stage_data = once_stage
             .get_or_try_init(|| async {
                 let stage_proto = doget
                     .stage_proto
@@ -134,10 +145,203 @@ impl ArrowFlightEndpoint {
                 config.set_extension(stage.clone());
                 config.set_extension(Arc::new(ContextGrpcMetadata(headers)));
 
-                Ok::<_, Status>((state, stage))
+                // Initialize partition count to the number of partitions in the stage
+                let total_partitions = stage.plan.properties().partitioning.partition_count();
+                Ok::<_, Status>(TaskData {
+                    state,
+                    stage,
+                    partition_count: Arc::new(AtomicUsize::new(total_partitions)),
+                })
             })
             .await?;
 
-        Ok((state.clone(), stage.clone()))
+        let remaining_partitions = stage_data.partition_count.fetch_sub(1, Ordering::SeqCst);
+        if remaining_partitions <= 1 {
+            self.stages.remove(key.clone());
+        }
+
+        Ok(stage_data.clone())
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use arrow_flight::Ticket;
+    use prost::{bytes::Bytes, Message};
+    use uuid::Uuid;
+
+    #[tokio::test]
+    async fn test_task_data_partition_counting() {
+        use crate::task::ExecutionTask;
+        use arrow_flight::Ticket;
+        use prost::{bytes::Bytes, Message};
+        use tonic::Request;
+        use url::Url;
+
+        // Create a mock channel resolver for ArrowFlightEndpoint
+        #[derive(Clone)]
+        struct MockChannelResolver;
+
+        #[async_trait::async_trait]
+        impl crate::ChannelResolver for MockChannelResolver {
+            fn get_urls(&self) -> Result<Vec<Url>, datafusion::error::DataFusionError> {
+                Ok(vec![])
+            }
+
+            async fn get_channel_for_url(
+                &self,
+                _url: &Url,
+            ) -> Result<crate::BoxCloneSyncChannel, datafusion::error::DataFusionError>
+            {
+                Err(datafusion::error::DataFusionError::NotImplemented(
+                    "Mock resolver".to_string(),
+                ))
+            }
+        }
+
+        // Create ArrowFlightEndpoint
+        let endpoint =
+            ArrowFlightEndpoint::new(MockChannelResolver).expect("Failed to create endpoint");
+
+        // Test parameters: 2 stages, each with 3 tasks, each task has 10 partitions
+        let num_tasks = 3;
+        let num_partitions_per_task = 3;
+        let stage_id = 1;
+        let query_id_uuid = Uuid::new_v4();
+        let query_id = query_id_uuid.as_bytes().to_vec();
+
+        // Create ExecutionStage proto with multiple tasks
+        let mut tasks = Vec::new();
+        for i in 0..num_tasks {
+            tasks.push(ExecutionTask {
+                url_str: None,
+                partition_group: vec![i], // Different partition group for each task
+            });
+        }
+
+        let stage_proto = ExecutionStageProto {
+            query_id: query_id.clone(),
+            num: 1,
+            name: format!("test_stage_{}", 1),
+            plan: Some(Box::new(create_mock_physical_plan_proto(
+                num_partitions_per_task,
+            ))),
+            inputs: vec![],
+            tasks,
+        };
+
+        let task_keys = vec![
+            StageKey {
+                query_id: query_id_uuid.to_string(),
+                stage_id,
+                task_number: 0,
+            },
+            StageKey {
+                query_id: query_id_uuid.to_string(),
+                stage_id,
+                task_number: 1,
+            },
+            StageKey {
+                query_id: query_id_uuid.to_string(),
+                stage_id,
+                task_number: 2,
+            },
+        ];
+
+        let stage_proto_for_closure = stage_proto.clone();
+        let endpoint_ref = &endpoint;
+        let do_get = async move |partition: u64, task_number: u64, stage_key: StageKey| {
+            let stage_proto = stage_proto_for_closure.clone();
+            // Create DoGet message
+            let doget = DoGet {
+                stage_proto: Some(stage_proto),
+                task_number,
+                partition,
+                stage_key: Some(stage_key),
+            };
+
+            // Create Flight ticket
+            let ticket = Ticket {
+                ticket: Bytes::from(doget.encode_to_vec()),
+            };
+
+            // Call the actual get() method
+            let request = Request::new(ticket);
+            endpoint_ref.get(request).await
+        };
+
+        // For each task, call do_get() for each partition except the last.
+        for task_number in 0..num_tasks {
+            for partition in 0..num_partitions_per_task - 1 {
+                let result = do_get(
+                    partition as u64,
+                    task_number,
+                    task_keys[task_number as usize].clone(),
+                )
+                .await;
+                assert!(result.is_ok());
+            }
+        }
+
+        // Check that the endpoint has not evicted any tasks.
+        assert_eq!(endpoint.stages.len(), num_tasks as usize);
+
+        // Run the last partition of task 0. Any partition number works. Verify that the task state
+        // is evicted because all partitions have been processed.
+        let result = do_get(1, 0, task_keys[0].clone()).await;
+        assert!(result.is_ok());
+        let stored_stage_keys = endpoint.stages.keys().collect::<Vec<StageKey>>();
+        assert_eq!(stored_stage_keys.len(), 2);
+        assert!(stored_stage_keys.contains(&task_keys[1]));
+        assert!(stored_stage_keys.contains(&task_keys[2]));
+
+        // Run the last partition of task 1.
+        let result = do_get(1, 1, task_keys[1].clone()).await;
+        assert!(result.is_ok());
+        let stored_stage_keys = endpoint.stages.keys().collect::<Vec<StageKey>>();
+        assert_eq!(stored_stage_keys.len(), 1);
+        assert!(stored_stage_keys.contains(&task_keys[2]));
+
+        // Run the last partition of the last task.
+        let result = do_get(1, 2, task_keys[2].clone()).await;
+        assert!(result.is_ok());
+        let stored_stage_keys = endpoint.stages.keys().collect::<Vec<StageKey>>();
+        assert_eq!(stored_stage_keys.len(), 0);
+    }
+
+    // Helper to create a mock physical plan proto
+    fn create_mock_physical_plan_proto(
+        partitions: usize,
+    ) -> datafusion_proto::protobuf::PhysicalPlanNode {
+        use datafusion_proto::protobuf::partitioning::PartitionMethod;
+        use datafusion_proto::protobuf::{
+            Partitioning, PhysicalPlanNode, RepartitionExecNode, Schema,
+        };
+
+        // Create a repartition node that will have the desired partition count
+        PhysicalPlanNode {
+            physical_plan_type: Some(
+                datafusion_proto::protobuf::physical_plan_node::PhysicalPlanType::Repartition(
+                    Box::new(RepartitionExecNode {
+                        input: Some(Box::new(PhysicalPlanNode {
+                            physical_plan_type: Some(
+                                datafusion_proto::protobuf::physical_plan_node::PhysicalPlanType::Empty(
+                                    datafusion_proto::protobuf::EmptyExecNode {
+                                        schema: Some(Schema {
+                                            columns: vec![],
+                                            metadata: std::collections::HashMap::new(),
+                                        })
+                                    }
+                                )
+                            ),
+                        })),
+                        partitioning: Some(Partitioning {
+                            partition_method: Some(PartitionMethod::RoundRobin(partitions as u64)),
+                        }),
+                    })
+                )
+            ),
+        }
     }
 }

src/flight_service/service.rs

@@ -1,17 +1,17 @@
 use crate::channel_manager::ChannelManager;
+use crate::common::ttl_map::{TTLMap, TTLMapConfig};
+use crate::flight_service::do_get::TaskData;
 use crate::flight_service::session_builder::DefaultSessionBuilder;
 use crate::flight_service::DistributedSessionBuilder;
-use crate::stage::ExecutionStage;
 use crate::ChannelResolver;
 use arrow_flight::flight_service_server::FlightService;
 use arrow_flight::{
     Action, ActionType, Criteria, Empty, FlightData, FlightDescriptor, FlightInfo,
     HandshakeRequest, HandshakeResponse, PollInfo, PutResult, SchemaResult, Ticket,
 };
 use async_trait::async_trait;
-use dashmap::DashMap;
+use datafusion::error::DataFusionError;
 use datafusion::execution::runtime_env::RuntimeEnv;
-use datafusion::execution::SessionState;
 use futures::stream::BoxStream;
 use std::sync::Arc;
 use tokio::sync::OnceCell;
@@ -35,18 +35,21 @@ pub struct ArrowFlightEndpoint {
     pub(super) channel_manager: Arc<ChannelManager>,
     pub(super) runtime: Arc<RuntimeEnv>,
     #[allow(clippy::type_complexity)]
-    pub(super) stages: DashMap<StageKey, Arc<OnceCell<(SessionState, Arc<ExecutionStage>)>>>,
+    pub(super) stages: TTLMap<StageKey, Arc<OnceCell<TaskData>>>,
     pub(super) session_builder: Arc<dyn DistributedSessionBuilder + Send + Sync>,
 }
 
 impl ArrowFlightEndpoint {
-    pub fn new(channel_resolver: impl ChannelResolver + Send + Sync + 'static) -> Self {
-        Self {
+    pub fn new(
+        channel_resolver: impl ChannelResolver + Send + Sync + 'static,
+    ) -> Result<Self, DataFusionError> {
+        let ttl_map = TTLMap::try_new(TTLMapConfig::default())?;
+        Ok(Self {
             channel_manager: Arc::new(ChannelManager::new(channel_resolver)),
             runtime: Arc::new(RuntimeEnv::default()),
-            stages: DashMap::new(),
+            stages: ttl_map,
             session_builder: Arc::new(DefaultSessionBuilder),
-        }
+        })
     }
 
     pub fn with_session_builder(

src/test_utils/localhost.rs

@@ -106,7 +106,7 @@ pub async fn spawn_flight_service(
     session_builder: impl DistributedSessionBuilder + Send + Sync + 'static,
     incoming: TcpListener,
 ) -> Result<(), Box<dyn Error + Send + Sync>> {
-    let mut endpoint = ArrowFlightEndpoint::new(channel_resolver);
+    let mut endpoint = ArrowFlightEndpoint::new(channel_resolver)?;
     endpoint.with_session_builder(session_builder);
 
     let incoming = tokio_stream::wrappers::TcpListenerStream::new(incoming);