feat: background downloading

This commit is contained in:
Andreas Bigger
2025-12-29 17:05:07 -05:00
parent 7129747d5e
commit 835bddd2ee
23 changed files with 3683 additions and 294 deletions
+25
View File
@@ -0,0 +1,25 @@
[package]
name = "paracas-daemon"
version.workspace = true
edition.workspace = true
rust-version.workspace = true
license.workspace = true
homepage.workspace = true
repository.workspace = true
description = "Background job management for paracas tick data downloader"
[lints]
workspace = true
[dependencies]
serde = { workspace = true }
serde_json = { workspace = true }
chrono = { workspace = true }
thiserror = { workspace = true }
uuid = { workspace = true }
directories = { workspace = true }
tokio = { workspace = true }
[dev-dependencies]
tempfile = { workspace = true }
tokio = { workspace = true, features = ["rt-multi-thread", "macros"] }
+273
View File
@@ -0,0 +1,273 @@
//! Daemon process spawning for background downloads.
//!
//! This module provides functionality to spawn detached daemon processes
//! that can run downloads in the background, even after the parent process exits.
use crate::{DownloadJob, JobId, StateError, StateManager};
use std::fs::OpenOptions;
use std::path::PathBuf;
use std::process::{Command, Stdio};
/// Result type for daemon operations (re-exported from state module).
pub(crate) type Result<T> = std::result::Result<T, StateError>;
/// Environment variable name for the daemon job ID.
pub const DAEMON_JOB_ID_ENV: &str = "PARACAS_DAEMON_JOB_ID";
/// Command line argument for daemon mode.
pub const DAEMON_RUN_ARG: &str = "--daemon-run";
/// Spawns detached daemon processes for background downloads.
///
/// The spawner handles all the platform-specific details of creating
/// a detached background process that will continue running after
/// the parent process exits.
#[derive(Debug, Clone)]
pub struct DaemonSpawner {
state_manager: StateManager,
executable_path: PathBuf,
}
impl DaemonSpawner {
/// Create a new daemon spawner.
///
/// # Errors
///
/// Returns an error if the current executable path cannot be determined.
pub fn new(state_manager: StateManager) -> Result<Self> {
let executable_path = Self::executable_path()?;
Ok(Self {
state_manager,
executable_path,
})
}
/// Create a new daemon spawner with a custom executable path.
///
/// This is useful for testing or when spawning a different binary.
#[must_use]
pub const fn with_executable(state_manager: StateManager, executable_path: PathBuf) -> Self {
Self {
state_manager,
executable_path,
}
}
/// Spawn a background download job.
///
/// Returns the job ID for tracking. The job's PID and log file path
/// will be updated after spawning.
///
/// # Errors
///
/// Returns an error if the daemon process cannot be spawned.
pub fn spawn(&self, job: &mut DownloadJob) -> Result<JobId> {
let job_id = job.id;
// Set up log file path
let log_path = self.state_manager.job_log_path(job_id);
job.log_file = Some(log_path.clone());
// Save job state before spawning
self.state_manager.save_job(job)?;
// Open log file for stdout/stderr redirection
let log_file = OpenOptions::new()
.create(true)
.write(true)
.truncate(true)
.open(&log_path)
.map_err(|e| StateError::WriteFile {
path: log_path.clone(),
source: e,
})?;
let log_file_stderr = log_file.try_clone().map_err(|e| StateError::WriteFile {
path: log_path.clone(),
source: e,
})?;
// Spawn the daemon process
let child = self.spawn_detached(job_id, log_file, log_file_stderr)?;
// Update job with PID
let pid = child.id();
job.pid = Some(pid);
self.state_manager.save_job(job)?;
Ok(job_id)
}
/// Spawn a detached child process.
#[cfg(unix)]
fn spawn_detached(
&self,
job_id: JobId,
stdout: std::fs::File,
stderr: std::fs::File,
) -> Result<std::process::Child> {
use std::os::unix::process::CommandExt;
let child = Command::new(&self.executable_path)
.args([DAEMON_RUN_ARG, &job_id.to_string()])
.env(DAEMON_JOB_ID_ENV, job_id.to_string())
.stdin(Stdio::null())
.stdout(stdout)
.stderr(stderr)
.process_group(0) // Create new process group (detach from parent)
.spawn()
.map_err(|e| StateError::SpawnDaemon {
executable: self.executable_path.clone(),
source: e,
})?;
Ok(child)
}
/// Spawn a detached child process on Windows.
#[cfg(windows)]
fn spawn_detached(
&self,
job_id: JobId,
stdout: std::fs::File,
stderr: std::fs::File,
) -> Result<std::process::Child> {
use std::os::windows::process::CommandExt;
// CREATE_NEW_PROCESS_GROUP | DETACHED_PROCESS
const CREATE_NEW_PROCESS_GROUP: u32 = 0x00000200;
const DETACHED_PROCESS: u32 = 0x00000008;
let child = Command::new(&self.executable_path)
.args([DAEMON_RUN_ARG, &job_id.to_string()])
.env(DAEMON_JOB_ID_ENV, job_id.to_string())
.stdin(Stdio::null())
.stdout(stdout)
.stderr(stderr)
.creation_flags(CREATE_NEW_PROCESS_GROUP | DETACHED_PROCESS)
.spawn()
.map_err(|e| StateError::SpawnDaemon {
executable: self.executable_path.clone(),
source: e,
})?;
Ok(child)
}
/// Spawn a detached child process (fallback for other platforms).
#[cfg(not(any(unix, windows)))]
fn spawn_detached(
&self,
job_id: JobId,
stdout: std::fs::File,
stderr: std::fs::File,
) -> Result<std::process::Child> {
let child = Command::new(&self.executable_path)
.args([DAEMON_RUN_ARG, &job_id.to_string()])
.env(DAEMON_JOB_ID_ENV, job_id.to_string())
.stdin(Stdio::null())
.stdout(stdout)
.stderr(stderr)
.spawn()
.map_err(|e| StateError::SpawnDaemon {
executable: self.executable_path.clone(),
source: e,
})?;
Ok(child)
}
/// Get the path to the current executable.
fn executable_path() -> Result<PathBuf> {
std::env::current_exe().map_err(|e| StateError::ExecutablePath { source: e })
}
/// Returns the state manager reference.
#[must_use]
pub const fn state_manager(&self) -> &StateManager {
&self.state_manager
}
/// Returns the executable path.
#[must_use]
pub const fn executable(&self) -> &PathBuf {
&self.executable_path
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::InstrumentTask;
use std::path::PathBuf;
use tempfile::TempDir;
fn create_test_job() -> DownloadJob {
let tasks = vec![InstrumentTask::new(
"EURUSD".to_string(),
"2024-01-01".to_string(),
"2024-01-02".to_string(),
PathBuf::from("/tmp/eurusd.csv"),
"csv".to_string(),
"tick".to_string(),
48,
)];
DownloadJob::new(tasks, 4)
}
#[test]
fn test_daemon_spawner_creation() {
let temp_dir = TempDir::new().unwrap();
let state_manager = StateManager::new(temp_dir.path().to_path_buf()).unwrap();
// Use a known executable for testing
let spawner = DaemonSpawner::with_executable(state_manager, PathBuf::from("/bin/echo"));
assert_eq!(spawner.executable(), &PathBuf::from("/bin/echo"));
}
#[test]
fn test_daemon_spawner_with_executable() {
let temp_dir = TempDir::new().unwrap();
let state_manager = StateManager::new(temp_dir.path().to_path_buf()).unwrap();
let custom_path = PathBuf::from("/custom/paracas");
let spawner = DaemonSpawner::with_executable(state_manager, custom_path.clone());
assert_eq!(spawner.executable(), &custom_path);
}
#[test]
fn test_spawn_sets_log_file() {
let temp_dir = TempDir::new().unwrap();
let state_manager = StateManager::new(temp_dir.path().to_path_buf()).unwrap();
// Use /bin/true or /usr/bin/true for a quick-exit process
#[cfg(unix)]
let exe_path = if PathBuf::from("/bin/true").exists() {
PathBuf::from("/bin/true")
} else {
PathBuf::from("/usr/bin/true")
};
#[cfg(not(unix))]
let exe_path = PathBuf::from("cmd.exe");
let spawner = DaemonSpawner::with_executable(state_manager.clone(), exe_path);
let mut job = create_test_job();
let job_id = job.id;
// Spawn the job
let result = spawner.spawn(&mut job);
// On CI or systems where the binary doesn't exist, this might fail
if result.is_ok() {
assert!(job.log_file.is_some());
assert!(job.pid.is_some());
// Verify log file path is correct
let expected_log_path = state_manager.job_log_path(job_id);
assert_eq!(job.log_file.as_ref().unwrap(), &expected_log_path);
}
}
}
+317
View File
@@ -0,0 +1,317 @@
//! Download job definitions and types.
use chrono::{DateTime, Utc};
use serde::{Deserialize, Serialize};
use std::path::PathBuf;
use uuid::Uuid;
/// Unique identifier for a download job.
pub type JobId = Uuid;
/// Status of a download job or task.
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize, Default)]
#[serde(rename_all = "lowercase")]
pub enum JobStatus {
/// Job is queued but not yet started.
#[default]
Pending,
/// Job is currently running.
Running,
/// Job completed successfully.
Completed,
/// Job failed with an error.
Failed,
/// Job was cancelled by the user.
Cancelled,
}
impl JobStatus {
/// Returns true if the job is in a terminal state.
#[must_use]
pub const fn is_finished(&self) -> bool {
matches!(self, Self::Completed | Self::Failed | Self::Cancelled)
}
/// Returns the status as a string identifier.
#[must_use]
pub const fn as_str(&self) -> &'static str {
match self {
Self::Pending => "pending",
Self::Running => "running",
Self::Completed => "completed",
Self::Failed => "failed",
Self::Cancelled => "cancelled",
}
}
}
impl std::fmt::Display for JobStatus {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
write!(f, "{}", self.as_str())
}
}
/// A download task for a single instrument within a job.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct InstrumentTask {
/// The instrument identifier (e.g., "EURUSD").
pub instrument_id: String,
/// Start date for the download (inclusive).
pub start_date: String,
/// End date for the download (inclusive).
pub end_date: String,
/// Output file path for this instrument's data.
pub output_path: PathBuf,
/// Output format (e.g., "csv", "json", "parquet").
pub format: String,
/// Timeframe for aggregation (e.g., "tick", "m1", "h1").
pub timeframe: String,
/// Current status of this task.
pub status: JobStatus,
/// Number of hours completed for this task.
pub hours_completed: u32,
/// Total number of hours to download.
pub hours_total: u32,
/// Number of ticks downloaded so far.
pub ticks_downloaded: u64,
/// Number of bytes written to output file.
pub bytes_written: u64,
/// Error message if the task failed.
pub error_message: Option<String>,
}
impl InstrumentTask {
/// Creates a new instrument task.
#[must_use]
pub const fn new(
instrument_id: String,
start_date: String,
end_date: String,
output_path: PathBuf,
format: String,
timeframe: String,
hours_total: u32,
) -> Self {
Self {
instrument_id,
start_date,
end_date,
output_path,
format,
timeframe,
status: JobStatus::Pending,
hours_completed: 0,
hours_total,
ticks_downloaded: 0,
bytes_written: 0,
error_message: None,
}
}
/// Returns the progress percentage for this task.
#[must_use]
pub fn progress_percent(&self) -> f64 {
if self.hours_total == 0 {
return 0.0;
}
(self.hours_completed as f64 / self.hours_total as f64) * 100.0
}
}
/// A complete download job containing one or more instrument tasks.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct DownloadJob {
/// Unique identifier for this job.
pub id: JobId,
/// Timestamp when the job was created.
pub created_at: DateTime<Utc>,
/// Timestamp when the job started running.
pub started_at: Option<DateTime<Utc>>,
/// Timestamp when the job completed (success, failure, or cancellation).
pub completed_at: Option<DateTime<Utc>>,
/// Current status of the job.
pub status: JobStatus,
/// List of instrument download tasks.
pub tasks: Vec<InstrumentTask>,
/// Number of concurrent downloads.
pub concurrency: usize,
/// Process ID of the daemon running this job.
pub pid: Option<u32>,
/// Path to the log file for this job.
pub log_file: Option<PathBuf>,
}
impl DownloadJob {
/// Creates a new download job with the given tasks.
#[must_use]
pub fn new(tasks: Vec<InstrumentTask>, concurrency: usize) -> Self {
Self {
id: Uuid::new_v4(),
created_at: Utc::now(),
started_at: None,
completed_at: None,
status: JobStatus::Pending,
tasks,
concurrency,
pid: None,
log_file: None,
}
}
/// Returns the overall progress percentage across all tasks.
#[must_use]
pub fn progress_percent(&self) -> f64 {
let total_hours: u32 = self.tasks.iter().map(|t| t.hours_total).sum();
let completed_hours: u32 = self.tasks.iter().map(|t| t.hours_completed).sum();
if total_hours == 0 {
return 0.0;
}
(completed_hours as f64 / total_hours as f64) * 100.0
}
/// Returns true if the job is in a terminal state.
#[must_use]
pub const fn is_finished(&self) -> bool {
self.status.is_finished()
}
/// Marks the job as started with the current timestamp and process ID.
pub fn mark_started(&mut self, pid: u32) {
self.status = JobStatus::Running;
self.started_at = Some(Utc::now());
self.pid = Some(pid);
}
/// Marks the job as completed successfully.
pub fn mark_completed(&mut self) {
self.status = JobStatus::Completed;
self.completed_at = Some(Utc::now());
}
/// Marks the job as failed with an optional error message.
pub fn mark_failed(&mut self, error: Option<String>) {
self.status = JobStatus::Failed;
self.completed_at = Some(Utc::now());
// If an error message is provided, set it on any running tasks
if let Some(ref msg) = error {
for task in &mut self.tasks {
if task.status == JobStatus::Running {
task.status = JobStatus::Failed;
task.error_message = Some(msg.clone());
}
}
}
}
/// Marks the job as cancelled.
pub fn mark_cancelled(&mut self) {
self.status = JobStatus::Cancelled;
self.completed_at = Some(Utc::now());
// Cancel any pending or running tasks
for task in &mut self.tasks {
if !task.status.is_finished() {
task.status = JobStatus::Cancelled;
}
}
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_job_status_is_finished() {
assert!(!JobStatus::Pending.is_finished());
assert!(!JobStatus::Running.is_finished());
assert!(JobStatus::Completed.is_finished());
assert!(JobStatus::Failed.is_finished());
assert!(JobStatus::Cancelled.is_finished());
}
#[test]
fn test_instrument_task_progress() {
let mut task = InstrumentTask::new(
"EURUSD".to_string(),
"2024-01-01".to_string(),
"2024-01-02".to_string(),
PathBuf::from("/tmp/eurusd.csv"),
"csv".to_string(),
"tick".to_string(),
48,
);
assert_eq!(task.progress_percent(), 0.0);
task.hours_completed = 24;
assert!((task.progress_percent() - 50.0).abs() < 0.001);
task.hours_completed = 48;
assert!((task.progress_percent() - 100.0).abs() < 0.001);
}
#[test]
fn test_download_job_progress() {
let tasks = vec![
InstrumentTask::new(
"EURUSD".to_string(),
"2024-01-01".to_string(),
"2024-01-02".to_string(),
PathBuf::from("/tmp/eurusd.csv"),
"csv".to_string(),
"tick".to_string(),
48,
),
InstrumentTask::new(
"GBPUSD".to_string(),
"2024-01-01".to_string(),
"2024-01-02".to_string(),
PathBuf::from("/tmp/gbpusd.csv"),
"csv".to_string(),
"tick".to_string(),
48,
),
];
let mut job = DownloadJob::new(tasks, 4);
assert_eq!(job.progress_percent(), 0.0);
job.tasks[0].hours_completed = 48;
assert!((job.progress_percent() - 50.0).abs() < 0.001);
job.tasks[1].hours_completed = 48;
assert!((job.progress_percent() - 100.0).abs() < 0.001);
}
#[test]
fn test_download_job_lifecycle() {
let tasks = vec![InstrumentTask::new(
"EURUSD".to_string(),
"2024-01-01".to_string(),
"2024-01-02".to_string(),
PathBuf::from("/tmp/eurusd.csv"),
"csv".to_string(),
"tick".to_string(),
48,
)];
let mut job = DownloadJob::new(tasks, 4);
assert_eq!(job.status, JobStatus::Pending);
assert!(job.started_at.is_none());
assert!(!job.is_finished());
job.mark_started(12345);
assert_eq!(job.status, JobStatus::Running);
assert!(job.started_at.is_some());
assert_eq!(job.pid, Some(12345));
assert!(!job.is_finished());
job.mark_completed();
assert_eq!(job.status, JobStatus::Completed);
assert!(job.completed_at.is_some());
assert!(job.is_finished());
}
}
+27
View File
@@ -0,0 +1,27 @@
//! Background job management for paracas tick data downloader.
//!
//! This crate provides state management and job tracking for background
//! download operations:
//!
//! - [`JobId`] - Unique identifier for download jobs
//! - [`JobStatus`] - Current status of a job
//! - [`InstrumentTask`] - Download task for a single instrument
//! - [`DownloadJob`] - Complete download job with multiple tasks
//! - [`StateManager`] - Persistent state storage and retrieval
//! - [`DaemonSpawner`] - Spawns detached daemon processes for background downloads
//! - [`DaemonProgress`] - Thread-safe progress tracking for daemon jobs
#![doc(issue_tracker_base_url = "https://github.com/factordynamics/paracas/issues/")]
#![cfg_attr(docsrs, feature(doc_cfg, doc_auto_cfg))]
#![warn(missing_docs)]
#![forbid(unsafe_code)]
mod daemon;
mod job;
mod progress;
mod state;
pub use daemon::{DAEMON_JOB_ID_ENV, DAEMON_RUN_ARG, DaemonSpawner};
pub use job::{DownloadJob, InstrumentTask, JobId, JobStatus};
pub use progress::DaemonProgress;
pub use state::{Result, StateError, StateManager};
+490
View File
@@ -0,0 +1,490 @@
//! Progress tracking for daemon downloads.
//!
//! This module provides thread-safe progress tracking for daemon jobs,
//! including periodic checkpointing to disk for crash recovery.
use crate::{DownloadJob, JobStatus, StateError, StateManager};
use std::sync::Arc;
use std::time::{Duration, Instant};
use tokio::sync::RwLock;
/// Thread-safe progress tracker for daemon jobs.
///
/// The `DaemonProgress` struct provides a way to track download progress
/// from multiple concurrent tasks while ensuring periodic checkpoints
/// are saved to disk for crash recovery.
#[derive(Debug)]
pub struct DaemonProgress {
/// State manager for persistence.
state_manager: StateManager,
/// The job being tracked (protected by RwLock for concurrent access).
job: Arc<RwLock<DownloadJob>>,
/// Minimum interval between saves.
save_interval: Duration,
/// Last time state was saved to disk.
last_save: std::sync::Mutex<Instant>,
}
impl DaemonProgress {
/// Default save interval for checkpointing (10 seconds).
pub const DEFAULT_SAVE_INTERVAL: Duration = Duration::from_secs(10);
/// Create a new progress tracker.
///
/// The tracker will periodically save checkpoints to disk at the
/// default interval of 10 seconds.
#[must_use]
pub fn new(state_manager: StateManager, job: DownloadJob) -> Self {
Self {
state_manager,
job: Arc::new(RwLock::new(job)),
save_interval: Self::DEFAULT_SAVE_INTERVAL,
last_save: std::sync::Mutex::new(Instant::now()),
}
}
/// Create a new progress tracker with a custom save interval.
#[must_use]
pub fn with_save_interval(
state_manager: StateManager,
job: DownloadJob,
save_interval: Duration,
) -> Self {
Self {
state_manager,
job: Arc::new(RwLock::new(job)),
save_interval,
last_save: std::sync::Mutex::new(Instant::now()),
}
}
/// Update progress for a specific task.
///
/// This updates the hours completed and ticks downloaded for the task
/// at the given index. If enough time has passed since the last save,
/// the state will be checkpointed to disk.
///
/// # Arguments
///
/// * `task_idx` - Index of the task to update
/// * `hours` - Number of hours completed
/// * `ticks` - Number of ticks downloaded
pub async fn update_task_progress(&self, task_idx: usize, hours: u64, ticks: u64) {
{
let mut job = self.job.write().await;
if let Some(task) = job.tasks.get_mut(task_idx) {
task.hours_completed = hours as u32;
task.ticks_downloaded = ticks;
if task.status == JobStatus::Pending {
task.status = JobStatus::Running;
}
}
}
// Check if we should save
self.maybe_save_checkpoint().await;
}
/// Mark a task as completed.
///
/// This updates the task status to `Completed` and records the
/// final byte count.
///
/// # Arguments
///
/// * `task_idx` - Index of the task to mark as completed
/// * `bytes` - Total bytes written for this task
pub async fn mark_task_completed(&self, task_idx: usize, bytes: u64) {
{
let mut job = self.job.write().await;
if let Some(task) = job.tasks.get_mut(task_idx) {
task.status = JobStatus::Completed;
task.bytes_written = bytes;
task.hours_completed = task.hours_total;
}
}
// Always save on task completion
let _ = self.save_checkpoint().await;
}
/// Mark a task as failed.
///
/// This updates the task status to `Failed` and records the error message.
///
/// # Arguments
///
/// * `task_idx` - Index of the task to mark as failed
/// * `error` - Error message describing the failure
pub async fn mark_task_failed(&self, task_idx: usize, error: &str) {
{
let mut job = self.job.write().await;
if let Some(task) = job.tasks.get_mut(task_idx) {
task.status = JobStatus::Failed;
task.error_message = Some(error.to_string());
}
}
// Always save on task failure
let _ = self.save_checkpoint().await;
}
/// Mark a task as running.
///
/// This updates the task status to `Running`.
///
/// # Arguments
///
/// * `task_idx` - Index of the task to mark as running
pub async fn mark_task_running(&self, task_idx: usize) {
{
let mut job = self.job.write().await;
if let Some(task) = job.tasks.get_mut(task_idx) {
task.status = JobStatus::Running;
}
}
// Save when task starts
let _ = self.save_checkpoint().await;
}
/// Mark the entire job as completed.
///
/// Call this when all tasks have finished successfully.
pub async fn mark_job_completed(&self) {
{
let mut job = self.job.write().await;
job.mark_completed();
}
// Always save on job completion
let _ = self.save_checkpoint().await;
}
/// Mark the entire job as failed.
///
/// Call this when the job fails due to a critical error.
///
/// # Arguments
///
/// * `error` - Error message describing the failure
pub async fn mark_job_failed(&self, error: &str) {
{
let mut job = self.job.write().await;
job.mark_failed(Some(error.to_string()));
}
// Always save on job failure
let _ = self.save_checkpoint().await;
}
/// Save current progress to disk (called periodically).
///
/// This forces a checkpoint save regardless of the save interval.
///
/// # Errors
///
/// Returns an error if the state cannot be saved to disk.
pub async fn save_checkpoint(&self) -> Result<(), StateError> {
let job = self.job.read().await;
self.state_manager.save_job(&job)?;
// Update last save time
if let Ok(mut last_save) = self.last_save.lock() {
*last_save = Instant::now();
}
Ok(())
}
/// Check if enough time has passed and save if needed.
async fn maybe_save_checkpoint(&self) {
let should_save = self
.last_save
.lock()
.map_or(true, |last_save| last_save.elapsed() >= self.save_interval);
if should_save {
let _ = self.save_checkpoint().await;
}
}
/// Get current job state.
///
/// Returns a clone of the current job state.
pub async fn job(&self) -> DownloadJob {
self.job.read().await.clone()
}
/// Get the number of completed tasks.
pub async fn completed_tasks(&self) -> usize {
let job = self.job.read().await;
job.tasks
.iter()
.filter(|t| t.status == JobStatus::Completed)
.count()
}
/// Get the number of failed tasks.
pub async fn failed_tasks(&self) -> usize {
let job = self.job.read().await;
job.tasks
.iter()
.filter(|t| t.status == JobStatus::Failed)
.count()
}
/// Get the total number of tasks.
pub async fn total_tasks(&self) -> usize {
let job = self.job.read().await;
job.tasks.len()
}
/// Get the current progress percentage.
pub async fn progress_percent(&self) -> f64 {
let job = self.job.read().await;
job.progress_percent()
}
/// Check if all tasks are finished.
pub async fn all_tasks_finished(&self) -> bool {
let job = self.job.read().await;
job.tasks.iter().all(|t| t.status.is_finished())
}
/// Returns a reference to the state manager.
#[must_use]
pub const fn state_manager(&self) -> &StateManager {
&self.state_manager
}
}
impl Clone for DaemonProgress {
fn clone(&self) -> Self {
Self {
state_manager: self.state_manager.clone(),
job: Arc::clone(&self.job),
save_interval: self.save_interval,
last_save: std::sync::Mutex::new(self.last_save.lock().map_or(Instant::now(), |g| *g)),
}
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::InstrumentTask;
use std::path::PathBuf;
use tempfile::TempDir;
fn create_test_job() -> DownloadJob {
let tasks = vec![
InstrumentTask::new(
"EURUSD".to_string(),
"2024-01-01".to_string(),
"2024-01-02".to_string(),
PathBuf::from("/tmp/eurusd.csv"),
"csv".to_string(),
"tick".to_string(),
48,
),
InstrumentTask::new(
"GBPUSD".to_string(),
"2024-01-01".to_string(),
"2024-01-02".to_string(),
PathBuf::from("/tmp/gbpusd.csv"),
"csv".to_string(),
"tick".to_string(),
48,
),
];
DownloadJob::new(tasks, 4)
}
#[tokio::test]
async fn test_progress_tracker_creation() {
let temp_dir = TempDir::new().unwrap();
let state_manager = StateManager::new(temp_dir.path().to_path_buf()).unwrap();
let job = create_test_job();
let progress = DaemonProgress::new(state_manager, job);
let current = progress.job().await;
assert_eq!(current.tasks.len(), 2);
assert_eq!(current.status, JobStatus::Pending);
}
#[tokio::test]
async fn test_update_task_progress() {
let temp_dir = TempDir::new().unwrap();
let state_manager = StateManager::new(temp_dir.path().to_path_buf()).unwrap();
let job = create_test_job();
let progress = DaemonProgress::new(state_manager, job);
progress.update_task_progress(0, 24, 1_000_000).await;
let current = progress.job().await;
assert_eq!(current.tasks[0].hours_completed, 24);
assert_eq!(current.tasks[0].ticks_downloaded, 1_000_000);
assert_eq!(current.tasks[0].status, JobStatus::Running);
}
#[tokio::test]
async fn test_mark_task_completed() {
let temp_dir = TempDir::new().unwrap();
let state_manager = StateManager::new(temp_dir.path().to_path_buf()).unwrap();
let job = create_test_job();
let job_id = job.id;
let progress = DaemonProgress::new(state_manager.clone(), job);
progress.mark_task_completed(0, 1024 * 1024).await;
let current = progress.job().await;
assert_eq!(current.tasks[0].status, JobStatus::Completed);
assert_eq!(current.tasks[0].bytes_written, 1024 * 1024);
assert_eq!(current.tasks[0].hours_completed, 48); // Should be set to total
// Verify saved to disk
let loaded = state_manager.load_job(job_id).unwrap();
assert_eq!(loaded.tasks[0].status, JobStatus::Completed);
}
#[tokio::test]
async fn test_mark_task_failed() {
let temp_dir = TempDir::new().unwrap();
let state_manager = StateManager::new(temp_dir.path().to_path_buf()).unwrap();
let job = create_test_job();
let job_id = job.id;
let progress = DaemonProgress::new(state_manager.clone(), job);
progress.mark_task_failed(0, "Connection timeout").await;
let current = progress.job().await;
assert_eq!(current.tasks[0].status, JobStatus::Failed);
assert_eq!(
current.tasks[0].error_message,
Some("Connection timeout".to_string())
);
// Verify saved to disk
let loaded = state_manager.load_job(job_id).unwrap();
assert_eq!(loaded.tasks[0].status, JobStatus::Failed);
}
#[tokio::test]
async fn test_mark_job_completed() {
let temp_dir = TempDir::new().unwrap();
let state_manager = StateManager::new(temp_dir.path().to_path_buf()).unwrap();
let job = create_test_job();
let progress = DaemonProgress::new(state_manager, job);
progress.mark_job_completed().await;
let current = progress.job().await;
assert_eq!(current.status, JobStatus::Completed);
assert!(current.completed_at.is_some());
}
#[tokio::test]
async fn test_mark_job_failed() {
let temp_dir = TempDir::new().unwrap();
let state_manager = StateManager::new(temp_dir.path().to_path_buf()).unwrap();
let job = create_test_job();
let progress = DaemonProgress::new(state_manager, job);
progress.mark_job_failed("Fatal error").await;
let current = progress.job().await;
assert_eq!(current.status, JobStatus::Failed);
assert!(current.completed_at.is_some());
}
#[tokio::test]
async fn test_completed_tasks_count() {
let temp_dir = TempDir::new().unwrap();
let state_manager = StateManager::new(temp_dir.path().to_path_buf()).unwrap();
let job = create_test_job();
let progress = DaemonProgress::new(state_manager, job);
assert_eq!(progress.completed_tasks().await, 0);
progress.mark_task_completed(0, 1024).await;
assert_eq!(progress.completed_tasks().await, 1);
progress.mark_task_completed(1, 2048).await;
assert_eq!(progress.completed_tasks().await, 2);
}
#[tokio::test]
async fn test_all_tasks_finished() {
let temp_dir = TempDir::new().unwrap();
let state_manager = StateManager::new(temp_dir.path().to_path_buf()).unwrap();
let job = create_test_job();
let progress = DaemonProgress::new(state_manager, job);
assert!(!progress.all_tasks_finished().await);
progress.mark_task_completed(0, 1024).await;
assert!(!progress.all_tasks_finished().await);
progress.mark_task_failed(1, "error").await;
assert!(progress.all_tasks_finished().await);
}
#[tokio::test]
async fn test_progress_percent() {
let temp_dir = TempDir::new().unwrap();
let state_manager = StateManager::new(temp_dir.path().to_path_buf()).unwrap();
let job = create_test_job();
let progress = DaemonProgress::new(state_manager, job);
// Initially 0%
assert!((progress.progress_percent().await - 0.0).abs() < 0.001);
// Update first task to 50%
progress.update_task_progress(0, 24, 100).await;
// 24/96 total hours = 25%
assert!((progress.progress_percent().await - 25.0).abs() < 0.001);
}
#[tokio::test]
async fn test_custom_save_interval() {
let temp_dir = TempDir::new().unwrap();
let state_manager = StateManager::new(temp_dir.path().to_path_buf()).unwrap();
let job = create_test_job();
let progress =
DaemonProgress::with_save_interval(state_manager, job, Duration::from_secs(1));
assert_eq!(progress.save_interval, Duration::from_secs(1));
}
#[tokio::test]
async fn test_progress_clone() {
let temp_dir = TempDir::new().unwrap();
let state_manager = StateManager::new(temp_dir.path().to_path_buf()).unwrap();
let job = create_test_job();
let progress = DaemonProgress::new(state_manager, job);
let cloned = progress.clone();
// Both should share the same job state
progress.update_task_progress(0, 10, 100).await;
let original_job = progress.job().await;
let cloned_job = cloned.job().await;
assert_eq!(
original_job.tasks[0].hours_completed,
cloned_job.tasks[0].hours_completed
);
}
}
+499
View File
@@ -0,0 +1,499 @@
//! State management for persistent job storage.
use crate::{DownloadJob, JobId, JobStatus};
use directories::ProjectDirs;
use std::fs;
use std::path::{Path, PathBuf};
use thiserror::Error;
/// Errors that can occur during state management operations.
#[derive(Error, Debug)]
pub enum StateError {
/// Failed to determine the application data directory.
#[error("Failed to determine application data directory")]
NoDataDir,
/// Failed to create a directory.
#[error("Failed to create directory '{path}': {source}")]
CreateDir {
/// The path that could not be created.
path: PathBuf,
/// The underlying I/O error.
source: std::io::Error,
},
/// Failed to read a file.
#[error("Failed to read file '{path}': {source}")]
ReadFile {
/// The path that could not be read.
path: PathBuf,
/// The underlying I/O error.
source: std::io::Error,
},
/// Failed to write a file.
#[error("Failed to write file '{path}': {source}")]
WriteFile {
/// The path that could not be written.
path: PathBuf,
/// The underlying I/O error.
source: std::io::Error,
},
/// Failed to delete a file.
#[error("Failed to delete file '{path}': {source}")]
DeleteFile {
/// The path that could not be deleted.
path: PathBuf,
/// The underlying I/O error.
source: std::io::Error,
},
/// Failed to parse JSON.
#[error("Failed to parse job file '{path}': {source}")]
ParseJson {
/// The path that could not be parsed.
path: PathBuf,
/// The underlying JSON error.
source: serde_json::Error,
},
/// Failed to serialize JSON.
#[error("Failed to serialize job: {0}")]
SerializeJson(#[from] serde_json::Error),
/// Job not found.
#[error("Job not found: {0}")]
JobNotFound(JobId),
/// Failed to read directory.
#[error("Failed to read directory '{path}': {source}")]
ReadDir {
/// The path that could not be read.
path: PathBuf,
/// The underlying I/O error.
source: std::io::Error,
},
/// Failed to spawn daemon process.
#[error("Failed to spawn daemon process '{executable}': {source}")]
SpawnDaemon {
/// The executable that could not be spawned.
executable: PathBuf,
/// The underlying I/O error.
source: std::io::Error,
},
/// Failed to determine executable path.
#[error("Failed to determine executable path: {source}")]
ExecutablePath {
/// The underlying I/O error.
source: std::io::Error,
},
}
/// Result type for state operations.
pub type Result<T> = std::result::Result<T, StateError>;
/// Manages persistent state for download jobs.
///
/// Jobs are stored as JSON files in `~/.paracas/jobs/` with log files
/// stored in `~/.paracas/logs/`.
#[derive(Debug, Clone)]
pub struct StateManager {
/// Base directory for state storage.
base_path: PathBuf,
/// Directory for job JSON files.
jobs_path: PathBuf,
/// Directory for job log files.
logs_path: PathBuf,
}
impl StateManager {
/// Creates a new state manager with the given base path.
///
/// Creates the necessary subdirectories if they don't exist.
///
/// # Errors
///
/// Returns an error if the directories cannot be created.
pub fn new(base_path: PathBuf) -> Result<Self> {
let jobs_path = base_path.join("jobs");
let logs_path = base_path.join("logs");
// Create directories if they don't exist
for path in [&base_path, &jobs_path, &logs_path] {
if !path.exists() {
fs::create_dir_all(path).map_err(|e| StateError::CreateDir {
path: path.clone(),
source: e,
})?;
}
}
Ok(Self {
base_path,
jobs_path,
logs_path,
})
}
/// Returns the default path for paracas state storage.
///
/// Uses the `directories` crate to find the appropriate location:
/// - Linux: `~/.local/share/paracas/`
/// - macOS: `~/Library/Application Support/paracas/`
/// - Windows: `C:\Users\<User>\AppData\Roaming\paracas\`
///
/// Falls back to `~/.paracas/` if the platform-specific location
/// cannot be determined.
#[must_use]
pub fn default_path() -> PathBuf {
ProjectDirs::from("", "", "paracas").map_or_else(dirs_fallback, |proj_dirs| {
proj_dirs.data_dir().to_path_buf()
})
}
/// Creates a state manager at the default path.
///
/// # Errors
///
/// Returns an error if the directories cannot be created.
pub fn with_default_path() -> Result<Self> {
Self::new(Self::default_path())
}
/// Returns the base path for state storage.
#[must_use]
pub fn base_path(&self) -> &Path {
&self.base_path
}
/// Returns the path to a job's state file.
#[must_use]
pub fn job_state_path(&self, job_id: JobId) -> PathBuf {
self.jobs_path.join(format!("{job_id}.json"))
}
/// Returns the path to a job's log file.
#[must_use]
pub fn job_log_path(&self, job_id: JobId) -> PathBuf {
self.logs_path.join(format!("{job_id}.log"))
}
/// Saves a job to persistent storage.
///
/// # Errors
///
/// Returns an error if the job cannot be serialized or written to disk.
pub fn save_job(&self, job: &DownloadJob) -> Result<()> {
let path = self.job_state_path(job.id);
let json = serde_json::to_string_pretty(job)?;
fs::write(&path, json).map_err(|e| StateError::WriteFile { path, source: e })
}
/// Loads a job from persistent storage.
///
/// # Errors
///
/// Returns an error if the job file cannot be read or parsed.
pub fn load_job(&self, job_id: JobId) -> Result<DownloadJob> {
let path = self.job_state_path(job_id);
if !path.exists() {
return Err(StateError::JobNotFound(job_id));
}
let content = fs::read_to_string(&path).map_err(|e| StateError::ReadFile {
path: path.clone(),
source: e,
})?;
serde_json::from_str(&content).map_err(|e| StateError::ParseJson { path, source: e })
}
/// Lists all jobs in persistent storage.
///
/// Returns jobs sorted by creation time (newest first).
///
/// # Errors
///
/// Returns an error if the jobs directory cannot be read.
pub fn list_jobs(&self) -> Result<Vec<DownloadJob>> {
let entries = fs::read_dir(&self.jobs_path).map_err(|e| StateError::ReadDir {
path: self.jobs_path.clone(),
source: e,
})?;
let mut jobs = Vec::new();
for entry in entries {
let entry = entry.map_err(|e| StateError::ReadDir {
path: self.jobs_path.clone(),
source: e,
})?;
let path = entry.path();
if path.extension().is_some_and(|ext| ext == "json") {
let content = fs::read_to_string(&path).map_err(|e| StateError::ReadFile {
path: path.clone(),
source: e,
})?;
match serde_json::from_str::<DownloadJob>(&content) {
Ok(job) => jobs.push(job),
Err(e) => {
// Log warning but continue - don't fail on corrupt files
eprintln!("Warning: Failed to parse job file {:?}: {}", path, e);
}
}
}
}
// Sort by creation time, newest first
jobs.sort_by(|a, b| b.created_at.cmp(&a.created_at));
Ok(jobs)
}
/// Deletes a job from persistent storage.
///
/// Also deletes the associated log file if it exists.
///
/// # Errors
///
/// Returns an error if the job file cannot be deleted.
pub fn delete_job(&self, job_id: JobId) -> Result<()> {
let state_path = self.job_state_path(job_id);
if !state_path.exists() {
return Err(StateError::JobNotFound(job_id));
}
fs::remove_file(&state_path).map_err(|e| StateError::DeleteFile {
path: state_path,
source: e,
})?;
// Also delete log file if it exists
let log_path = self.job_log_path(job_id);
if log_path.exists() {
let _ = fs::remove_file(&log_path); // Ignore errors for log file
}
Ok(())
}
/// Returns all active (pending or running) jobs.
///
/// # Errors
///
/// Returns an error if jobs cannot be listed.
pub fn active_jobs(&self) -> Result<Vec<DownloadJob>> {
let jobs = self.list_jobs()?;
Ok(jobs.into_iter().filter(|j| !j.is_finished()).collect())
}
/// Checks if a process with the given PID is still running.
#[must_use]
pub fn is_process_running(pid: u32) -> bool {
// Use kill with signal 0 to check if process exists
// This doesn't actually send a signal, just checks if the process exists
#[cfg(unix)]
{
use std::process::Command;
Command::new("kill")
.args(["-0", &pid.to_string()])
.output()
.map(|output| output.status.success())
.unwrap_or(false)
}
#[cfg(windows)]
{
// On Windows, use tasklist to check if process exists
use std::process::Command;
Command::new("tasklist")
.args(["/FI", &format!("PID eq {}", pid)])
.output()
.map(|output| {
let stdout = String::from_utf8_lossy(&output.stdout);
stdout.contains(&pid.to_string())
})
.unwrap_or(false)
}
#[cfg(not(any(unix, windows)))]
{
// On other platforms, assume the process is not running
false
}
}
/// Cleans up stale jobs where the process is no longer running.
///
/// Marks running jobs as failed if their daemon process has died.
///
/// # Errors
///
/// Returns an error if jobs cannot be listed or updated.
pub fn cleanup_stale_jobs(&self) -> Result<Vec<JobId>> {
let jobs = self.list_jobs()?;
let mut cleaned = Vec::new();
for mut job in jobs {
if job.status == JobStatus::Running {
let is_stale = job.pid.is_none_or(|pid| !Self::is_process_running(pid));
if is_stale {
job.mark_failed(Some("Daemon process died unexpectedly".to_string()));
self.save_job(&job)?;
cleaned.push(job.id);
}
}
}
Ok(cleaned)
}
}
/// Fallback for determining home directory.
fn dirs_fallback() -> PathBuf {
std::env::var("HOME")
.map(PathBuf::from)
.unwrap_or_else(|_| PathBuf::from("."))
.join(".paracas")
}
#[cfg(test)]
mod tests {
use super::*;
use crate::InstrumentTask;
use tempfile::TempDir;
fn create_test_job() -> DownloadJob {
let tasks = vec![InstrumentTask::new(
"EURUSD".to_string(),
"2024-01-01".to_string(),
"2024-01-02".to_string(),
PathBuf::from("/tmp/eurusd.csv"),
"csv".to_string(),
"tick".to_string(),
48,
)];
DownloadJob::new(tasks, 4)
}
#[test]
fn test_state_manager_creation() {
let temp_dir = TempDir::new().unwrap();
let manager = StateManager::new(temp_dir.path().to_path_buf()).unwrap();
assert!(manager.base_path().exists());
assert!(temp_dir.path().join("jobs").exists());
assert!(temp_dir.path().join("logs").exists());
}
#[test]
fn test_save_and_load_job() {
let temp_dir = TempDir::new().unwrap();
let manager = StateManager::new(temp_dir.path().to_path_buf()).unwrap();
let job = create_test_job();
let job_id = job.id;
manager.save_job(&job).unwrap();
let loaded = manager.load_job(job_id).unwrap();
assert_eq!(loaded.id, job_id);
assert_eq!(loaded.status, JobStatus::Pending);
assert_eq!(loaded.tasks.len(), 1);
}
#[test]
fn test_list_jobs() {
let temp_dir = TempDir::new().unwrap();
let manager = StateManager::new(temp_dir.path().to_path_buf()).unwrap();
let job1 = create_test_job();
let job2 = create_test_job();
manager.save_job(&job1).unwrap();
manager.save_job(&job2).unwrap();
let jobs = manager.list_jobs().unwrap();
assert_eq!(jobs.len(), 2);
}
#[test]
fn test_delete_job() {
let temp_dir = TempDir::new().unwrap();
let manager = StateManager::new(temp_dir.path().to_path_buf()).unwrap();
let job = create_test_job();
let job_id = job.id;
manager.save_job(&job).unwrap();
assert!(manager.load_job(job_id).is_ok());
manager.delete_job(job_id).unwrap();
assert!(matches!(
manager.load_job(job_id),
Err(StateError::JobNotFound(_))
));
}
#[test]
fn test_active_jobs() {
let temp_dir = TempDir::new().unwrap();
let manager = StateManager::new(temp_dir.path().to_path_buf()).unwrap();
let mut pending_job = create_test_job();
let mut completed_job = create_test_job();
completed_job.mark_completed();
manager.save_job(&pending_job).unwrap();
manager.save_job(&completed_job).unwrap();
pending_job.mark_started(12345);
manager.save_job(&pending_job).unwrap();
let active = manager.active_jobs().unwrap();
assert_eq!(active.len(), 1);
assert_eq!(active[0].id, pending_job.id);
}
#[test]
fn test_job_not_found() {
let temp_dir = TempDir::new().unwrap();
let manager = StateManager::new(temp_dir.path().to_path_buf()).unwrap();
let result = manager.load_job(uuid::Uuid::new_v4());
assert!(matches!(result, Err(StateError::JobNotFound(_))));
}
#[test]
fn test_job_state_path() {
let temp_dir = TempDir::new().unwrap();
let manager = StateManager::new(temp_dir.path().to_path_buf()).unwrap();
let job_id = uuid::Uuid::new_v4();
let path = manager.job_state_path(job_id);
assert!(path.to_string_lossy().contains("jobs"));
assert!(path.to_string_lossy().ends_with(".json"));
}
#[test]
fn test_job_log_path() {
let temp_dir = TempDir::new().unwrap();
let manager = StateManager::new(temp_dir.path().to_path_buf()).unwrap();
let job_id = uuid::Uuid::new_v4();
let path = manager.job_log_path(job_id);
assert!(path.to_string_lossy().contains("logs"));
assert!(path.to_string_lossy().ends_with(".log"));
}
}
+18
View File
@@ -0,0 +1,18 @@
[package]
name = "paracas-estimate"
version.workspace = true
edition.workspace = true
rust-version.workspace = true
license.workspace = true
homepage.workspace = true
repository.workspace = true
description = "Download size and time estimation for paracas tick data downloader"
[lints]
workspace = true
[dependencies]
paracas-types = { workspace = true }
serde = { workspace = true }
serde_json = { workspace = true }
chrono = { workspace = true }
@@ -0,0 +1,39 @@
{
"categories": {
"forex": {
"avg_compressed_bytes_per_hour": 75000,
"avg_ticks_per_hour": 5000,
"peak_multiplier": 2.5
},
"crypto": {
"avg_compressed_bytes_per_hour": 150000,
"avg_ticks_per_hour": 10000,
"peak_multiplier": 3.0
},
"index": {
"avg_compressed_bytes_per_hour": 50000,
"avg_ticks_per_hour": 3000,
"peak_multiplier": 2.0
},
"commodity": {
"avg_compressed_bytes_per_hour": 40000,
"avg_ticks_per_hour": 2500,
"peak_multiplier": 2.0
},
"stock": {
"avg_compressed_bytes_per_hour": 30000,
"avg_ticks_per_hour": 2000,
"peak_multiplier": 1.5
},
"etf": {
"avg_compressed_bytes_per_hour": 25000,
"avg_ticks_per_hour": 1500,
"peak_multiplier": 1.5
},
"bond": {
"avg_compressed_bytes_per_hour": 20000,
"avg_ticks_per_hour": 1000,
"peak_multiplier": 1.5
}
}
}
+186
View File
@@ -0,0 +1,186 @@
//! Estimate database with historical averages.
use std::collections::HashMap;
use std::sync::OnceLock;
use serde::{Deserialize, Serialize};
/// Embedded JSON data with historical size estimates.
const SIZE_ESTIMATES_JSON: &str = include_str!("../data/size_estimates.json");
/// Static estimate database instance.
static ESTIMATES: OnceLock<EstimateDatabase> = OnceLock::new();
/// Size estimate for a single instrument category.
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
pub struct CategoryEstimate {
/// Category name (e.g., "forex", "crypto").
pub category: String,
/// Average compressed bytes per hour of data.
pub avg_compressed_bytes_per_hour: u64,
/// Average number of ticks per hour.
pub avg_ticks_per_hour: u64,
/// Multiplier for peak trading hours.
pub peak_multiplier: f64,
}
impl CategoryEstimate {
/// Creates a new category estimate.
#[must_use]
pub fn new(
category: impl Into<String>,
avg_compressed_bytes_per_hour: u64,
avg_ticks_per_hour: u64,
peak_multiplier: f64,
) -> Self {
Self {
category: category.into(),
avg_compressed_bytes_per_hour,
avg_ticks_per_hour,
peak_multiplier,
}
}
/// Returns the maximum compressed bytes per hour (at peak).
#[must_use]
pub fn max_compressed_bytes_per_hour(&self) -> u64 {
(self.avg_compressed_bytes_per_hour as f64 * self.peak_multiplier) as u64
}
/// Returns the maximum ticks per hour (at peak).
#[must_use]
pub fn max_ticks_per_hour(&self) -> u64 {
(self.avg_ticks_per_hour as f64 * self.peak_multiplier) as u64
}
}
/// Raw JSON structure for deserialization.
#[derive(Debug, Deserialize)]
struct RawEstimateData {
categories: HashMap<String, RawCategoryEstimate>,
}
/// Raw category estimate from JSON.
#[derive(Debug, Deserialize)]
struct RawCategoryEstimate {
avg_compressed_bytes_per_hour: u64,
avg_ticks_per_hour: u64,
peak_multiplier: f64,
}
/// Database of historical size estimates per instrument category.
#[derive(Debug, Clone)]
pub struct EstimateDatabase {
categories: HashMap<String, CategoryEstimate>,
}
impl EstimateDatabase {
/// Returns the global estimate database instance.
///
/// This lazily initializes the database from embedded JSON on first access.
#[must_use]
pub fn global() -> &'static Self {
ESTIMATES.get_or_init(|| {
Self::from_json(SIZE_ESTIMATES_JSON)
.expect("embedded size_estimates.json should be valid")
})
}
/// Creates an estimate database from JSON string.
///
/// # Errors
///
/// Returns an error if the JSON is invalid.
pub fn from_json(json: &str) -> Result<Self, serde_json::Error> {
let raw: RawEstimateData = serde_json::from_str(json)?;
let categories = raw
.categories
.into_iter()
.map(|(name, raw_est)| {
let estimate = CategoryEstimate::new(
name.clone(),
raw_est.avg_compressed_bytes_per_hour,
raw_est.avg_ticks_per_hour,
raw_est.peak_multiplier,
);
(name, estimate)
})
.collect();
Ok(Self { categories })
}
/// Returns the estimate for a category by name.
#[must_use]
pub fn get(&self, category: &str) -> Option<&CategoryEstimate> {
self.categories.get(category)
}
/// Returns all available categories.
pub fn categories(&self) -> impl Iterator<Item = &str> {
self.categories.keys().map(String::as_str)
}
/// Returns the number of categories in the database.
#[must_use]
pub fn len(&self) -> usize {
self.categories.len()
}
/// Returns true if the database is empty.
#[must_use]
pub fn is_empty(&self) -> bool {
self.categories.is_empty()
}
/// Returns a default estimate for unknown categories.
#[must_use]
pub fn default_estimate() -> CategoryEstimate {
CategoryEstimate::new("unknown", 50000, 3000, 2.0)
}
}
impl Default for EstimateDatabase {
fn default() -> Self {
Self::global().clone()
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_global_database_loads() {
let db = EstimateDatabase::global();
assert!(!db.is_empty());
assert_eq!(db.len(), 7);
}
#[test]
fn test_category_estimates_present() {
let db = EstimateDatabase::global();
let forex = db.get("forex").expect("forex should exist");
assert_eq!(forex.avg_compressed_bytes_per_hour, 75000);
assert_eq!(forex.avg_ticks_per_hour, 5000);
let crypto = db.get("crypto").expect("crypto should exist");
assert_eq!(crypto.avg_compressed_bytes_per_hour, 150000);
assert_eq!(crypto.avg_ticks_per_hour, 10000);
}
#[test]
fn test_peak_calculations() {
let estimate = CategoryEstimate::new("test", 100000, 5000, 2.0);
assert_eq!(estimate.max_compressed_bytes_per_hour(), 200000);
assert_eq!(estimate.max_ticks_per_hour(), 10000);
}
#[test]
fn test_default_estimate() {
let default = EstimateDatabase::default_estimate();
assert_eq!(default.category, "unknown");
assert_eq!(default.avg_compressed_bytes_per_hour, 50000);
}
}
+425
View File
@@ -0,0 +1,425 @@
//! Download estimation logic.
use std::sync::OnceLock;
use std::time::Duration;
use paracas_types::{DateRange, Instrument};
use crate::data::EstimateDatabase;
/// Default download speed assumption in Mbps.
const DEFAULT_DOWNLOAD_SPEED_MBPS: f64 = 10.0;
/// Compression ratio (uncompressed / compressed).
const COMPRESSION_RATIO: f64 = 10.0;
/// Static estimator instance.
static ESTIMATOR: OnceLock<Estimator> = OnceLock::new();
/// Confidence level of the estimate.
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
pub enum EstimateConfidence {
/// High confidence - well-known instrument category with good historical data.
High,
/// Medium confidence - known category but less data or more variability.
Medium,
/// Low confidence - unknown category or limited historical data.
Low,
}
impl EstimateConfidence {
/// Returns the confidence as a string slice.
#[must_use]
pub const fn as_str(&self) -> &'static str {
match self {
Self::High => "high",
Self::Medium => "medium",
Self::Low => "low",
}
}
}
impl std::fmt::Display for EstimateConfidence {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
write!(f, "{}", self.as_str())
}
}
/// Estimated download metrics.
#[derive(Debug, Clone, PartialEq)]
pub struct DownloadEstimate {
/// Total hours of data to download.
pub total_hours: usize,
/// Estimated compressed bytes to download.
pub estimated_compressed_bytes: u64,
/// Estimated uncompressed bytes (compressed * compression ratio).
pub estimated_uncompressed_bytes: u64,
/// Estimated output file size in bytes.
pub estimated_output_bytes: u64,
/// Estimated number of ticks.
pub estimated_ticks: u64,
/// Estimated download duration.
pub estimated_duration: Duration,
/// Confidence level of the estimate.
pub confidence: EstimateConfidence,
}
impl DownloadEstimate {
/// Creates a new download estimate.
#[must_use]
#[allow(clippy::too_many_arguments)]
pub const fn new(
total_hours: usize,
estimated_compressed_bytes: u64,
estimated_uncompressed_bytes: u64,
estimated_output_bytes: u64,
estimated_ticks: u64,
estimated_duration: Duration,
confidence: EstimateConfidence,
) -> Self {
Self {
total_hours,
estimated_compressed_bytes,
estimated_uncompressed_bytes,
estimated_output_bytes,
estimated_ticks,
estimated_duration,
confidence,
}
}
/// Creates an empty estimate (zero hours, zero bytes).
#[must_use]
pub const fn empty() -> Self {
Self {
total_hours: 0,
estimated_compressed_bytes: 0,
estimated_uncompressed_bytes: 0,
estimated_output_bytes: 0,
estimated_ticks: 0,
estimated_duration: Duration::ZERO,
confidence: EstimateConfidence::High,
}
}
}
/// Download size and time estimator.
#[derive(Debug, Clone)]
pub struct Estimator {
/// Assumed download speed in Mbps.
assumed_download_speed_mbps: f64,
}
impl Estimator {
/// Creates a new estimator with the specified download speed.
#[must_use]
pub const fn new(assumed_download_speed_mbps: f64) -> Self {
Self {
assumed_download_speed_mbps,
}
}
/// Returns the global estimator instance with default settings.
#[must_use]
pub fn global() -> &'static Self {
ESTIMATOR.get_or_init(|| Self::new(DEFAULT_DOWNLOAD_SPEED_MBPS))
}
/// Returns the assumed download speed in Mbps.
#[must_use]
pub const fn download_speed_mbps(&self) -> f64 {
self.assumed_download_speed_mbps
}
/// Estimates download metrics for a single instrument and date range.
#[must_use]
pub fn estimate_single(
&self,
instrument: &Instrument,
date_range: &DateRange,
) -> DownloadEstimate {
let total_hours = date_range.total_hours();
let category = instrument.category().as_str();
let db = EstimateDatabase::global();
let (cat_estimate, confidence) = db.get(category).map_or_else(
|| {
(
EstimateDatabase::default_estimate(),
EstimateConfidence::Low,
)
},
|est| (est.clone(), EstimateConfidence::High),
);
self.calculate_estimate(total_hours, &cat_estimate, confidence)
}
/// Estimates download metrics for multiple instruments and date range.
#[must_use]
pub fn estimate_batch(
&self,
instruments: &[&Instrument],
date_range: &DateRange,
) -> DownloadEstimate {
if instruments.is_empty() {
return DownloadEstimate::empty();
}
let total_hours = date_range.total_hours();
let db = EstimateDatabase::global();
let mut total_compressed_bytes: u64 = 0;
let mut total_ticks: u64 = 0;
let mut min_confidence = EstimateConfidence::High;
for instrument in instruments {
let category = instrument.category().as_str();
let (cat_estimate, confidence) = db.get(category).map_or_else(
|| {
(
EstimateDatabase::default_estimate(),
EstimateConfidence::Low,
)
},
|est| (est.clone(), EstimateConfidence::High),
);
total_compressed_bytes +=
cat_estimate.avg_compressed_bytes_per_hour * total_hours as u64;
total_ticks += cat_estimate.avg_ticks_per_hour * total_hours as u64;
// Use the lowest confidence among all instruments
if matches!(confidence, EstimateConfidence::Low) {
min_confidence = EstimateConfidence::Low;
} else if matches!(confidence, EstimateConfidence::Medium)
&& !matches!(min_confidence, EstimateConfidence::Low)
{
min_confidence = EstimateConfidence::Medium;
}
}
let estimated_uncompressed_bytes =
(total_compressed_bytes as f64 * COMPRESSION_RATIO) as u64;
let estimated_output_bytes = estimated_uncompressed_bytes;
let estimated_duration = self.calculate_duration(total_compressed_bytes);
DownloadEstimate::new(
total_hours * instruments.len(),
total_compressed_bytes,
estimated_uncompressed_bytes,
estimated_output_bytes,
total_ticks,
estimated_duration,
min_confidence,
)
}
/// Calculates estimate for a given number of hours and category.
fn calculate_estimate(
&self,
total_hours: usize,
cat_estimate: &crate::data::CategoryEstimate,
confidence: EstimateConfidence,
) -> DownloadEstimate {
let estimated_compressed_bytes =
cat_estimate.avg_compressed_bytes_per_hour * total_hours as u64;
let estimated_uncompressed_bytes =
(estimated_compressed_bytes as f64 * COMPRESSION_RATIO) as u64;
let estimated_output_bytes = estimated_uncompressed_bytes;
let estimated_ticks = cat_estimate.avg_ticks_per_hour * total_hours as u64;
let estimated_duration = self.calculate_duration(estimated_compressed_bytes);
DownloadEstimate::new(
total_hours,
estimated_compressed_bytes,
estimated_uncompressed_bytes,
estimated_output_bytes,
estimated_ticks,
estimated_duration,
confidence,
)
}
/// Calculates download duration based on compressed bytes and speed.
fn calculate_duration(&self, compressed_bytes: u64) -> Duration {
// Convert Mbps to bytes per second
let bytes_per_second = self.assumed_download_speed_mbps * 1_000_000.0 / 8.0;
let seconds = compressed_bytes as f64 / bytes_per_second;
Duration::from_secs_f64(seconds)
}
/// Formats an estimate as a human-readable summary.
#[must_use]
pub fn format_estimate(estimate: &DownloadEstimate) -> String {
format!(
"Download: {} compressed, {} uncompressed\n\
Ticks: ~{}\n\
Duration: {} (at assumed speed)\n\
Confidence: {}",
Self::format_bytes(estimate.estimated_compressed_bytes),
Self::format_bytes(estimate.estimated_uncompressed_bytes),
Self::format_ticks(estimate.estimated_ticks),
Self::format_duration(estimate.estimated_duration),
estimate.confidence,
)
}
/// Formats bytes in human-readable form (e.g., "1.5 GB", "250 MB").
#[must_use]
pub fn format_bytes(bytes: u64) -> String {
const KB: u64 = 1024;
const MB: u64 = 1024 * KB;
const GB: u64 = 1024 * MB;
const TB: u64 = 1024 * GB;
if bytes >= TB {
format!("{:.2} TB", bytes as f64 / TB as f64)
} else if bytes >= GB {
format!("{:.2} GB", bytes as f64 / GB as f64)
} else if bytes >= MB {
format!("{:.2} MB", bytes as f64 / MB as f64)
} else if bytes >= KB {
format!("{:.2} KB", bytes as f64 / KB as f64)
} else {
format!("{} B", bytes)
}
}
/// Formats duration in human-readable form (e.g., "2h 30m", "45m").
#[must_use]
pub fn format_duration(duration: Duration) -> String {
let total_secs = duration.as_secs();
let hours = total_secs / 3600;
let minutes = (total_secs % 3600) / 60;
let seconds = total_secs % 60;
if hours > 0 {
if minutes > 0 {
format!("{}h {}m", hours, minutes)
} else {
format!("{}h", hours)
}
} else if minutes > 0 {
if seconds > 0 && minutes < 10 {
format!("{}m {}s", minutes, seconds)
} else {
format!("{}m", minutes)
}
} else {
format!("{}s", seconds)
}
}
/// Formats tick count in human-readable form.
fn format_ticks(ticks: u64) -> String {
if ticks >= 1_000_000_000 {
format!("{:.2}B", ticks as f64 / 1_000_000_000.0)
} else if ticks >= 1_000_000 {
format!("{:.2}M", ticks as f64 / 1_000_000.0)
} else if ticks >= 1_000 {
format!("{:.2}K", ticks as f64 / 1_000.0)
} else {
format!("{}", ticks)
}
}
}
impl Default for Estimator {
fn default() -> Self {
Self::new(DEFAULT_DOWNLOAD_SPEED_MBPS)
}
}
#[cfg(test)]
mod tests {
use super::*;
use chrono::NaiveDate;
use paracas_types::Category;
fn create_test_instrument(category: Category) -> Instrument {
Instrument::new(
"test",
"Test Instrument",
"Test description",
category,
100_000,
None,
)
}
#[test]
fn test_estimate_single_forex() {
let estimator = Estimator::default();
let instrument = create_test_instrument(Category::Forex);
let start = NaiveDate::from_ymd_opt(2024, 1, 1).unwrap();
let end = NaiveDate::from_ymd_opt(2024, 1, 1).unwrap();
let date_range = DateRange::new(start, end).unwrap();
let estimate = estimator.estimate_single(&instrument, &date_range);
assert_eq!(estimate.total_hours, 24);
assert_eq!(estimate.estimated_compressed_bytes, 75000 * 24);
assert_eq!(estimate.estimated_ticks, 5000 * 24);
assert_eq!(estimate.confidence, EstimateConfidence::High);
}
#[test]
fn test_estimate_batch() {
let estimator = Estimator::default();
let forex = create_test_instrument(Category::Forex);
let crypto = create_test_instrument(Category::Crypto);
let instruments: Vec<&Instrument> = vec![&forex, &crypto];
let start = NaiveDate::from_ymd_opt(2024, 1, 1).unwrap();
let date_range = DateRange::single_day(start);
let estimate = estimator.estimate_batch(&instruments, &date_range);
// 24 hours * 2 instruments
assert_eq!(estimate.total_hours, 48);
// forex: 75000 * 24 + crypto: 150000 * 24
assert_eq!(estimate.estimated_compressed_bytes, (75000 + 150000) * 24);
}
#[test]
fn test_format_bytes() {
assert_eq!(Estimator::format_bytes(500), "500 B");
assert_eq!(Estimator::format_bytes(1536), "1.50 KB");
assert_eq!(Estimator::format_bytes(1_572_864), "1.50 MB");
assert_eq!(Estimator::format_bytes(1_610_612_736), "1.50 GB");
}
#[test]
fn test_format_duration() {
assert_eq!(Estimator::format_duration(Duration::from_secs(30)), "30s");
assert_eq!(
Estimator::format_duration(Duration::from_secs(90)),
"1m 30s"
);
assert_eq!(Estimator::format_duration(Duration::from_secs(3600)), "1h");
assert_eq!(
Estimator::format_duration(Duration::from_secs(5400)),
"1h 30m"
);
}
#[test]
fn test_empty_batch() {
let estimator = Estimator::default();
let instruments: Vec<&Instrument> = vec![];
let start = NaiveDate::from_ymd_opt(2024, 1, 1).unwrap();
let date_range = DateRange::single_day(start);
let estimate = estimator.estimate_batch(&instruments, &date_range);
assert_eq!(estimate.total_hours, 0);
assert_eq!(estimate.estimated_compressed_bytes, 0);
}
#[test]
fn test_estimate_confidence() {
assert_eq!(EstimateConfidence::High.as_str(), "high");
assert_eq!(EstimateConfidence::Medium.as_str(), "medium");
assert_eq!(EstimateConfidence::Low.as_str(), "low");
}
}
+21
View File
@@ -0,0 +1,21 @@
//! Download size and time estimation for paracas tick data downloader.
//!
//! This crate provides utilities for estimating download sizes and times
//! based on historical averages for different instrument categories:
//!
//! - [`EstimateDatabase`] - Database of historical size estimates per category
//! - [`CategoryEstimate`] - Size estimates for a single category
//! - [`Estimator`] - Computes download estimates for instruments and date ranges
//! - [`DownloadEstimate`] - Estimated download metrics
//! - [`EstimateConfidence`] - Confidence level of the estimate
#![doc(issue_tracker_base_url = "https://github.com/factordynamics/paracas/issues/")]
#![cfg_attr(docsrs, feature(doc_cfg, doc_auto_cfg))]
#![warn(missing_docs)]
#![forbid(unsafe_code)]
mod data;
mod estimator;
pub use data::{CategoryEstimate, EstimateDatabase};
pub use estimator::{DownloadEstimate, EstimateConfidence, Estimator};