do_get: use TTL map to store task state

The do_get call now evicts task state from the map after N calls where N
is the number of partitions. This is an approximation because we don't
track that unique partition ids are used, so we might evict early in case
of retries.

The TTLMap will also GC orphaned task state after the configured TTL period.

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);