checksum poc

polkadot/node/core/pvf/common/src/prepare.rs

@@ -23,6 +23,8 @@ pub struct PrepareStats {
 	pub cpu_time_elapsed: std::time::Duration,
 	/// The observed memory statistics for the preparation job.
 	pub memory_stats: MemoryStats,
+	/// checksum
+	pub checksum: [u8; 32],
 }
 
 /// Helper struct to contain all the memory stats, including `MemoryAllocationStats` and, if

polkadot/node/core/pvf/prepare-worker/Cargo.toml

@@ -7,6 +7,7 @@ edition.workspace = true
 license.workspace = true
 
 [dependencies]
+blake3 = "1.5"
 cfg-if = "1.0"
 futures = "0.3.21"
 gum = { package = "tracing-gum", path = "../../../gum" }

polkadot/node/core/pvf/prepare-worker/src/lib.rs

@@ -257,7 +257,10 @@ pub fn worker_entrypoint(
 								);
 								tokio::fs::write(&temp_artifact_dest, &artifact).await?;
 
-								Ok(PrepareStats { cpu_time_elapsed, memory_stats })
+								let mut gen_hash = blake3::hash(artifact.as_ref());
+								let checksum = gen_hash.finalize();
+
+								Ok(PrepareStats { cpu_time_elapsed, memory_stats, checksum })
 							},
 						}
 					},

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

@@ -111,28 +111,59 @@ const RUNTIME_PREFIX: &str = "wasmtime_";
 const NODE_PREFIX: &str = "polkadot_v";
 const ARTIFACT_PREFIX: &str = concat_const!(RUNTIME_PREFIX, NODE_PREFIX, NODE_VERSION);
 
+pub type ChecksumInner = [u8; 32];
+
+#[derive(Clone, Debug, PartialEq, Eq, PartialOrd, Ord, Hash)]
+pub struct Checksum(ChecksumInner);
+
+impl Checksum {
+	pub fn new() -> Self {
+		Self([0; 32])
+	}
+
+	pub fn update(&mut self, src: &ChecksumInner) {
+		self.0.copy_from_slice(src);
+	}
+}
+
+impl std::fmt::LowerHex for Checksum {
+	fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+		for byte in self.0 {
+			write!(f, "{}", byte)?;
+		}
+		Ok(())
+	}
+}
+
 /// Identifier of an artifact. Encodes a code hash of the PVF and a hash of executor parameter set.
 #[derive(Clone, Debug, PartialEq, Eq, PartialOrd, Ord, Hash)]
 pub struct ArtifactId {
 	pub(crate) code_hash: ValidationCodeHash,
 	pub(crate) executor_params_hash: ExecutorParamsHash,
+	pub(crate) checksum: Checksum,
 }
 
 impl ArtifactId {
 	/// Creates a new artifact ID with the given hash.
 	pub fn new(code_hash: ValidationCodeHash, executor_params_hash: ExecutorParamsHash) -> Self {
-		Self { code_hash, executor_params_hash }
+		Self { code_hash, executor_params_hash, checksum: Checksum::new() }
 	}
 
 	/// Returns an artifact ID that corresponds to the PVF with given executor params.
 	pub fn from_pvf_prep_data(pvf: &PvfPrepData) -> Self {
 		Self::new(pvf.code_hash(), pvf.executor_params().hash())
 	}
 
+	pub fn update_checksum(&mut self, src: &ChecksumInner) {
+		self.checksum.update(src);
+	}
+
 	/// Returns the expected path to this artifact given the root of the cache.
 	pub fn path(&self, cache_path: &Path) -> PathBuf {
-		let file_name =
-			format!("{}_{:#x}_{:#x}", ARTIFACT_PREFIX, self.code_hash, self.executor_params_hash);
+		let file_name = format!(
+			"{}_{:#x}_{:#x}_{:#x}",
+			ARTIFACT_PREFIX, self.code_hash, self.executor_params_hash, self.checksum
+		);
 		cache_path.join(file_name)
 	}
 
@@ -153,7 +184,7 @@ impl ArtifactId {
 		let executor_params_hash =
 			ExecutorParamsHash::from_hash(Hash::from_str(executor_params_hash_str).ok()?);
 
-		Some(Self { code_hash, executor_params_hash })
+		Some(Self::new(code_hash, executor_params_hash))
 	}
 }
 
@@ -371,6 +402,13 @@ impl Artifacts {
 			.is_none());
 	}
 
+	pub(crate) fn remove_artifact(
+		&mut self,
+		artifact_id: &ArtifactId,
+	) -> Option<(ArtifactId, ArtifactState)> {
+		self.inner.remove_entry(&artifact_id)
+	}
+
 	/// Remove and retrieve the artifacts from the table that are older than the supplied
 	/// Time-To-Live.
 	pub fn prune(&mut self, artifact_ttl: Duration) -> Vec<ArtifactId> {

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

@@ -757,9 +757,14 @@ async fn handle_prepare_done(
 		.await?;
 	}
 
-	*state = match result {
-		Ok(prepare_stats) =>
-			ArtifactState::Prepared { last_time_needed: SystemTime::now(), prepare_stats },
+	match result {
+		Ok(prepare_stats) => match artifacts.remove_artifact(&artifact_id) {
+			Some((mut id, _)) => {
+				id.update_checksum(&prepare_stats.checksum);
+				artifacts.insert_prepared(artifact_id, SystemTime::now(), prepare_stats);
+			},
+			None => unreachable!(),
+		},
 		Err(error) => {
 			let last_time_failed = SystemTime::now();
 			let num_failures = *num_failures + 1;
@@ -772,9 +777,9 @@ async fn handle_prepare_done(
 				"artifact preparation failed: {}",
 				error
 			);
-			ArtifactState::FailedToProcess { last_time_failed, num_failures, error }
+			*state = ArtifactState::FailedToProcess { last_time_failed, num_failures, error };
 		},
-	};
+	}
 
 	Ok(())
 }

polkadot/node/core/pvf/src/prepare/worker_intf.rs

@@ -199,7 +199,7 @@ async fn handle_response(
 	artifact_path: PathBuf,
 	preparation_timeout: Duration,
 ) -> Outcome {
-	let PrepareStats { cpu_time_elapsed, memory_stats } = match result.clone() {
+	let PrepareStats { cpu_time_elapsed, memory_stats, .. } = match result.clone() {
 		Ok(result) => result,
 		// Timed out on the child. This should already be logged by the child.
 		Err(PrepareError::TimedOut) => return Outcome::TimedOut,