Cleanup PVF artifact by cache limit and stale time

polkadot/node/core/pvf/src/artifacts.rs

@@ -169,6 +169,15 @@ pub struct Artifacts {
 	inner: HashMap<ArtifactId, ArtifactState>,
 }
 
+/// A condition which we use to cleanup artifacts
+#[derive(Debug)]
+pub enum CleanupBy {
+	// Inactive time after which artefact is deleted
+	Time(Duration),
+	// Max size in bytes. Reaching it older artefacts are deleted
+	Size(u64),
+}
+
 impl Artifacts {
 	#[cfg(test)]
 	pub(crate) fn empty() -> Self {
@@ -251,21 +260,47 @@ impl Artifacts {
 		})
 	}
 
-	/// Remove artifacts older than the given TTL and return id and path of the removed ones.
-	pub fn prune(&mut self, artifact_ttl: Duration) -> Vec<(ArtifactId, PathBuf)> {
-		let now = SystemTime::now();
-
+	/// Remove artifacts older than the given TTL or the total artifacts size limit and return id
+	/// and path of the removed ones.
+	pub fn prune(&mut self, cleanup_by: &CleanupBy) -> Vec<(ArtifactId, PathBuf)> {
 		let mut to_remove = vec![];
-		for (k, v) in self.inner.iter() {
-			if let ArtifactState::Prepared { last_time_needed, ref path, .. } = *v {
-				if now
-					.duration_since(last_time_needed)
-					.map(|age| age > artifact_ttl)
-					.unwrap_or(false)
-				{
-					to_remove.push((k.clone(), path.clone()));
+
+		match cleanup_by {
+			CleanupBy::Time(artifact_ttl) => {
+				let now = SystemTime::now();
+				for (k, v) in self.inner.iter() {
+					if let ArtifactState::Prepared { last_time_needed, ref path, .. } = *v {
+						if now
+							.duration_since(last_time_needed)
+							.map(|age| age > *artifact_ttl)
+							.unwrap_or(false)
+						{
+							to_remove.push((k.clone(), path.clone()));
+						}
+					}
 				}
-			}
+			},
+			CleanupBy::Size(size_limit) => {
+				let mut total_size = 0;
+				let mut artifact_sizes = vec![];
+
+				for (k, v) in self.inner.iter() {
+					if let ArtifactState::Prepared { ref path, last_time_needed, .. } = *v {
+						if let Ok(metadata) = fs::metadata(path) {
+							let size = metadata.len();
+							total_size += size;
+							artifact_sizes.push((k.clone(), path.clone(), size, last_time_needed));
+						}
+					}
+				}
+				artifact_sizes.sort_by_key(|&(_, _, _, last_time_needed)| last_time_needed);
+
+				while total_size > *size_limit {
+					let Some((artifact_id, path, size, _)) = artifact_sizes.pop() else { break };
+					to_remove.push((artifact_id, path));
+					total_size -= size;
+				}
+			},
 		}
 
 		for artifact in &to_remove {

polkadot/node/core/pvf/src/host.rs

@@ -21,7 +21,7 @@
 //! [`ValidationHost`], that allows communication with that event-loop.
 
 use crate::{
-	artifacts::{ArtifactId, ArtifactPathId, ArtifactState, Artifacts},
+	artifacts::{ArtifactId, ArtifactPathId, ArtifactState, Artifacts, CleanupBy},
 	execute::{self, PendingExecutionRequest},
 	metrics::Metrics,
 	prepare, Priority, SecurityStatus, ValidationError, LOG_TARGET,
@@ -293,7 +293,7 @@ pub async fn start(
 	let run_host = async move {
 		run(Inner {
 			cleanup_pulse_interval: Duration::from_secs(3600),
-			artifact_ttl: Duration::from_secs(3600 * 24),
+			cleanup_by: CleanupBy::Time(Duration::from_secs(3600 * 24)),
 			artifacts,
 			to_host_rx,
 			to_prepare_queue_tx,
@@ -337,7 +337,7 @@ impl AwaitingPrepare {
 
 struct Inner {
 	cleanup_pulse_interval: Duration,
-	artifact_ttl: Duration,
+	cleanup_by: CleanupBy,
 	artifacts: Artifacts,
 
 	to_host_rx: mpsc::Receiver<ToHost>,
@@ -359,7 +359,7 @@ struct Fatal;
 async fn run(
 	Inner {
 		cleanup_pulse_interval,
-		artifact_ttl,
+		cleanup_by,
 		mut artifacts,
 		to_host_rx,
 		from_prepare_queue_rx,
@@ -415,7 +415,7 @@ async fn run(
 				break_if_fatal!(handle_cleanup_pulse(
 					&mut to_sweeper_tx,
 					&mut artifacts,
-					artifact_ttl,
+					&cleanup_by,
 				).await);
 			},
 			to_host = to_host_rx.next() => {
@@ -859,9 +859,9 @@ async fn enqueue_prepare_for_execute(
 async fn handle_cleanup_pulse(
 	sweeper_tx: &mut mpsc::Sender<PathBuf>,
 	artifacts: &mut Artifacts,
-	artifact_ttl: Duration,
+	cleanup_by: &CleanupBy,
 ) -> Result<(), Fatal> {
-	let to_remove = artifacts.prune(artifact_ttl);
+	let to_remove = artifacts.prune(cleanup_by);
 	gum::debug!(
 		target: LOG_TARGET,
 		"PVF pruning: {} artifacts reached their end of life",
@@ -1032,7 +1032,7 @@ pub(crate) mod tests {
 
 			let run = run(Inner {
 				cleanup_pulse_interval,
-				artifact_ttl,
+				cleanup_by: CleanupBy::Time(artifact_ttl),
 				artifacts,
 				to_host_rx,
 				to_prepare_queue_tx,