hook everything up

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

src/flight_service/do_get.rs

@@ -14,6 +14,7 @@ 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;
@@ -38,6 +39,13 @@ pub struct DoGet {
     pub stage_key: Option<StageKey>,
 }
 
+#[derive(Clone)]
+pub struct TaskData {
+    pub(super) state: SessionState,
+    pub(super) stage: Arc<ExecutionStage>,
+    partition_count: Arc<AtomicUsize>,
+}
+
 impl ArrowFlightEndpoint {
     pub(super) async fn get(
         &self,
@@ -51,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
@@ -91,15 +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 = self.stages.get_or_init(key, || {
-            OnceCell::<(SessionState, Arc<ExecutionStage>)>::new()
-        });
+        let once_stage = self
+            .stages
+            .get_or_init(key.clone(), || OnceCell::<TaskData>::new());
 
-        let stage_data = once_stage
+        let mut stage_data = once_stage
             .get_or_try_init(|| async {
                 let stage_proto = doget
                     .stage_proto
@@ -134,10 +145,19 @@ impl ArrowFlightEndpoint {
                 config.set_extension(stage.clone());
                 config.set_extension(Arc::new(ContextGrpcMetadata(headers)));
 
-                Ok::<_, Status>((state, stage))
+                Ok::<_, Status>(TaskData {
+                    state,
+                    stage,
+                    partition_count: Arc::new(AtomicUsize::new(0)),
+                })
             })
             .await?;
 
+        stage_data.partition_count.fetch_sub(1, Ordering::SeqCst);
+        if stage_data.partition_count.load(Ordering::SeqCst) <= 0 {
+            self.stages.remove(key.clone());
+        }
+
         Ok(stage_data.clone())
     }
 }

src/flight_service/service.rs

@@ -1,8 +1,8 @@
 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::{
@@ -12,7 +12,6 @@ use arrow_flight::{
 use async_trait::async_trait;
 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;
@@ -36,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<(SessionState, Arc<ExecutionStage>)>>,
+    pub(super) stages: TTLMap<StageKey, OnceCell<TaskData>>,
     pub(super) session_builder: Arc<dyn DistributedSessionBuilder + Send + Sync>,
 }