Merge pull request #500 from NillionNetwork/chore/rename-heartbeats

mfontanini · web-flow · commit 99d0d4a0c502 · 2025-12-11T14:24:55.000-08:00
chore: rename `heartbeats` to `heartbeat`
diff --git a/crates/nilcc-agent-models/src/lib.rs b/crates/nilcc-agent-models/src/lib.rs
@@ -193,12 +193,12 @@ pub mod workloads {
             #[validate(regex(path  = DOMAIN_REGEX))]
             pub domain: String,
 
-            pub heartbeats: Option<CreateWorkloadHeartbeats>,
+            pub heartbeat: Option<CreateWorkloadHeartbeat>,
         }
 
         #[derive(Clone, Debug, Serialize, Deserialize, PartialEq)]
         #[serde(rename_all = "camelCase")]
-        pub struct CreateWorkloadHeartbeats {
+        pub struct CreateWorkloadHeartbeat {
             pub measurement_hash_url: String,
         }
 
diff --git a/nilcc-agent-cli/src/main.rs b/nilcc-agent-cli/src/main.rs
@@ -26,7 +26,7 @@ use nilcc_agent_models::system::InstallArtifactVersionRequest;
 use nilcc_agent_models::system::LastUpgrade;
 use nilcc_agent_models::system::UpgradeState;
 use nilcc_agent_models::system::VerifierKey;
-use nilcc_agent_models::workloads::create::CreateWorkloadHeartbeats;
+use nilcc_agent_models::workloads::create::CreateWorkloadHeartbeat;
 use nilcc_agent_models::workloads::restart::RestartWorkloadRequest;
 use nilcc_agent_models::workloads::start::StartWorkloadRequest;
 use nilcc_agent_models::workloads::stop::StopWorkloadRequest;
@@ -442,7 +442,7 @@ fn launch(client: ApiClient, args: LaunchArgs) -> anyhow::Result<()> {
         gpus,
         disk_space_gb,
         domain,
-        heartbeats: measurement_hash_url.map(|measurement_hash_url| CreateWorkloadHeartbeats { measurement_hash_url }),
+        heartbeat: measurement_hash_url.map(|measurement_hash_url| CreateWorkloadHeartbeat { measurement_hash_url }),
     };
     let response: CreateWorkloadResponse = client.post("/api/v1/workloads/create", &request)?;
     let CreateWorkloadResponse { id } = response;
diff --git a/nilcc-agent/migrations/20251211214244_rename_heartbeats.sql b/nilcc-agent/migrations/20251211214244_rename_heartbeats.sql
@@ -0,0 +1,3 @@
+-- Rename `workloads.hearbeats` to `heartbeat`
+
+ALTER TABLE workloads RENAME COLUMN heartbeats TO heartbeat;
diff --git a/nilcc-agent/src/repositories/workload.rs b/nilcc-agent/src/repositories/workload.rs
@@ -37,7 +37,7 @@ pub struct Workload {
     pub domain: String,
     pub last_reported_event: Option<String>,
     #[sqlx(json)]
-    pub heartbeats: Option<WorkloadHeartbeats>,
+    pub heartbeat: Option<WorkloadHeartbeat>,
 }
 
 impl Workload {
@@ -73,7 +73,7 @@ impl fmt::Debug for Workload {
             enabled: running,
             docker_credentials,
             last_reported_event,
-            heartbeats,
+            heartbeat,
         } = self;
         // Hide this one since it can have sensitive data
         let environment_variables: BTreeMap<_, _> = env_vars.keys().map(|key| (key, "...")).collect();
@@ -94,14 +94,14 @@ impl fmt::Debug for Workload {
             .field("running", running)
             .field("docker_credentials", docker_credentials)
             .field("last_reported_event", last_reported_event)
-            .field("heartbeats", heartbeats)
+            .field("heartbeat", heartbeat)
             .finish()
     }
 }
 
 #[serde_as]
 #[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
-pub struct WorkloadHeartbeats {
+pub struct WorkloadHeartbeat {
     #[serde_as(as = "Option<Hex>")]
     pub wallet_public_key: Option<Vec<u8>>,
 
@@ -135,10 +135,10 @@ pub trait WorkloadRepository: Send + Sync {
     async fn set_enabled(&mut self, id: Uuid, value: bool) -> Result<(), WorkloadRepositoryError>;
 
     /// Set the `heartbeats` column for a workload.
-    async fn set_heartbeats(
+    async fn set_heartbeat(
         &mut self,
         id: Uuid,
-        heartbeats: Option<WorkloadHeartbeats>,
+        heartbeats: Option<WorkloadHeartbeat>,
     ) -> Result<(), WorkloadRepositoryError>;
 
     /// Set the `gpus` column for a workload.
@@ -221,7 +221,7 @@ INSERT INTO workloads (
     ports,
     domain,
     last_reported_event,
-    heartbeats,
+    heartbeat,
     enabled,
     created_at
 )
@@ -244,7 +244,7 @@ VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10, $11, $12, $13, $14, $15, $16, $
             domain,
             last_reported_event,
             enabled,
-            heartbeats,
+            heartbeat,
         } = workload;
 
         sqlx::query(query)
@@ -263,7 +263,7 @@ VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10, $11, $12, $13, $14, $15, $16, $
             .bind(sqlx::types::Json(ports))
             .bind(domain)
             .bind(last_reported_event)
-            .bind(sqlx::types::Json(heartbeats))
+            .bind(sqlx::types::Json(heartbeat))
             .bind(enabled)
             .bind(Utc::now())
             .execute(&mut *self.ctx)
@@ -299,13 +299,13 @@ VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10, $11, $12, $13, $14, $15, $16, $
         Ok(())
     }
 
-    async fn set_heartbeats(
+    async fn set_heartbeat(
         &mut self,
         id: Uuid,
-        heartbeats: Option<WorkloadHeartbeats>,
+        heartbeat: Option<WorkloadHeartbeat>,
     ) -> Result<(), WorkloadRepositoryError> {
-        let query = "UPDATE workloads SET heartbeats = ? WHERE id = ?";
-        sqlx::query(query).bind(sqlx::types::Json(heartbeats)).bind(id).execute(&mut *self.ctx).await?;
+        let query = "UPDATE workloads SET heartbeat = ? WHERE id = ?";
+        sqlx::query(query).bind(sqlx::types::Json(heartbeat)).bind(id).execute(&mut *self.ctx).await?;
         Ok(())
     }
 
@@ -368,7 +368,7 @@ mod tests {
             domain: "example.com".into(),
             last_reported_event: None,
             enabled: true,
-            heartbeats: None,
+            heartbeat: None,
         };
         repo.create(&workload).await.expect("failed to insert");
 
@@ -385,15 +385,15 @@ mod tests {
         repo.set_enabled(workload.id, false).await.expect("failed to update");
         assert_eq!(repo.find(workload.id).await.expect("failed to find").enabled, false);
 
-        repo.set_heartbeats(
+        repo.set_heartbeat(
             workload.id,
-            Some(WorkloadHeartbeats { wallet_public_key: Some(vec![1, 2, 3]), measurement_hash_url: "a".into() }),
+            Some(WorkloadHeartbeat { wallet_public_key: Some(vec![1, 2, 3]), measurement_hash_url: "a".into() }),
         )
         .await
         .expect("failed to update");
         assert_eq!(
-            repo.find(workload.id).await.expect("failed to find").heartbeats,
-            Some(WorkloadHeartbeats { wallet_public_key: Some(vec![1, 2, 3]), measurement_hash_url: "a".into() })
+            repo.find(workload.id).await.expect("failed to find").heartbeat,
+            Some(WorkloadHeartbeat { wallet_public_key: Some(vec![1, 2, 3]), measurement_hash_url: "a".into() })
         );
 
         repo.set_gpus(workload.id, &["cc:dd".into()]).await.expect("failed to update");
diff --git a/nilcc-agent/src/resources.rs b/nilcc-agent/src/resources.rs
@@ -318,7 +318,7 @@ mod tests {
             domain: domain.into(),
             last_reported_event: None,
             enabled: true,
-            heartbeats: None,
+            heartbeat: None,
         }
     }
 
diff --git a/nilcc-agent/src/services/vm.rs b/nilcc-agent/src/services/vm.rs
@@ -477,7 +477,7 @@ mod tests {
             domain: "example.com".into(),
             last_reported_event: None,
             enabled: true,
-            heartbeats: None,
+            heartbeat: None,
         };
         let mut builder = Builder::default();
         let base_disk_contents = b"totally a disk";
diff --git a/nilcc-agent/src/services/workload.rs b/nilcc-agent/src/services/workload.rs
@@ -3,7 +3,7 @@ use crate::{
     repositories::{
         artifacts::ArtifactsRepositoryError,
         sqlite::{ProviderError, ProviderMode, RepositoryProvider},
-        workload::{Workload, WorkloadHeartbeats, WorkloadRepositoryError},
+        workload::{Workload, WorkloadHeartbeat, WorkloadRepositoryError},
     },
     resources::{GpuAddress, SystemResources},
     services::{
@@ -227,7 +227,7 @@ impl DefaultWorkloadService {
         request: CreateWorkloadRequest,
         resources: &AvailableResources,
         artifacts_version: String,
-        heartbeats: Option<WorkloadHeartbeats>,
+        heartbeat: Option<WorkloadHeartbeat>,
     ) -> Workload {
         let CreateWorkloadRequest {
             id,
@@ -265,13 +265,13 @@ impl DefaultWorkloadService {
             domain,
             last_reported_event: None,
             enabled: true,
-            heartbeats,
+            heartbeat,
         }
     }
 
     fn workload_key(&self, workload: &Workload) -> anyhow::Result<Option<VerifierKey>> {
         let id = workload.id;
-        let key = match &workload.heartbeats {
+        let key = match &workload.heartbeat {
             Some(config) => {
                 // Lookup the existing key if one is already set. This will only happen during the
                 // transition period where this is still null for some workloads.
@@ -339,12 +339,12 @@ impl WorkloadService for DefaultWorkloadService {
             return Err(InsufficientResources("open ports"));
         }
 
-        let (heartbeats, wallet_key) = match &request.heartbeats {
+        let (heartbeat, wallet_key) = match &request.heartbeat {
             Some(config) => {
                 let wallet_key = self.verifier_keys.next_key().map_err(|_| NotEnoughKeys)?;
                 let wallet_public_key = Some(wallet_key.public_key().to_vec());
                 (
-                    Some(WorkloadHeartbeats {
+                    Some(WorkloadHeartbeat {
                         wallet_public_key,
                         measurement_hash_url: config.measurement_hash_url.clone(),
                     }),
@@ -353,7 +353,7 @@ impl WorkloadService for DefaultWorkloadService {
             }
             None => (None, None),
         };
-        let workload = self.build_workload(request, &resources, artifacts.version.clone(), heartbeats);
+        let workload = self.build_workload(request, &resources, artifacts.version.clone(), heartbeat);
         let id = workload.id;
         info!("Storing workload {id} in database");
         let mut repo = self.repository_provider.workloads(ProviderMode::Transactional).await?;
@@ -433,13 +433,13 @@ impl WorkloadService for DefaultWorkloadService {
             return Ok(());
         }
         // The workload key is going back to the pool
-        let mut heartbeats = workload.heartbeats;
+        let mut heartbeats = workload.heartbeat;
         if let Some(config) = &mut heartbeats {
             config.wallet_public_key = None;
         }
         info!("Disabling workload {id}");
         repo.set_enabled(id, false).await?;
-        repo.set_heartbeats(id, heartbeats).await?;
+        repo.set_heartbeat(id, heartbeats).await?;
         repo.commit().await?;
         self.vm_service.delete_vm(id).await;
         Ok(())
@@ -453,12 +453,12 @@ impl WorkloadService for DefaultWorkloadService {
             return Ok(());
         }
         let key = self.verifier_keys.next_key().map_err(|e| WorkloadLookupError::Internal(e.to_string()))?;
-        if let Some(config) = &mut workload.heartbeats {
+        if let Some(config) = &mut workload.heartbeat {
             config.wallet_public_key = Some(key.public_key().to_vec());
         }
         info!("Starting workload {id} using wallet key {}", hex::encode(key.public_key()));
         repo.set_enabled(id, true).await?;
-        repo.set_heartbeats(id, workload.heartbeats.clone()).await?;
+        repo.set_heartbeat(id, workload.heartbeat.clone()).await?;
         repo.commit().await?;
         self.vm_service
             .create_vm(workload, Some(key))
@@ -490,7 +490,7 @@ mod tests {
         },
     };
     use mockall::predicate::{always, eq};
-    use nilcc_agent_models::workloads::create::CreateWorkloadHeartbeats;
+    use nilcc_agent_models::workloads::create::CreateWorkloadHeartbeat;
     use rstest::rstest;
     use uuid::Uuid;
 
@@ -586,7 +586,7 @@ mod tests {
             domain: "example.com".into(),
             last_reported_event: None,
             enabled: true,
-            heartbeats: None,
+            heartbeat: None,
         }
     }
 
@@ -687,7 +687,7 @@ mod tests {
             gpus: 1,
             disk_space_gb: 1.try_into().unwrap(),
             domain: "example.com".into(),
-            heartbeats: Some(CreateWorkloadHeartbeats { measurement_hash_url: "url".into() }),
+            heartbeat: Some(CreateWorkloadHeartbeat { measurement_hash_url: "url".into() }),
         };
         let expected_key = VerifierKeys::dummy().next_key().unwrap().public_key().to_vec();
         let workload = Workload {
@@ -707,7 +707,7 @@ mod tests {
             domain: request.domain.clone(),
             last_reported_event: None,
             enabled: true,
-            heartbeats: Some(WorkloadHeartbeats {
+            heartbeat: Some(WorkloadHeartbeat {
                 wallet_public_key: Some(expected_key),
                 measurement_hash_url: "url".into(),
             }),

Original file line number	Diff line number	Diff line change
`@@ -193,12 +193,12 @@ pub mod workloads {`
`193`	`193`	`#[validate(regex(path = DOMAIN_REGEX))]`
`194`	`194`	`pub domain: String,`
`195`	`195`
`196`		`- pub heartbeats: Option<CreateWorkloadHeartbeats>,`
	`196`	`+ pub heartbeat: Option<CreateWorkloadHeartbeat>,`
`197`	`197`	`}`
`198`	`198`
`199`	`199`	`#[derive(Clone, Debug, Serialize, Deserialize, PartialEq)]`
`200`	`200`	`#[serde(rename_all = "camelCase")]`
`201`		`- pub struct CreateWorkloadHeartbeats {`
	`201`	`+ pub struct CreateWorkloadHeartbeat {`
`202`	`202`	`pub measurement_hash_url: String,`
`203`	`203`	`}`
`204`	`204`
Original file line number	Diff line number	Diff line change
`@@ -0,0 +1,3 @@`
	`1`	+-- Rename `workloads.hearbeats` to `heartbeat`
	`2`	`+`
	`3`	`+ALTER TABLE workloads RENAME COLUMN heartbeats TO heartbeat;`
Original file line number	Diff line number	Diff line change
`@@ -318,7 +318,7 @@ mod tests {`
`318`	`318`	`domain: domain.into(),`
`319`	`319`	`last_reported_event: None,`
`320`	`320`	`enabled: true,`
`321`		`- heartbeats: None,`
	`321`	`+ heartbeat: None,`
`322`	`322`	`}`
`323`	`323`	`}`
`324`	`324`