Better long term memory management

CHANGELOG.md

@@ -5,10 +5,15 @@
 ### New features
 
 * Python API `crash_limit` now accepts `"unlimited"` and `"never-restart"` (and `0` is rejected)
+* `hq server start --auto-forget=<duration>` makes the server automatically forget completed jobs
+  that finished more than the given duration ago (e.g. `--auto-forget=7d`)
+* `hq server start --journal-prune-period=<duration>` makes the server periodically prune the journal
 
-## Fixes
+### Fixes
 
 * Fixed crash in auto alloc query in the case of multi-node tasks
+* Job and worker IDs are no longer reused after the server is restarted from a pruned journal
+
 
 ## v0.27.0
 

crates/hyperqueue/src/client/commands/journal/output.rs

@@ -146,6 +146,16 @@ fn format_payload(event: EventPayload) -> serde_json::Value {
                 "type": "server-stop",
             })
         }
+        EventPayload::IdCounters {
+            max_job_id,
+            max_worker_id,
+        } => {
+            json!({
+                "type": "id-counters",
+                "max_job_id": max_job_id,
+                "max_worker_id": max_worker_id,
+            })
+        }
         EventPayload::JobOpen(job_id, job_desc) => {
             json!({
                 "type": "job-open",

crates/hyperqueue/src/client/commands/journal/report.rs

@@ -330,7 +330,8 @@ impl JournalStats {
                 EventPayload::AllocationStarted(_, _)
                 | EventPayload::JobIdle(_)
                 | EventPayload::TaskNotify(_)
-                | EventPayload::ServerStart { .. } => {}
+                | EventPayload::ServerStart { .. }
+                | EventPayload::IdCounters { .. } => {}
                 EventPayload::ServerStop => {
                     for trace in jstats.running_workers.values_mut() {
                         trace.add_value(time, 0);

crates/hyperqueue/src/client/commands/server.rs

@@ -113,6 +113,14 @@ pub struct ServerStartOpts {
     )]
     journal_flush_period: Duration,
 
+    #[arg(
+        long,
+        requires = "journal",
+        value_parser = parse_hms_or_human_time,
+        help = duration_doc!("Periodically prune the journal (remove completed jobs and lost workers).\n\nIf `--auto-forget` is set, only completed jobs that were already forgotten are removed.\n\nIf not set, the journal is never pruned automatically.")
+    )]
+    journal_prune_period: Option<Duration>,
+
     /// The path to an access file
     #[arg(long)]
     access_file: Option<PathBuf>,
@@ -138,6 +146,13 @@ pub struct ServerStartOpts {
         help = duration_doc!("Maximum time the scheduler's placement solve may run per round.")
     )]
     scheduler_time_limit: Duration,
+
+    #[arg(
+        long,
+        value_parser = parse_hms_or_human_time,
+        help = duration_doc!("Automatically forget completed jobs that finished more than this duration ago.")
+    )]
+    auto_forget: Option<Duration>,
 }
 
 #[derive(Parser)]
@@ -209,6 +224,7 @@ async fn start_server(gsettings: &GlobalSettings, opts: ServerStartOpts) -> anyh
         worker_port,
         journal_path: opts.journal,
         journal_flush_period: opts.journal_flush_period,
+        journal_prune_period: opts.journal_prune_period,
         worker_secret_key: access_file
             .as_ref()
             .map(|a| a.worker_key().cloned())
@@ -231,6 +247,7 @@ async fn start_server(gsettings: &GlobalSettings, opts: ServerStartOpts) -> anyh
             }),
         server_uid: access_file.as_ref().map(|a| a.server_uid().to_string()),
         scheduler_mip_time_limit: opts.scheduler_time_limit,
+        auto_forget: opts.auto_forget,
     };
 
     init_hq_server(gsettings, server_cfg).await

crates/hyperqueue/src/client/commands/wait.rs

