Re: [PATCH datacenter-manager 10/15] task cache: introduce ArchiveFileWriter

[Dominik Csapak <d.csapak@proxmox.com>]

comments inline

On 7/2/26 11:23 AM, Lukas Wagner wrote:
> This is nicely abstracts handling of compressed and uncompressed archive
> files behind a common interface, as well as atomically replacing any
> existing archive file by using the usual tmpfile + rename approach.
> 
> The 'active' file has the same file structure as regular archive files,
> so we can use ArchiveFileWriter for these as well.
> 
> Signed-off-by: Lukas Wagner <l.wagner@proxmox.com>
> ---
>   server/src/remote_tasks/task_cache.rs | 251 +++++++++++++++++---------
>   1 file changed, 163 insertions(+), 88 deletions(-)
> 
> diff --git a/server/src/remote_tasks/task_cache.rs b/server/src/remote_tasks/task_cache.rs
> index c7057786..969a8124 100644
> --- a/server/src/remote_tasks/task_cache.rs
> +++ b/server/src/remote_tasks/task_cache.rs
> @@ -12,6 +12,7 @@ use std::{
>   
>   use anyhow::{Context, Error};
>   use serde::{Deserialize, Serialize};
> +use zstd::Encoder;
>   
>   use proxmox_sys::fs::CreateOptions;
>   
> @@ -665,36 +666,13 @@ impl<'a> WritableTaskCache<'a> {
>       /// The tasks are first written to a temporary file, which is then used
>       /// to atomically replace the original.
>       fn write_active_tasks(&self, tasks: impl Iterator<Item = TaskCacheItem>) -> Result<(), Error> {
> -        let target = self.cache.active_path();
> +        let file = self.cache.active_file();
> +        let mut writer = file.writer(self.cache.create_options)?;
>   
> -        let (fd, path) = proxmox_sys::fs::make_tmp_file(target, self.cache.create_options)?;
> -        let mut fd = BufWriter::new(fd);
> -
> -        Self::write_tasks(&mut fd, tasks)?;
> -
> -        if let Err(err) = fd.flush() {
> -            log::error!("could not flush 'active' file: {err:#}");
> -        }
> -        drop(fd);
> -
> -        let res = std::fs::rename(&path, target).with_context(|| {
> -            format!(
> -                "failed to replace {} with {}",
> -                target.display(),
> -                path.display(),
> -            )
> -        });
> -
> -        if let Err(err) = res {
> -            if let Err(err) = std::fs::remove_file(&path) {
> -                log::error!(
> -                    "failed to cleanup temporary file {}: {err:#}",
> -                    path.display()
> -                );
> -            }
> -
> -            return Err(err);
> -        }
> +        writer.write_tasks(tasks)?;
> +        writer
> +            .finalize_and_replace()
> +            .context("could not finalize 'active' file")?;
>   
>           Ok(())
>       }
> @@ -712,17 +690,7 @@ impl<'a> WritableTaskCache<'a> {
>               return Ok(());
>           }
>   
> -        // TODO: Might be nice to also move this to ArchiveFile
> -        let (temp_file, temp_file_path) =
> -            proxmox_sys::fs::make_tmp_file(&file.path, self.cache.create_options)?;
> -        let mut writer = if file.compressed {
> -            let encoder =
> -                zstd::stream::write::Encoder::new(temp_file, zstd::DEFAULT_COMPRESSION_LEVEL)?
> -                    .auto_finish();
> -            Box::new(BufWriter::new(encoder)) as Box<dyn Write>
> -        } else {
> -            Box::new(BufWriter::new(temp_file)) as Box<dyn Write>
> -        };
> +        let mut writer = file.writer(self.cache.create_options)?;
>   
>           let archive_iter = file
>               .iter()?
> @@ -742,43 +710,8 @@ impl<'a> WritableTaskCache<'a> {
>               .peekable();
>           let task_iter = tasks.into_iter().peekable();
>   
> -        Self::write_tasks(&mut writer, MergeTaskIterator::new(archive_iter, task_iter))?;
> -
> -        if let Err(err) = writer.flush() {
> -            log::error!("could not flush BufWriter for {file:?}: {err:#}");
> -        }
> -        drop(writer);
> -
> -        if let Err(err) = std::fs::rename(&temp_file_path, &file.path).with_context(|| {
> -            format!(
> -                "failed to replace {} with {}",
> -                file.path.display(),
> -                temp_file_path.display()
> -            )
> -        }) {
> -            if let Err(err) = std::fs::remove_file(&temp_file_path) {
> -                log::error!(
> -                    "failed to clean up temporary file {}: {err:#}",
> -                    temp_file_path.display()
> -                );
> -            }
> -
> -            return Err(err);
> -        }
> -
> -        Ok(())
> -    }
> -
> -    /// Write an iterator of [`TaskCacheItem`] to a something that implements [`Write`].
> -    /// The individual items are encoded as JSON followed by a newline.
> -    fn write_tasks(
> -        writer: &mut impl Write,
> -        tasks: impl Iterator<Item = TaskCacheItem>,
> -    ) -> Result<(), Error> {
> -        for task in tasks {
> -            serde_json::to_writer(&mut *writer, &task)?;
> -            writeln!(writer)?;
> -        }
> +        writer.write_tasks(MergeTaskIterator::new(archive_iter, task_iter))?;
> +        writer.finalize_and_replace()?;
>   
>           Ok(())
>       }
> @@ -872,21 +805,12 @@ impl TaskCache {
>               GetTasks::All => {
>                   let mut archive_files = self.archive_files(lock)?;
>                   archive_files.reverse();
> -
> -                archive_files.push(ArchiveFile {
> -                    path: self.active_path().into(),
> -                    compressed: false,
> -                    starttime: 0,
> -                });
> +                archive_files.push(self.active_file());
>   
>                   TaskArchiveIterator::new(Some(journal_file.into()), archive_files, lock)
>               }
>               GetTasks::Active => {
> -                let archive_files = vec![ArchiveFile {
> -                    path: self.active_path().into(),
> -                    compressed: false,
> -                    starttime: 0,
> -                }];
> +                let archive_files = vec![self.active_file()];
>   
>                   TaskArchiveIterator::new(None, archive_files, lock)
>               }
> @@ -977,6 +901,15 @@ impl TaskCache {
>           self.base_path
>               .join(format!("{ARCHIVE_FILENAME_PREFIX}{starttime}{suffix}"))
>       }
> +
> +    /// Return the [`ArchiveFile`] instance for the `active` file.
> +    fn active_file(&self) -> ArchiveFile {
> +        ArchiveFile {
> +            path: self.active_path().into(),
> +            compressed: false,
> +            starttime: 0,
> +        }
> +    }
>   }
>   
>   /// Comparison function for sorting tasks.
> @@ -1174,6 +1107,28 @@ impl ArchiveFile {
>           Ok(Some(iter))
>       }
>   
> +    /// Create an [`ArchiveFileWriter`] for this archive file.
> +    fn writer(&self, create_options: CreateOptions) -> Result<ArchiveFileWriter<'_>, Error> {
> +        let (temp_file, temp_file_path) =
> +            proxmox_sys::fs::make_tmp_file(&self.path, create_options)?;
> +
> +        if self.compressed {
> +            let encoder =
> +                zstd::stream::write::Encoder::new(temp_file, zstd::DEFAULT_COMPRESSION_LEVEL)?;
> +            Ok(ArchiveFileWriter {
> +                archive_file: self,
> +                writer: Some(ArchiveFileWriterInner::ZstdEncoder(encoder)),
> +                path: temp_file_path,
> +            })
> +        } else {
> +            Ok(ArchiveFileWriter {
> +                archive_file: self,
> +                writer: Some(ArchiveFileWriterInner::Plain(BufWriter::new(temp_file))),
> +                path: temp_file_path,
> +            })
> +        }
> +    }
> +
>       fn compress(&mut self, options: CreateOptions) -> Result<(), Error> {
>           let uncompressed_file_path = &self.path;
>   
> @@ -1207,6 +1162,98 @@ impl ArchiveFile {
>       }
>   }
>   
> +/// Writer for an [`ArchiveFile`].
> +///
> +/// Instantiate this via [`ArchiveFile::writer`]. When calling [`Self::finalize_and_replace`], the
> +/// original archive file is replaced with the contents that were written to the writer.
> +struct ArchiveFileWriter<'a> {
> +    archive_file: &'a ArchiveFile,
> +    // writer: Option<Box<dyn Write>>,
> +    writer: Option<ArchiveFileWriterInner<'a>>,
> +    path: PathBuf,
> +}
> +
> +enum ArchiveFileWriterInner<'a> {
> +    ZstdEncoder(Encoder<'a, File>),
> +    Plain(BufWriter<File>),
> +}
> +
> +impl<'a> ArchiveFileWriter<'a> {
> +    /// Write the contents of an iterator of [`TaskCacheItem`] to this archive file.
> +    /// The individual items are encoded as JSON followed by a newline.
> +    fn write_tasks(&mut self, tasks: impl Iterator<Item = TaskCacheItem>) -> Result<(), Error> {
> +        if let Some(writer) = self.writer.as_mut() {
> +            let writer = match writer {
> +                ArchiveFileWriterInner::ZstdEncoder(encoder) => encoder as &mut dyn Write,
> +                ArchiveFileWriterInner::Plain(buf_writer) => buf_writer as &mut dyn Write,
> +            };
> +
> +            for task in tasks {
> +                serde_json::to_writer(&mut *writer, &task)?;
> +                writeln!(writer)?;
> +            }
> +        }
> +
> +        Ok(())
> +    }
> +
> +    /// Call [`ArchiveFileWriter::finalize_and_replace`] or just drop `self`, if you don't need to
> +    /// handle the error.
> +    fn finalize_impl(&mut self) -> Result<(), Error> {
> +        if let Some(writer) = self.writer.take() {
> +            let file = match writer {
> +                ArchiveFileWriterInner::ZstdEncoder(mut encoder) => {
> +                    encoder.flush()?;
> +                    encoder.finish()?
> +                }
> +                ArchiveFileWriterInner::Plain(mut buf_writer) => {
> +                    buf_writer.flush()?;
> +                    buf_writer.into_inner()?
> +                }
> +            };
> +
> +            file.sync_all()
> +                .with_context(|| format!("could not fsync {}", self.path.display()))?;
> +
> +            if let Err(err) =
> +                std::fs::rename(&self.path, &self.archive_file.path).with_context(|| {
> +                    format!(
> +                        "failed to replace {} with {}",
> +                        self.path.display(),
> +                        self.archive_file.path.display()
> +                    )
> +                })
> +            {
> +                if let Err(err) = std::fs::remove_file(&self.path) {
> +                    log::error!(
> +                        "failed to clean up temporary file {path}: {err:#}",
> +                        path = self.path.display()
> +                    );
> +                }
> +
> +                return Err(err);
> +            }
> +        }
> +
> +        Ok(())
> +    }
> +
> +    /// Finalize and drop the internal writer and replace the original
> +    /// archive file with the new contents. This method consumes `self`.
> +    ///
> +    /// [`ArchiveFileWriter::drop`] also finalizes the writer properly, but does not
> +    /// allow to catch errors.
> +    fn finalize_and_replace(mut self) -> Result<(), Error> {
> +        self.finalize_impl()
> +    }
> +}
> +
> +impl<'a> Drop for ArchiveFileWriter<'a> {
> +    fn drop(&mut self) {
> +        let _ = self.finalize_impl();
> +    }
> +}
> +
isn't this a bit dangerous?

