Make DiskManager max_temp_directory_size dynamically adjustable

Change `max_temp_directory_size` from `u64` to `AtomicU64` so the spill
disk limit can be adjusted at runtime without requiring exclusive (`&mut`)
access to the DiskManager.

Before this change, `set_max_temp_directory_size` required `&mut self`
which was unavailable after the DiskManager was shared via `Arc` (as it
always is in production through RuntimeEnv). The only workaround was
`set_arc_max_temp_directory_size` which required `Arc::get_mut` — always
failing when multiple sessions held references.

After this change, `set_max_temp_directory_size` takes `&self` and uses
an atomic store. Any thread can adjust the limit, and subsequent spill
writes immediately see the new value.

Use cases:
- Adaptive spill limits based on available disk space
- Runtime cluster setting changes without restart
- Graceful degradation under disk pressure

The `set_arc_max_temp_directory_size` method is deprecated but kept for
backward compatibility — it now delegates to the `&self` method.

Performance: `AtomicU64::load(Acquire)` adds ~1ns per spill write check.
Negligible since spill writes take milliseconds.

Signed-off-by: Bukhtawar Khan <bukhtawa@amazon.com>

Author: Bukhtawar

datafusion/execution/src/disk_manager.rs

@@ -75,7 +75,7 @@ impl DiskManagerBuilder {
         match self.mode {
             DiskManagerMode::OsTmpDirectory => Ok(DiskManager {
                 local_dirs: Mutex::new(Some(vec![])),
-                max_temp_directory_size: self.max_temp_directory_size,
+                max_temp_directory_size: AtomicU64::new(self.max_temp_directory_size),
                 used_disk_space: Arc::new(AtomicU64::new(0)),
                 active_files_count: Arc::new(AtomicUsize::new(0)),
             }),
@@ -86,14 +86,14 @@ impl DiskManagerBuilder {
                 );
                 Ok(DiskManager {
                     local_dirs: Mutex::new(Some(local_dirs)),
-                    max_temp_directory_size: self.max_temp_directory_size,
+                    max_temp_directory_size: AtomicU64::new(self.max_temp_directory_size),
                     used_disk_space: Arc::new(AtomicU64::new(0)),
                     active_files_count: Arc::new(AtomicUsize::new(0)),
                 })
             }
             DiskManagerMode::Disabled => Ok(DiskManager {
                 local_dirs: Mutex::new(None),
-                max_temp_directory_size: self.max_temp_directory_size,
+                max_temp_directory_size: AtomicU64::new(self.max_temp_directory_size),
                 used_disk_space: Arc::new(AtomicU64::new(0)),
                 active_files_count: Arc::new(AtomicUsize::new(0)),
             }),
@@ -167,8 +167,9 @@ pub struct DiskManager {
     /// If `None` an error will be returned (configured not to spill)
     local_dirs: Mutex<Option<Vec<Arc<TempDir>>>>,
     /// The maximum amount of data (in bytes) stored inside the temporary directories.
-    /// Default to 100GB
-    max_temp_directory_size: u64,
+    /// Default to 100GB. Stored as `AtomicU64` so it can be adjusted at runtime
+    /// without requiring exclusive (`&mut`) access to the `DiskManager`.
+    max_temp_directory_size: AtomicU64,
     /// Used disk space in the temporary directories. Now only spilled data for
     /// external executors are counted.
     used_disk_space: Arc<AtomicU64>,
@@ -199,7 +200,7 @@ impl DiskManager {
             DiskManagerConfig::Existing(manager) => Ok(manager),
             DiskManagerConfig::NewOs => Ok(Arc::new(Self {
                 local_dirs: Mutex::new(Some(vec![])),
-                max_temp_directory_size: DEFAULT_MAX_TEMP_DIRECTORY_SIZE,
+                max_temp_directory_size: AtomicU64::new(DEFAULT_MAX_TEMP_DIRECTORY_SIZE),
                 used_disk_space: Arc::new(AtomicU64::new(0)),
                 active_files_count: Arc::new(AtomicUsize::new(0)),
             })),
@@ -210,45 +211,48 @@ impl DiskManager {
                 );
                 Ok(Arc::new(Self {
                     local_dirs: Mutex::new(Some(local_dirs)),
-                    max_temp_directory_size: DEFAULT_MAX_TEMP_DIRECTORY_SIZE,
+                    max_temp_directory_size: AtomicU64::new(DEFAULT_MAX_TEMP_DIRECTORY_SIZE),
                     used_disk_space: Arc::new(AtomicU64::new(0)),
                     active_files_count: Arc::new(AtomicUsize::new(0)),
                 }))
             }
             DiskManagerConfig::Disabled => Ok(Arc::new(Self {
                 local_dirs: Mutex::new(None),
-                max_temp_directory_size: DEFAULT_MAX_TEMP_DIRECTORY_SIZE,
+                max_temp_directory_size: AtomicU64::new(DEFAULT_MAX_TEMP_DIRECTORY_SIZE),
                 used_disk_space: Arc::new(AtomicU64::new(0)),
                 active_files_count: Arc::new(AtomicUsize::new(0)),
             })),
         }
     }
 