@@ -212,7 +212,8 @@ pub async fn wait_for_jobs_with_progress(
                     | EventPayload::AllocationFinished(_, _)
                     | EventPayload::ServerStart { .. }
                     | EventPayload::ServerStop
-                    | EventPayload::TaskNotify(..) => {}
+                    | EventPayload::TaskNotify(..)
+                    | EventPayload::IdCounters { .. } => {}
                 },
                 _ => {
                     log::warn!("Unexpected message from server: {:?}", msg);

crates/hyperqueue/src/dashboard/data/timelines/job_timeline.rs

@@ -150,7 +150,8 @@ impl JobTimeline {
                 | EventPayload::AllocationFinished(_, _)
                 | EventPayload::ServerStart { .. }
                 | EventPayload::ServerStop
-                | EventPayload::TaskNotify(_) => {}
+                | EventPayload::TaskNotify(_)
+                | EventPayload::IdCounters { .. } => {}
             }
         }
     }

crates/hyperqueue/src/dashboard/data/timelines/worker_timeline.rs

@@ -88,7 +88,8 @@ impl WorkerTimeline {
                 | EventPayload::AllocationFinished(_, _)
                 | EventPayload::ServerStart { .. }
                 | EventPayload::ServerStop
-                | EventPayload::TaskNotify(_) => {}
+                | EventPayload::TaskNotify(_)
+                | EventPayload::IdCounters { .. } => {}
             }
         }
     }

crates/hyperqueue/src/server/autoclean.rs

@@ -0,0 +1,54 @@
+use crate::server::Senders;
+use crate::server::client::{forget_jobs, prune_journal};
+use crate::server::state::StateRef;
+use chrono::{TimeDelta, Utc};
+use std::time::Duration;
+use tako::JobId;
+use tokio::time::Instant;
+
+pub async fn auto_forget_process(state_ref: StateRef, senders: Senders, max_age: Duration) {
+    let max_age_delta = TimeDelta::from_std(max_age).unwrap_or(TimeDelta::MAX);
+    let check_period = (max_age / 10).clamp(Duration::from_secs(1), Duration::from_secs(3600));
+    let mut interval = tokio::time::interval(check_period);
+    loop {
+        interval.tick().await;
+        let forgotten_jobs = {
+            let mut state = state_ref.get_mut();
+            let now = Utc::now();
+            let job_ids: Vec<JobId> = state
+                .jobs()
+                .filter(|job| {
+                    job.is_terminated()
+                        && job
+                            .completion_date
+                            .is_some_and(|date| now - date >= max_age_delta)
+                })
+                .map(|job| job.job_id)
+                .collect();
+            if job_ids.is_empty() {
+                continue;
+            }
+            log::info!("Auto-forgetting {} job(s)", job_ids.len());
+            forget_jobs(&mut state, &senders, &job_ids)
+        };
+        // Dropping jobs can take a long time, so we do it outside the event loop.
+        let _ = tokio::task::spawn_blocking(move || {
+            drop(forgotten_jobs);
+        })
+        .await;
+    }
+}
+
+pub async fn journal_prune_process(
+    state_ref: StateRef,
+    senders: Senders,
+    period: Duration,
+    keep_completed_jobs: bool,
+) {
+    let mut interval = tokio::time::interval_at(Instant::now() + period, period);
+    loop {
+        interval.tick().await;
+        log::debug!("Pruning journal");
+        prune_journal(&state_ref, &senders, keep_completed_jobs).await;
+    }
+}

crates/hyperqueue/src/server/bootstrap.rs

@@ -17,6 +17,7 @@ use crate::common::serverdir::{
 };
 use crate::server::Senders;
 use crate::server::autoalloc::{QueueId, create_autoalloc_service};
+use crate::server::autoclean::{auto_forget_process, journal_prune_process};
 use crate::server::event::journal::JournalWriter;
 use crate::server::event::journal::start_event_streaming;
 use crate::server::event::streamer::EventStreamer;