this will be called whenever the writer is dropped regardless of error
conditions and will rotate/overwrite the files.

e.g. if any write in between fails, the drop handler will be called and 
overwrite an existing archive with the partially written one?

or am i missing something here?

>   /// Iterator that merges two _sorted_ `Iterator<Item = TaskCacheItem>`, returning the items
>   /// from both iterators sorted.
>   /// The two iterators are expected to be sorted descendingly based on the task's starttime and
> @@ -1715,4 +1762,32 @@ mod tests {
>   
>           assert_eq!(cache.get_tasks(GetTasks::Archived).unwrap().count(), 0);
>       }
> +
> +    #[test]
> +    fn temporary_file_when_writer_lives_is_not_enumerated() {
> +        let (_tmp_dir, mut cache) = make_cache().unwrap();
> +        cache.rotate_after = 100;
> +        cache.uncompressed_files = 1;
> +        cache.max_files = 2;
> +
> +        let cache = cache.write().unwrap();
> +        cache.init(1000).unwrap();
> +
> +        let files = cache.cache.archive_files(&cache.lock).unwrap();
> +
> +        assert_eq!(files.len(), 2);
> +
> +        let _writer = files
> +            .get(0)
> +            .unwrap()
> +            .writer(cache.cache.create_options)
> +            .unwrap();
> +
> +        let files = cache.cache.archive_files(&cache.lock).unwrap();
> +        assert_eq!(
> +            files.len(),
> +            2,
> +            "the tmp file created by the writer should not show up in the list of archive files"
> +        );
> +    }
>   }