refactor(state): trace snapshot write phases and strengthen refresh coverage

wakey-control-plane/src/state/store.rs

@@ -4,7 +4,7 @@ use std::time::{Duration, SystemTime, UNIX_EPOCH};
 use anyhow::{Context, Result};
 use sqlx::sqlite::{SqliteConnectOptions, SqliteJournalMode, SqlitePoolOptions, SqliteSynchronous};
 use sqlx::{Row, Sqlite, SqlitePool, Transaction};
-use tracing::{info, warn};
+use tracing::{Instrument, info, warn};
 use uuid::Uuid;
 
 use crate::state::types::{
@@ -33,9 +33,17 @@ impl Store {
     /// Begins a new transaction with an IMMEDIATE write lock.
     ///
     /// Use this for all transactions that perform writes to avoid SQLite upgrade deadlocks.
+    /// Separate spans distinguish pool wait/setup from acquiring SQLite's writer
+    /// transaction. Keep SQLx ownership and rollback-on-drop behavior intact.
     pub async fn begin_write(&self) -> Result<Transaction<'static, Sqlite>> {
-        self.pool
-            .begin_with("BEGIN IMMEDIATE")
+        let connection = self
+            .pool
+            .acquire()
+            .instrument(tracing::info_span!("db_pool_acquire"))
+            .await
+            .context("failed to begin write transaction")?;
+        Transaction::begin(connection, Some("BEGIN IMMEDIATE".into()))
+            .instrument(tracing::info_span!("db_write_transaction_begin"))
             .await
             .context("failed to begin write transaction")
     }

wakey-control-plane/src/state/store/agent_devices.rs

@@ -1,4 +1,5 @@
 use anyhow::{Context, Result};
+use tracing::Instrument;
 use wakey_core::{Device, DeviceEndpoint, DeviceId, EndpointSource};
 
 use super::Store;
@@ -53,47 +54,16 @@ impl Store {
     ) -> Result<usize> {
         let mut tx = self.begin_write().await?;
 
-        let mut incoming_keys = std::collections::HashSet::with_capacity(devices.len());
-        let snapshot_time = now_unix();
-
-        let existing_keys: Vec<String> = sqlx::query_scalar!(
-            "SELECT device_key FROM agent_devices WHERE agent_id = ?1",
-            agent_id
-        )
-        .fetch_all(&mut *tx)
-        .await
-        .context("failed fetching existing keys")?;
-
-        for device in devices {
-            let Some(device_id) = &device.id else {
-                continue;
-            };
-            let device_key = device_key_from_id(device_id);
-            incoming_keys.insert(device_key.clone());
-
-            snapshot::write_device(&mut tx, agent_id, &device_key, device, snapshot_time).await?;
-        }
-
-        for old_key in existing_keys {
-            if !incoming_keys.contains(&old_key) {
-                sqlx::query!(
-                    "DELETE FROM agent_devices WHERE agent_id = ?1 AND device_key = ?2",
-                    agent_id,
-                    old_key
-                )
-                .execute(&mut *tx)
-                .await
-                .context("failed pruning old agent device")?;
-            }
-        }
+        let accepted = snapshot::replace(&mut tx, agent_id, devices).await?;
 
         tx.commit()
+            .instrument(tracing::info_span!("db_snapshot_commit"))
             .await
             .context("failed committing device snapshot")?;
         self.touch_fleet();
 
-        tracing::Span::current().record("accepted", incoming_keys.len());
-        Ok(incoming_keys.len())
+        tracing::Span::current().record("accepted", accepted);
+        Ok(accepted)
     }
 
     /// List all agent devices and their child rows from one consistent read view.

wakey-control-plane/src/state/store/agent_devices/snapshot.rs

@@ -1,4 +1,4 @@
-//! Writes one snapshot device using the caller's transaction. No helper commits.
+//! Replaces snapshot rows using the caller's transaction. No helper commits.
 
 use anyhow::{Context, Result};
 use sqlx::SqliteConnection;
@@ -9,7 +9,51 @@ use super::{
     presence_to_str,
 };
 
-pub(super) async fn write_device(
+/// Includes key lookup, normalization, parent/child writes and missing-device
+/// deletion. This measures replacement work, not pure SQL execution time.
+#[tracing::instrument(name = "db_snapshot_write_rows", skip_all, fields(devices = devices.len()))]
+pub(super) async fn replace(
+    connection: &mut SqliteConnection,
+    agent_id: &str,
+    devices: &[Device],
+) -> Result<usize> {
+    let mut incoming_keys = std::collections::HashSet::with_capacity(devices.len());
+    let snapshot_time = super::now_unix();
+
+    let existing_keys: Vec<String> = sqlx::query_scalar!(
+        "SELECT device_key FROM agent_devices WHERE agent_id = ?1",
+        agent_id
+    )
+    .fetch_all(&mut *connection)
+    .await
+    .context("failed fetching existing keys")?;
+
+    for device in devices {
+        let Some(device_id) = &device.id else {
+            continue;
+        };
+        let device_key = super::device_key_from_id(device_id);
+        incoming_keys.insert(device_key.clone());
+
+        write_device(connection, agent_id, &device_key, device, snapshot_time).await?;
+    }
+
+    for old_key in existing_keys {
+        if !incoming_keys.contains(&old_key) {
+            sqlx::query!(
+                "DELETE FROM agent_devices WHERE agent_id = ?1 AND device_key = ?2",
+                agent_id,
+                old_key
+            )
+            .execute(&mut *connection)
+            .await
+            .context("failed pruning old agent device")?;
+        }
+    }
+    Ok(incoming_keys.len())
+}
+
+async fn write_device(
     connection: &mut SqliteConnection,
     agent_id: &str,
     device_key: &str,