DS4Server/src/app/model_manager.rs

use super::*;

#[derive(Debug)]
pub(super) enum ModelDownload {
    Idle,
    Active(ActiveDownload),
    Complete(ManagedArtifactId, ModelOperation, DownloadProgress),
    Failed(ManagedArtifactId, String, DownloadProgress),
}

#[derive(Clone, Copy, Debug, Eq, PartialEq)]
pub(super) enum ModelOperation {
    Download,
    Validate,
}

#[derive(Debug)]
pub(super) struct ActiveDownload {
    pub(super) artifact: ManagedArtifactId,
    pub(super) operation: ModelOperation,
    pub(super) progress: DownloadProgress,
    pub(super) sampled_at: Instant,
    pub(super) sampled_bytes: u64,
    pub(super) bytes_per_second: f64,
    pub(super) verified_bytes: Arc<AtomicU64>,
    pub(super) cancel: Arc<AtomicBool>,
    pub(super) result: mpsc::Receiver<Result<DownloadOutcome, String>>,
    pub(super) stopping: bool,
}

impl App {
    pub(super) fn open_model_manager(&mut self) -> Task<Message> {
        if let Some(id) = self.model_manager_window {
            return window::gain_focus(id);
        }
        let (id, open) = window::open(window::Settings {
            size: Size::new(760.0, 560.0),
            min_size: Some(Size::new(620.0, 420.0)),
            icon: Some(app_icon()),
            ..Default::default()
        });
        self.model_manager_window = Some(id);
        open.map(Message::ModelManagerOpened)
    }

    pub(super) fn start_model_operation(
        &mut self,
        artifact: ManagedArtifactId,
        operation: ModelOperation,
    ) {
        if matches!(self.model_download, ModelDownload::Active(_)) {
            return;
        }
        let models_path = models_path();
        let progress = match operation {
            ModelOperation::Download => model::artifact_download_progress(artifact, &models_path),
            ModelOperation::Validate => model::artifact_verification_progress(artifact, 0),
        };
        let cancel = Arc::new(AtomicBool::new(false));
        let worker_cancel = Arc::clone(&cancel);
        let verified_bytes = Arc::new(AtomicU64::new(0));
        let worker_verified_bytes = Arc::clone(&verified_bytes);
        let (result_sender, result_receiver) = mpsc::channel();
        let thread_name = match operation {
            ModelOperation::Download => "model-download",
            ModelOperation::Validate => "model-validation",
        };
        if let Err(error) = thread::Builder::new()
            .name(thread_name.to_owned())
            .spawn(move || {
                let result = match operation {
                    ModelOperation::Download => model::download_managed_artifact(
                        artifact,
                        &models_path,
                        &worker_cancel,
                        &worker_verified_bytes,
                    ),
                    ModelOperation::Validate => model::validate_managed_artifact(
                        artifact,
                        &models_path,
                        &worker_cancel,
                        &worker_verified_bytes,
                    ),
                };
                let _ = result_sender.send(result);
            })
        {
            self.error = Some(format!("Could not start {thread_name}: {error}"));
            return;
        }
        self.pending_model_delete = None;
        self.model_download = ModelDownload::Active(ActiveDownload {
            artifact,
            operation,
            sampled_at: Instant::now(),
            sampled_bytes: progress.completed(),
            progress,
            bytes_per_second: 0.0,
            verified_bytes,
            cancel,
            result: result_receiver,
            stopping: false,
        });
    }
}

impl App {
    pub(super) fn update_download_progress(&mut self) {
        let ModelDownload::Active(download) = &mut self.model_download else {
            return;
        };
        let now = Instant::now();
        let verified = download.verified_bytes.load(Ordering::Relaxed);
        let mut progress = match download.operation {
            ModelOperation::Download => {
                model::artifact_download_progress(download.artifact, &models_path())
            }
            ModelOperation::Validate => {
                model::artifact_verification_progress(download.artifact, verified)
            }
        };
        if download.operation == ModelOperation::Download
            && let Some(verification) = &mut progress.verification
        {
            verification.verified = verified.min(verification.total);
        }
        let elapsed = now.duration_since(download.sampled_at).as_secs_f64();
        let completed = progress.completed();
        let phase_changed = progress.phase != download.progress.phase;
        let transferred = completed.saturating_sub(download.sampled_bytes);
        if phase_changed {
            download.bytes_per_second = 0.0;
        } else if transferred > 0 && elapsed > 0.0 {
            let current = transferred as f64 / elapsed;
            download.bytes_per_second = if download.bytes_per_second == 0.0 {
                current
            } else {
                download.bytes_per_second * 0.75 + current * 0.25
            };
        }
        download.progress = progress;
        download.sampled_at = now;
        download.sampled_bytes = completed;

        let result = match download.result.try_recv() {
            Ok(result) => Some(result),
            Err(TryRecvError::Empty) => None,
            Err(TryRecvError::Disconnected) => {
                Some(Err("Download worker stopped unexpectedly.".into()))
            }
        };
        let artifact = download.artifact;
        let operation = download.operation;
        let progress = download.progress.clone();
        if let Some(result) = result {
            match result {
                Ok(DownloadOutcome::Complete) => {
                    let progress = model::artifact_download_progress(artifact, &models_path());
                    self.model_download = ModelDownload::Complete(artifact, operation, progress);
                    self.error = None;
                }
                Ok(DownloadOutcome::Stopped) => {
                    self.model_download = ModelDownload::Idle;
                    self.error = None;
                }
                Err(error) => {
                    self.error = Some(error.clone());
                    self.model_download = ModelDownload::Failed(artifact, error, progress);
                }
            }
        }
    }
}