Updates the global disk usage counter after modifications to the underlying file. # Errors - Returns an error if the global disk usage exceeds the configured limit.
(&mut self)
| 399 | /// # Errors |
| 400 | /// - Returns an error if the global disk usage exceeds the configured limit. |
| 401 | pub fn update_disk_usage(&mut self) -> Result<()> { |
| 402 | // Get new file size from OS |
| 403 | let metadata = self.tempfile.as_file().metadata()?; |
| 404 | let new_disk_usage = metadata.len(); |
| 405 | |
| 406 | // Get the old disk usage |
| 407 | let old_disk_usage = self.current_file_disk_usage.load(Ordering::Relaxed); |
| 408 | |
| 409 | // Update the global disk usage by: |
| 410 | // 1. Subtracting the old file size from the global counter |
| 411 | self.disk_manager |
| 412 | .used_disk_space |
| 413 | .fetch_sub(old_disk_usage, Ordering::Relaxed); |
| 414 | // 2. Adding the new file size to the global counter |
| 415 | self.disk_manager |
| 416 | .used_disk_space |
| 417 | .fetch_add(new_disk_usage, Ordering::Relaxed); |
| 418 | |
| 419 | // 3. Check if the updated global disk usage exceeds the configured limit |
| 420 | let global_disk_usage = self.disk_manager.used_disk_space.load(Ordering::Relaxed); |
| 421 | if global_disk_usage > self.disk_manager.max_temp_directory_size { |
| 422 | return resources_err!( |
| 423 | "The used disk space during the spilling process has exceeded the allowable limit of {}. \ |
| 424 | Please try increasing the config: `datafusion.runtime.max_temp_directory_size`.", |
| 425 | human_readable_size(self.disk_manager.max_temp_directory_size as usize) |
| 426 | ); |
| 427 | } |
| 428 | |
| 429 | // 4. Update the local file size tracking |
| 430 | self.current_file_disk_usage |
| 431 | .store(new_disk_usage, Ordering::Relaxed); |
| 432 | |
| 433 | Ok(()) |
| 434 | } |
| 435 | |
| 436 | pub fn current_disk_usage(&self) -> u64 { |
| 437 | self.current_file_disk_usage.load(Ordering::Relaxed) |