+    /// Atomically update the max temp directory size. Takes effect immediately
+    /// for subsequent spill writes. Safe to call concurrently from any thread
+    /// without requiring exclusive (`&mut`) access.
     pub fn set_max_temp_directory_size(
-        &mut self,
+        &self,
         max_temp_directory_size: u64,
     ) -> 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 {
             return config_err!(
                 "Cannot set max temp directory size for a disk manager that spilling is disabled"
             );
         }
 
-        self.max_temp_directory_size = max_temp_directory_size;
+        self.max_temp_directory_size.store(max_temp_directory_size, Ordering::Release);
         Ok(())
     }
 
+    #[deprecated(since = "46.0.0", note = "Use set_max_temp_directory_size directly on &self instead")]
     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)?;
+            inner.max_temp_directory_size.store(max_temp_directory_size, Ordering::Release);
             Ok(())
         } else {
-            config_err!("DiskManager should be a single instance")
+            // With AtomicU64 we can now update through shared reference
+            this.set_max_temp_directory_size(max_temp_directory_size)
         }
     }
 
@@ -266,7 +270,7 @@ impl DiskManager {
 
     /// Returns the maximum temporary directory size in bytes
     pub fn max_temp_directory_size(&self) -> u64 {
-        self.max_temp_directory_size
+        self.max_temp_directory_size.load(Ordering::Acquire)
     }
 
     /// Returns the current spilling progress
@@ -418,11 +422,12 @@ impl RefCountedTempFile {
 
         // 3. Check if the updated global disk usage exceeds the configured limit
         let global_disk_usage = self.disk_manager.used_disk_space.load(Ordering::Relaxed);
-        if global_disk_usage > self.disk_manager.max_temp_directory_size {
+        let limit = self.disk_manager.max_temp_directory_size.load(Ordering::Acquire);
+        if global_disk_usage > limit {
             return resources_err!(
                 "The used disk space during the spilling process has exceeded the allowable limit of {}. \
                 Please try increasing the config: `datafusion.runtime.max_temp_directory_size`.",
-                human_readable_size(self.disk_manager.max_temp_directory_size as usize)
+                human_readable_size(limit as usize)
             );
         }
 
@@ -796,4 +801,78 @@ mod tests {
 
         Ok(())
     }
+
+    #[test]
+    fn test_dynamic_limit_adjustment_through_shared_ref() -> Result<()> {
+        // Verify that set_max_temp_directory_size works through &self (not &mut self).
+        // This is the key behavioral change: the limit can be adjusted at runtime
+        // without exclusive access, enabling dynamic resize while queries are running.
+        let dm = DiskManager::builder()
+            .with_max_temp_directory_size(1024)
+            .build()?;
+        let dm = Arc::new(dm);
+
+        assert_eq!(dm.max_temp_directory_size(), 1024);
+
+        // Adjust through shared reference (simulates concurrent access via Arc)
+        dm.set_max_temp_directory_size(2048)?;
+        assert_eq!(dm.max_temp_directory_size(), 2048);
+
+        // Can also decrease
+        dm.set_max_temp_directory_size(512)?;
+        assert_eq!(dm.max_temp_directory_size(), 512);
+
+        Ok(())
+    }
+
+    #[test]
+    fn test_dynamic_limit_concurrent_access() -> Result<()> {
+        // Verify that multiple threads can read and write the limit concurrently
+        let dm = Arc::new(
+            DiskManager::builder()
+                .with_max_temp_directory_size(1000)
+                .build()?,
+        );
+
+        let handles: Vec<_> = (0..8)
+            .map(|i| {
+                let dm = Arc::clone(&dm);
+                std::thread::spawn(move || {
+                    // Each thread sets a different limit and reads it back
+                    let new_limit = (i + 1) * 1000;
+                    dm.set_max_temp_directory_size(new_limit).unwrap();
+                    // Read should return SOME value set by one of the threads
+                    let current = dm.max_temp_directory_size();
+                    assert!(current >= 1000 && current <= 8000);
+                })
+            })
+            .collect();
+
+        for h in handles {
+            h.join().unwrap();
+        }
+
+        // Final value should be one of the values set by threads
+        let final_val = dm.max_temp_directory_size();
+        assert!(final_val >= 1000 && final_val <= 8000);
+
+        Ok(())
+    }
+
+    #[test]
+    fn test_disabled_disk_manager_rejects_nonzero_limit() -> Result<()> {
+        let dm = DiskManager::builder()
+            .with_mode(DiskManagerMode::Disabled)
+            .build()?;
+        let dm = Arc::new(dm);
+
+        // Setting non-zero limit on disabled DiskManager should error
+        let result = dm.set_max_temp_directory_size(1024);
+        assert!(result.is_err());
+
+        // Setting zero is OK
+        assert!(dm.set_max_temp_directory_size(0).is_ok());
+
+        Ok(())
+    }
 }