feat: create helpers to set the max_temp_directory_size (apache#15919)

* feat: create helpers to set the max_temp_directory_size

Signed-off-by: Jérémie Drouet <jeremie.drouet@gmail.com>

* refactor: use helper in cli

Signed-off-by: Jérémie Drouet <jeremie.drouet@gmail.com>

* refactor: update error message

Signed-off-by: Jérémie Drouet <jeremie.drouet@gmail.com>

* refactor: use setter in tests

Signed-off-by: Jérémie Drouet <jeremie.drouet@gmail.com>

---------

Signed-off-by: Jérémie Drouet <jeremie.drouet@gmail.com>
Co-authored-by: Andrew Lamb <andrew@nerdnetworks.org>

Author: jdrouet

datafusion-cli/src/main.rs

@@ -177,13 +177,14 @@ async fn main_inner() -> Result<()> {
 
     // set disk limit
     if let Some(disk_limit) = args.disk_limit {
-        let disk_manager = DiskManager::try_new(DiskManagerConfig::NewOs)?;
+        let mut disk_manager = DiskManager::try_new(DiskManagerConfig::NewOs)?;
 
-        let disk_manager = Arc::try_unwrap(disk_manager)
-            .expect("DiskManager should be a single instance")
-            .with_max_temp_directory_size(disk_limit.try_into().unwrap())?;
+        DiskManager::set_arc_max_temp_directory_size(
+            &mut disk_manager,
+            disk_limit.try_into().unwrap(),
+        )?;
 
-        let disk_config = DiskManagerConfig::new_existing(Arc::new(disk_manager));
+        let disk_config = DiskManagerConfig::new_existing(disk_manager);
         rt_builder = rt_builder.with_disk_manager(disk_config);
     }
 

datafusion/core/tests/memory_limit/mod.rs

@@ -534,11 +534,9 @@ async fn setup_context(
     disk_limit: u64,
     memory_pool_limit: usize,
 ) -> Result<SessionContext> {
-    let disk_manager = DiskManager::try_new(DiskManagerConfig::NewOs)?;
+    let mut disk_manager = DiskManager::try_new(DiskManagerConfig::NewOs)?;
 
-    let disk_manager = Arc::try_unwrap(disk_manager)
-        .expect("DiskManager should be a single instance")
-        .with_max_temp_directory_size(disk_limit)?;
+    DiskManager::set_arc_max_temp_directory_size(&mut disk_manager, disk_limit)?;
 
     let runtime = RuntimeEnvBuilder::new()
         .with_memory_pool(Arc::new(FairSpillPool::new(memory_pool_limit)))
@@ -547,7 +545,7 @@ async fn setup_context(
 
     let runtime = Arc::new(RuntimeEnv {
         memory_pool: runtime.memory_pool.clone(),
-        disk_manager: Arc::new(disk_manager),
+        disk_manager,
         cache_manager: runtime.cache_manager.clone(),
         object_store_registry: runtime.object_store_registry.clone(),
     });

datafusion/execution/src/disk_manager.rs

@@ -120,10 +120,10 @@ impl DiskManager {
         }
     }
 
-    pub fn with_max_temp_directory_size(
-        mut self,
+    pub fn set_max_temp_directory_size(
+        &mut self,
         max_temp_directory_size: u64,
-    ) -> Result<Self> {
+    ) -> Result<()> {
         // If the disk manager is disabled and `max_temp_directory_size` is not 0,
         // this operation is not meaningful, fail early.
         if self.local_dirs.lock().is_none() && max_temp_directory_size != 0 {
@@ -133,6 +133,26 @@ impl DiskManager {
         }
 
         self.max_temp_directory_size = max_temp_directory_size;
+        Ok(())
+    }
+
+    pub fn set_arc_max_temp_directory_size(
+        this: &mut Arc<Self>,
+        max_temp_directory_size: u64,
+    ) -> Result<()> {
+        if let Some(inner) = Arc::get_mut(this) {
+            inner.set_max_temp_directory_size(max_temp_directory_size)?;
+            Ok(())
+        } else {
+            config_err!("DiskManager should be a single instance")
+        }
+    }
+
+    pub fn with_max_temp_directory_size(
+        mut self,
+        max_temp_directory_size: u64,
+    ) -> Result<Self> {
+        self.set_max_temp_directory_size(max_temp_directory_size)?;
         Ok(self)
     }