do_get: use TTLMap in ArrowFlightEndpoint instead of DashMap

Author: jayshrivastava

src/flight_service/do_get.rs

@@ -15,6 +15,7 @@ use datafusion::execution::SessionState;
 use futures::TryStreamExt;
 use prost::Message;
 use std::sync::Arc;
+use tokio::sync::OnceCell;
 use tonic::metadata::MetadataMap;
 use tonic::{Request, Response, Status};
 
@@ -94,12 +95,11 @@ impl ArrowFlightEndpoint {
         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, || {
+            OnceCell::<(SessionState, Arc<ExecutionStage>)>::new()
+        });
 
-        let (state, stage) = once_stage
+        let stage_data = once_stage
             .get_or_try_init(|| async {
                 let stage_proto = doget
                     .stage_proto
@@ -138,6 +138,6 @@ impl ArrowFlightEndpoint {
             })
             .await?;
 
-        Ok((state.clone(), stage.clone()))
+        Ok(stage_data.clone())
     }
 }

src/flight_service/service.rs

@@ -1,4 +1,5 @@
 use crate::channel_manager::ChannelManager;
+use crate::common::ttl_map::{TTLMap, TTLMapConfig};
 use crate::flight_service::session_builder::DefaultSessionBuilder;
 use crate::flight_service::DistributedSessionBuilder;
 use crate::stage::ExecutionStage;
@@ -9,7 +10,7 @@ use arrow_flight::{
     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;
@@ -35,18 +36,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, OnceCell<(SessionState, Arc<ExecutionStage>)>>,
     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);