add test

Author: jayshrivastava

src/common/ttl_map.rs

@@ -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

@@ -39,7 +39,7 @@ pub struct DoGet {
     pub stage_key: Option<StageKey>,
 }
 
-#[derive(Clone)]
+#[derive(Clone, Debug)]
 pub struct TaskData {
     pub(super) state: SessionState,
     pub(super) stage: Arc<ExecutionStage>,
@@ -108,9 +108,9 @@ impl ArrowFlightEndpoint {
             .ok_or(Status::invalid_argument("DoGet is missing the stage key"))?;
         let once_stage = self
             .stages
-            .get_or_init(key.clone(), || OnceCell::<TaskData>::new());
+            .get_or_init(key.clone(), || Arc::new(OnceCell::<TaskData>::new()));
 
-        let mut stage_data = once_stage
+        let stage_data = once_stage
             .get_or_try_init(|| async {
                 let stage_proto = doget
                     .stage_proto
@@ -145,19 +145,203 @@ impl ArrowFlightEndpoint {
                 config.set_extension(stage.clone());
                 config.set_extension(Arc::new(ContextGrpcMetadata(headers)));
 
+                // 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(0)),
+                    partition_count: Arc::new(AtomicUsize::new(total_partitions)),
                 })
             })
             .await?;
 
-        stage_data.partition_count.fetch_sub(1, Ordering::SeqCst);
-        if stage_data.partition_count.load(Ordering::SeqCst) <= 0 {
+        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

@@ -35,7 +35,7 @@ pub struct ArrowFlightEndpoint {
     pub(super) channel_manager: Arc<ChannelManager>,
     pub(super) runtime: Arc<RuntimeEnv>,
     #[allow(clippy::type_complexity)]
-    pub(super) stages: TTLMap<StageKey, OnceCell<TaskData>>,
+    pub(super) stages: TTLMap<StageKey, Arc<OnceCell<TaskData>>>,
     pub(super) session_builder: Arc<dyn DistributedSessionBuilder + Send + Sync>,
 }