@@ -45,11 +46,14 @@ pub struct ServerConfig {
     pub worker_port: Option<u16>,
     pub journal_path: Option<PathBuf>,
     pub journal_flush_period: Duration,
+    pub journal_prune_period: Option<Duration>,
     pub worker_secret_key: Option<Arc<SecretKey>>,
     pub client_secret_key: Option<Arc<SecretKey>>,
     pub server_uid: Option<String>,
     /// See `tako::server::SchedulerConfig::mip_time_limit`.
     pub scheduler_mip_time_limit: Duration,
+    /// Completed jobs older than this duration are automatically forgotten.
+    pub auto_forget: Option<Duration>,
 }
 
 /// This function initializes the HQ server.
@@ -267,6 +271,25 @@ pub async fn initialize_server(
         senders.clone(),
     )));
 
+    let auto_forget_fut: Pin<Box<dyn Future<Output = ()>>> = match server_cfg.auto_forget {
+        Some(max_age) => Box::pin(auto_forget_process(
+            state_ref.clone(),
+            senders.clone(),
+            max_age,
+        )),
+        None => Box::pin(futures::future::pending()),
+    };
+    let journal_prune_fut: Pin<Box<dyn Future<Output = ()>>> = match server_cfg.journal_prune_period
+    {
+        Some(period) => Box::pin(journal_prune_process(
+            state_ref.clone(),
+            senders.clone(),
+            period,
+            server_cfg.auto_forget.is_some(),
+        )),
+        None => Box::pin(futures::future::pending()),
+    };
+
     let senders2 = senders.clone();
     let fut = async move {
         tokio::pin! {
@@ -287,6 +310,8 @@ pub async fn initialize_server(
                 key
             ) => { Ok(()) }
             _ = &mut autoalloc_process => { Ok(()) }
+            _ = auto_forget_fut => { Ok(()) }
+            _ = journal_prune_fut => { Ok(()) }
             r = comm_fut => { r.map_err(|e| e.into()) }
         };
 

crates/hyperqueue/src/server/client/mod.rs

@@ -545,12 +545,24 @@ fn reconstruct_historical_events(
 
 async fn handle_prune_journal(state_ref: &StateRef, senders: &Senders) -> ToClientMessage {
     log::debug!("Client asked for journal prunning");
-    let (live_jobs, live_workers) = {
+    prune_journal(state_ref, senders, false).await;
+    ToClientMessage::Finished
+}
+
+/// Removes completed jobs and lost workers from the journal (if journal is enabled).
+/// If `keep_completed_jobs` is true, completed jobs that are still in the server memory
+/// (i.e. not forgotten yet) are kept in the journal.
+pub(crate) async fn prune_journal(
+    state_ref: &StateRef,
+    senders: &Senders,
+    keep_completed_jobs: bool,
+) {
+    let (live_jobs, live_workers, max_job_id, max_worker_id) = {
         let state = state_ref.get();
         let live_jobs: Set<_> = state
             .jobs()
             .filter_map(|job| {
-                if job.is_terminated() {
+                if job.is_terminated() && !keep_completed_jobs {
                     None
                 } else {
                     Some(job.job_id)
@@ -568,12 +580,20 @@ async fn handle_prune_journal(state_ref: &StateRef, senders: &Senders) -> ToClie
                 }
             })
             .collect();
-        (live_jobs, live_workers)
+        (
+            live_jobs,
+            live_workers,
+            state.last_job_id(),
+            state.max_worker_id(),
+        )
     };
-    if let Some(receiver) = senders.events.prune_journal(live_jobs, live_workers) {
+    if let Some(receiver) =
+        senders
+            .events
+            .prune_journal(live_jobs, live_workers, max_job_id, max_worker_id)
+    {
         let _ = receiver.await;
     }
-    ToClientMessage::Finished
 }
 
 fn handle_worker_stop(
@@ -797,25 +817,22 @@ fn handle_job_forget(
 ) -> (ToClientMessage, Vec<Job>) {
     let mut state = state_ref.get_mut();
     let job_ids: Vec<JobId> = get_job_ids(&state, selector);
-    let mut forgotten: usize = 0;
 
-    let mut forgotten_jobs = Vec::with_capacity(job_ids.len());
-    for &job_id in &job_ids {
-        let can_be_forgotten = state
-            .get_job(job_id)
-            .map(|j| {
-                j.is_terminated() && allowed_statuses.contains(&job_status(&j.make_job_info(false)))
-            })
-            .unwrap_or(false);
-        if can_be_forgotten {
-            if let Some(job) = state.forget_job(job_id) {
-                forgotten_jobs.push(job);
-            }
-            forgotten += 1;
-        }
-    }
-    state.try_release_memory();
-    senders.server_control.try_release_memory();
+    let to_forget: Vec<JobId> = job_ids
+        .iter()
+        .copied()
+        .filter(|&job_id| {
+            state
+                .get_job(job_id)
+                .map(|j| {
+                    j.is_terminated()
+                        && allowed_statuses.contains(&job_status(&j.make_job_info(false)))
+                })
+                .unwrap_or(false)
+        })
+        .collect();
+    let forgotten = to_forget.len();
+    let forgotten_jobs = forget_jobs(&mut state, senders, &to_forget);
 
     let ignored = job_ids.len() - forgotten;
 
@@ -825,6 +842,18 @@ fn handle_job_forget(
     )
 }
 
+/// Removes the given (terminated) jobs from the state and releases memory.
+/// The forgotten jobs are returned, so that the caller can drop them outside the event loop.
+pub(crate) fn forget_jobs(state: &mut State, senders: &Senders, job_ids: &[JobId]) -> Vec<Job> {
+    let forgotten_jobs: Vec<Job> = job_ids
+        .iter()
+        .filter_map(|&job_id| state.forget_job(job_id))
+        .collect();
+    state.try_release_memory();
+    senders.server_control.try_release_memory();
+    forgotten_jobs
+}
+
 fn handle_get_list(state_ref: &StateRef, workers: bool) -> ToClientMessage {
     let state = state_ref.get();
 

crates/hyperqueue/src/server/event/journal/mod.rs

@@ -11,7 +11,7 @@ pub use write::JournalWriter;
 const HQ_JOURNAL_HEADER: &[u8] = b"hqjl0002";
 
 const HQ_JOURNAL_VERSION_MAJOR: u32 = 26;
-const HQ_JOURNAL_VERSION_MINOR: u32 = 0;
+const HQ_JOURNAL_VERSION_MINOR: u32 = 1;
 
 #[derive(Serialize, Deserialize)]
 struct JournalVersion {

crates/hyperqueue/src/server/event/journal/prune.rs

@@ -1,5 +1,7 @@
+use crate::server::event::Event;
 use crate::server::event::journal::{JournalReader, JournalWriter};
 use crate::server::event::payload::EventPayload;
+use chrono::Utc;
 use tako::JobId;
 use tako::{Set, WorkerId};
 
@@ -8,7 +10,16 @@ pub(crate) fn prune_journal(
     writer: &mut JournalWriter,
     live_job_ids: &Set<JobId>,
     live_worker_ids: &Set<WorkerId>,
+    max_job_id: JobId,
+    max_worker_id: WorkerId,
 ) -> crate::Result<()> {
+    writer.store(Event {
+        time: Utc::now(),
+        payload: EventPayload::IdCounters {
+            max_job_id,
+            max_worker_id,
+        },
+    })?;
     for event in reader {
         let mut event = event?;
         let event = match &mut event.payload {
@@ -44,6 +55,8 @@ pub(crate) fn prune_journal(
             | EventPayload::TaskNotify(_)
             | EventPayload::JobIdle(_)
             | EventPayload::ServerStop => Some(event),
+            // A fresh one was written at the beginning
+            EventPayload::IdCounters { .. } => None,
         };
         if let Some(event) = event {
             writer.store(event)?;