Merge "Implement mls-rs-crypto-traits backed by BoringSSL." into main
diff --git a/keystore2/aconfig/flags.aconfig b/keystore2/aconfig/flags.aconfig
index b67bc6c..856b42e 100644
--- a/keystore2/aconfig/flags.aconfig
+++ b/keystore2/aconfig/flags.aconfig
@@ -23,4 +23,12 @@
   description: "Include support for importing keys that were previously software-emulated into KeyMint"
   bug: "283077822"
   is_fixed_read_only: true
+}
+
+flag {
+  name: "database_loop_timeout"
+  namespace: "hardware_backed_security"
+  description: "Abandon Keystore database retry loop after an interval"
+  bug: "319563050"
+  is_fixed_read_only: true
 }
\ No newline at end of file
diff --git a/keystore2/aidl/android/security/apc/IConfirmationCallback.aidl b/keystore2/aidl/android/security/apc/IConfirmationCallback.aidl
index 5b22be0..277b9dd 100644
--- a/keystore2/aidl/android/security/apc/IConfirmationCallback.aidl
+++ b/keystore2/aidl/android/security/apc/IConfirmationCallback.aidl
@@ -27,10 +27,6 @@
     /**
      * This callback gets called by the implementing service when a pending confirmation prompt
      * gets finalized.
-     * @deprecated Android Protected Confirmation had a low adoption rate among Android device
-     *             makers and developers alike. Given the lack of devices supporting the feature,
-     *             it is deprecated. Developers can use auth-bound Keystore keys as a partial
-     *             replacement.
      *
      * @param result
      *  - ResponseCode.OK On success. In this case dataConfirmed must be non null.
diff --git a/keystore2/aidl/android/security/apc/IProtectedConfirmation.aidl b/keystore2/aidl/android/security/apc/IProtectedConfirmation.aidl
index 9f97847..3162224 100644
--- a/keystore2/aidl/android/security/apc/IProtectedConfirmation.aidl
+++ b/keystore2/aidl/android/security/apc/IProtectedConfirmation.aidl
@@ -35,10 +35,6 @@
     /**
      * Present the confirmation prompt. The caller must implement IConfirmationCallback and pass
      * it to this function as listener.
-     * @deprecated Android Protected Confirmation had a low adoption rate among Android device
-     *             makers and developers alike. Given the lack of devices supporting the
-     *             feature, it is deprecated. Developers can use auth-bound Keystore keys
-     *             as a partial replacement.
      *
      * @param listener Must implement IConfirmationCallback. Doubles as session identifier when
      *           passed to cancelPrompt.
@@ -59,11 +55,6 @@
 
     /**
      * Cancel an ongoing prompt.
-     * @deprecated Android Protected Confirmation had a low adoption rate among Android device
-     *             makers and developers alike. Given the lack of devices supporting the 
-     *             feature, it is deprecated. Developers can use auth-bound Keystore keys as
-     *             a partial replacement.
-     *
      *
      * @param listener Must implement IConfirmationCallback, although in this context this binder
      *            token is only used to identify the session that is to be cancelled.
@@ -75,10 +66,6 @@
 
     /**
      * Returns true if the device supports Android Protected Confirmation.
-     * @deprecated Android Protected Confirmation had a low adoption rate among Android device
-     *             makers and developers alike. Given the lack of devices supporting the
-     *             feature, it is deprecated. Developers can use auth-bound Keystore keys
-     *             as a partial replacement.
      */
     boolean isSupported();
 }
diff --git a/keystore2/legacykeystore/lib.rs b/keystore2/legacykeystore/lib.rs
index db3eff6..d407416 100644
--- a/keystore2/legacykeystore/lib.rs
+++ b/keystore2/legacykeystore/lib.rs
@@ -551,19 +551,19 @@
 
 impl ILegacyKeystore for LegacyKeystoreService {
     fn get(&self, alias: &str, uid: i32) -> BinderResult<Vec<u8>> {
-        let _wp = wd::watch_millis("ILegacyKeystore::get", 500);
+        let _wp = wd::watch("ILegacyKeystore::get");
         map_or_log_err(self.legacy_keystore.get(alias, uid), Ok)
     }
     fn put(&self, alias: &str, uid: i32, entry: &[u8]) -> BinderResult<()> {
-        let _wp = wd::watch_millis("ILegacyKeystore::put", 500);
+        let _wp = wd::watch("ILegacyKeystore::put");
         map_or_log_err(self.legacy_keystore.put(alias, uid, entry), Ok)
     }
     fn remove(&self, alias: &str, uid: i32) -> BinderResult<()> {
-        let _wp = wd::watch_millis("ILegacyKeystore::remove", 500);
+        let _wp = wd::watch("ILegacyKeystore::remove");
         map_or_log_err(self.legacy_keystore.remove(alias, uid), Ok)
     }
     fn list(&self, prefix: &str, uid: i32) -> BinderResult<Vec<String>> {
-        let _wp = wd::watch_millis("ILegacyKeystore::list", 500);
+        let _wp = wd::watch("ILegacyKeystore::list");
         map_or_log_err(self.legacy_keystore.list(prefix, uid), Ok)
     }
 }
diff --git a/keystore2/src/apc.rs b/keystore2/src/apc.rs
index fbf9464..bdde5ae 100644
--- a/keystore2/src/apc.rs
+++ b/keystore2/src/apc.rs
@@ -62,16 +62,6 @@
         Error::Rc(ResponseCode::OPERATION_PENDING)
     }
 
-    /// Short hand for `Error::Rc(ResponseCode::CANCELLED)`
-    pub fn cancelled() -> Self {
-        Error::Rc(ResponseCode::CANCELLED)
-    }
-
-    /// Short hand for `Error::Rc(ResponseCode::ABORTED)`
-    pub fn aborted() -> Self {
-        Error::Rc(ResponseCode::ABORTED)
-    }
-
     /// Short hand for `Error::Rc(ResponseCode::IGNORED)`
     pub fn ignored() -> Self {
         Error::Rc(ResponseCode::IGNORED)
@@ -373,11 +363,11 @@
         &self,
         listener: &binder::Strong<dyn IConfirmationCallback>,
     ) -> BinderResult<()> {
-        let _wp = wd::watch_millis("IProtectedConfirmation::cancelPrompt", 500);
+        let _wp = wd::watch("IProtectedConfirmation::cancelPrompt");
         map_or_log_err(self.cancel_prompt(listener), Ok)
     }
     fn isSupported(&self) -> BinderResult<bool> {
-        let _wp = wd::watch_millis("IProtectedConfirmation::isSupported", 500);
+        let _wp = wd::watch("IProtectedConfirmation::isSupported");
         map_or_log_err(Self::is_supported(), Ok)
     }
 }
diff --git a/keystore2/src/authorization.rs b/keystore2/src/authorization.rs
index 243abf1..5810c35 100644
--- a/keystore2/src/authorization.rs
+++ b/keystore2/src/authorization.rs
@@ -271,12 +271,12 @@
 
 impl IKeystoreAuthorization for AuthorizationManager {
     fn addAuthToken(&self, auth_token: &HardwareAuthToken) -> BinderResult<()> {
-        let _wp = wd::watch_millis("IKeystoreAuthorization::addAuthToken", 500);
+        let _wp = wd::watch("IKeystoreAuthorization::addAuthToken");
         map_or_log_err(self.add_auth_token(auth_token), Ok)
     }
 
     fn onDeviceUnlocked(&self, user_id: i32, password: Option<&[u8]>) -> BinderResult<()> {
-        let _wp = wd::watch_millis("IKeystoreAuthorization::onDeviceUnlocked", 500);
+        let _wp = wd::watch("IKeystoreAuthorization::onDeviceUnlocked");
         map_or_log_err(self.on_device_unlocked(user_id, password.map(|pw| pw.into())), Ok)
     }
 
@@ -286,17 +286,17 @@
         unlocking_sids: &[i64],
         weak_unlock_enabled: bool,
     ) -> BinderResult<()> {
-        let _wp = wd::watch_millis("IKeystoreAuthorization::onDeviceLocked", 500);
+        let _wp = wd::watch("IKeystoreAuthorization::onDeviceLocked");
         map_or_log_err(self.on_device_locked(user_id, unlocking_sids, weak_unlock_enabled), Ok)
     }
 
     fn onWeakUnlockMethodsExpired(&self, user_id: i32) -> BinderResult<()> {
-        let _wp = wd::watch_millis("IKeystoreAuthorization::onWeakUnlockMethodsExpired", 500);
+        let _wp = wd::watch("IKeystoreAuthorization::onWeakUnlockMethodsExpired");
         map_or_log_err(self.on_weak_unlock_methods_expired(user_id), Ok)
     }
 
     fn onNonLskfUnlockMethodsExpired(&self, user_id: i32) -> BinderResult<()> {
-        let _wp = wd::watch_millis("IKeystoreAuthorization::onNonLskfUnlockMethodsExpired", 500);
+        let _wp = wd::watch("IKeystoreAuthorization::onNonLskfUnlockMethodsExpired");
         map_or_log_err(self.on_non_lskf_unlock_methods_expired(user_id), Ok)
     }
 
@@ -306,7 +306,7 @@
         secure_user_id: i64,
         auth_token_max_age_millis: i64,
     ) -> binder::Result<AuthorizationTokens> {
-        let _wp = wd::watch_millis("IKeystoreAuthorization::getAuthTokensForCredStore", 500);
+        let _wp = wd::watch("IKeystoreAuthorization::getAuthTokensForCredStore");
         map_or_log_err(
             self.get_auth_tokens_for_credstore(
                 challenge,
diff --git a/keystore2/src/database.rs b/keystore2/src/database.rs
index 0cc982a..50cd3ba 100644
--- a/keystore2/src/database.rs
+++ b/keystore2/src/database.rs
@@ -82,7 +82,7 @@
     types::FromSqlResult,
     types::ToSqlOutput,
     types::{FromSqlError, Value, ValueRef},
-    Connection, OptionalExtension, ToSql, Transaction, TransactionBehavior,
+    Connection, OptionalExtension, ToSql, Transaction,
 };
 
 use std::{
@@ -92,9 +92,55 @@
     time::{Duration, SystemTime},
 };
 
+use TransactionBehavior::Immediate;
+
 #[cfg(test)]
 use tests::random;
 
+/// Wrapper for `rusqlite::TransactionBehavior` which includes information about the transaction
+/// being performed.
+#[derive(Clone, Copy)]
+enum TransactionBehavior {
+    Deferred,
+    Immediate(&'static str),
+}
+
+impl From<TransactionBehavior> for rusqlite::TransactionBehavior {
+    fn from(val: TransactionBehavior) -> Self {
+        match val {
+            TransactionBehavior::Deferred => rusqlite::TransactionBehavior::Deferred,
+            TransactionBehavior::Immediate(_) => rusqlite::TransactionBehavior::Immediate,
+        }
+    }
+}
+
+impl TransactionBehavior {
+    fn name(&self) -> Option<&'static str> {
+        match self {
+            TransactionBehavior::Deferred => None,
+            TransactionBehavior::Immediate(v) => Some(v),
+        }
+    }
+}
+
+/// If the database returns a busy error code, retry after this interval.
+const DB_BUSY_RETRY_INTERVAL: Duration = Duration::from_micros(500);
+/// If the database returns a busy error code, keep retrying for this long.
+const MAX_DB_BUSY_RETRY_PERIOD: Duration = Duration::from_secs(15);
+
+/// Check whether a database lock has timed out.
+fn check_lock_timeout(start: &std::time::Instant, timeout: Duration) -> Result<()> {
+    if keystore2_flags::database_loop_timeout() {
+        let elapsed = start.elapsed();
+        if elapsed >= timeout {
+            error!("Abandon locked DB after {elapsed:?}");
+            return Err(&KsError::Rc(ResponseCode::BACKEND_BUSY))
+                .context(ks_err!("Abandon locked DB after {elapsed:?}",));
+        }
+    }
+    Ok(())
+}
+
 impl_metadata!(
     /// A set of metadata for key entries.
     #[derive(Debug, Default, Eq, PartialEq)]
@@ -382,11 +428,6 @@
     pub fn to_millis_epoch(self) -> i64 {
         self.0
     }
-
-    /// Returns unix epoch time in seconds.
-    pub fn to_secs_epoch(self) -> i64 {
-        self.0 / 1000
-    }
 }
 
 impl ToSql for DateTime {
@@ -523,7 +564,7 @@
     /// This function blocks until an exclusive lock for the given key entry id can
     /// be acquired. It returns a guard object, that represents the lifecycle of the
     /// acquired lock.
-    pub fn get(&self, key_id: i64) -> KeyIdGuard {
+    fn get(&self, key_id: i64) -> KeyIdGuard {
         let mut locked_keys = self.locked_keys.lock().unwrap();
         while locked_keys.contains(&key_id) {
             locked_keys = self.cond_var.wait(locked_keys).unwrap();
@@ -536,7 +577,7 @@
     /// given key id is already taken the function returns None immediately. If a lock
     /// can be acquired this function returns a guard object, that represents the
     /// lifecycle of the acquired lock.
-    pub fn try_get(&self, key_id: i64) -> Option<KeyIdGuard> {
+    fn try_get(&self, key_id: i64) -> Option<KeyIdGuard> {
         let mut locked_keys = self.locked_keys.lock().unwrap();
         if locked_keys.insert(key_id) {
             Some(KeyIdGuard(key_id))
@@ -665,10 +706,6 @@
     pub fn take_cert(&mut self) -> Option<Vec<u8>> {
         self.cert.take()
     }
-    /// Exposes the optional public certificate chain.
-    pub fn cert_chain(&self) -> &Option<Vec<u8>> {
-        &self.cert_chain
-    }
     /// Extracts the optional public certificate_chain.
     pub fn take_cert_chain(&mut self) -> Option<Vec<u8>> {
         self.cert_chain.take()
@@ -677,10 +714,6 @@
     pub fn km_uuid(&self) -> &Uuid {
         &self.km_uuid
     }
-    /// Exposes the key parameters of this key entry.
-    pub fn key_parameters(&self) -> &Vec<KeyParameter> {
-        &self.parameters
-    }
     /// Consumes this key entry and extracts the keyparameters from it.
     pub fn into_key_parameters(self) -> Vec<KeyParameter> {
         self.parameters
@@ -694,10 +727,6 @@
     pub fn pure_cert(&self) -> bool {
         self.pure_cert
     }
-    /// Consumes this key entry and extracts the keyparameters and metadata from it.
-    pub fn into_key_parameters_and_metadata(self) -> (Vec<KeyParameter>, KeyMetaData) {
-        (self.parameters, self.metadata)
-    }
 }
 
 /// Indicates the sub component of a key entry for persistent storage.
@@ -856,13 +885,13 @@
     /// KeystoreDB cannot be used by multiple threads.
     /// Each thread should open their own connection using `thread_local!`.
     pub fn new(db_root: &Path, gc: Option<Arc<Gc>>) -> Result<Self> {
-        let _wp = wd::watch_millis("KeystoreDB::new", 500);
+        let _wp = wd::watch("KeystoreDB::new");
 
         let persistent_path = Self::make_persistent_path(db_root)?;
         let conn = Self::make_connection(&persistent_path)?;
 
         let mut db = Self { conn, gc, perboot: perboot::PERBOOT_DB.clone() };
-        db.with_transaction(TransactionBehavior::Immediate, |tx| {
+        db.with_transaction(Immediate("TX_new"), |tx| {
             versioning::upgrade_database(tx, Self::CURRENT_DB_VERSION, Self::UPGRADERS)
                 .context(ks_err!("KeystoreDB::new: trying to upgrade database."))?;
             Self::init_tables(tx).context("Trying to initialize tables.").no_gc()
@@ -1030,7 +1059,7 @@
                 .context("Failed to attach database persistent.")
             {
                 if Self::is_locked_error(&e) {
-                    std::thread::sleep(std::time::Duration::from_micros(500));
+                    std::thread::sleep(DB_BUSY_RETRY_INTERVAL);
                     continue;
                 } else {
                     return Err(e);
@@ -1091,7 +1120,7 @@
     /// types that map to a table, information about the table's storage is
     /// returned. Requests for storage types that are not DB tables return None.
     pub fn get_storage_stat(&mut self, storage_type: MetricsStorage) -> Result<StorageStats> {
-        let _wp = wd::watch_millis("KeystoreDB::get_storage_stat", 500);
+        let _wp = wd::watch("KeystoreDB::get_storage_stat");
 
         match storage_type {
             MetricsStorage::DATABASE => self.get_total_size(),
@@ -1154,8 +1183,8 @@
         blob_ids_to_delete: &[i64],
         max_blobs: usize,
     ) -> Result<Vec<(i64, Vec<u8>, BlobMetaData)>> {
-        let _wp = wd::watch_millis("KeystoreDB::handle_next_superseded_blob", 500);
-        self.with_transaction(TransactionBehavior::Immediate, |tx| {
+        let _wp = wd::watch("KeystoreDB::handle_next_superseded_blob");
+        self.with_transaction(Immediate("TX_handle_next_superseded_blob"), |tx| {
             // Delete the given blobs.
             for blob_id in blob_ids_to_delete {
                 tx.execute(
@@ -1242,9 +1271,9 @@
     /// Unlike with `mark_unreferenced`, we don't need to purge grants, because only keys that made
     /// it to `KeyLifeCycle::Live` may have grants.
     pub fn cleanup_leftovers(&mut self) -> Result<usize> {
-        let _wp = wd::watch_millis("KeystoreDB::cleanup_leftovers", 500);
+        let _wp = wd::watch("KeystoreDB::cleanup_leftovers");
 
-        self.with_transaction(TransactionBehavior::Immediate, |tx| {
+        self.with_transaction(Immediate("TX_cleanup_leftovers"), |tx| {
             tx.execute(
                 "UPDATE persistent.keyentry SET state = ? WHERE state = ?;",
                 params![KeyLifeCycle::Unreferenced, KeyLifeCycle::Existing],
@@ -1263,9 +1292,9 @@
         alias: &str,
         key_type: KeyType,
     ) -> Result<bool> {
-        let _wp = wd::watch_millis("KeystoreDB::key_exists", 500);
+        let _wp = wd::watch("KeystoreDB::key_exists");
 
-        self.with_transaction(TransactionBehavior::Immediate, |tx| {
+        self.with_transaction(Immediate("TX_key_exists"), |tx| {
             let key_descriptor =
                 KeyDescriptor { domain, nspace, alias: Some(alias.to_string()), blob: None };
             let result = Self::load_key_entry_id(tx, &key_descriptor, key_type);
@@ -1290,9 +1319,9 @@
         blob_metadata: &BlobMetaData,
         key_metadata: &KeyMetaData,
     ) -> Result<KeyEntry> {
-        let _wp = wd::watch_millis("KeystoreDB::store_super_key", 500);
+        let _wp = wd::watch("KeystoreDB::store_super_key");
 
-        self.with_transaction(TransactionBehavior::Immediate, |tx| {
+        self.with_transaction(Immediate("TX_store_super_key"), |tx| {
             let key_id = Self::insert_with_retry(|id| {
                 tx.execute(
                     "INSERT into persistent.keyentry
@@ -1335,9 +1364,9 @@
         key_type: &SuperKeyType,
         user_id: u32,
     ) -> Result<Option<(KeyIdGuard, KeyEntry)>> {
-        let _wp = wd::watch_millis("KeystoreDB::load_super_key", 500);
+        let _wp = wd::watch("KeystoreDB::load_super_key");
 
-        self.with_transaction(TransactionBehavior::Immediate, |tx| {
+        self.with_transaction(Immediate("TX_load_super_key"), |tx| {
             let key_descriptor = KeyDescriptor {
                 domain: Domain::APP,
                 nspace: user_id as i64,
@@ -1361,99 +1390,6 @@
         .context(ks_err!())
     }
 
-    /// Atomically loads a key entry and associated metadata or creates it using the
-    /// callback create_new_key callback. The callback is called during a database
-    /// transaction. This means that implementers should be mindful about using
-    /// blocking operations such as IPC or grabbing mutexes.
-    pub fn get_or_create_key_with<F>(
-        &mut self,
-        domain: Domain,
-        namespace: i64,
-        alias: &str,
-        km_uuid: Uuid,
-        create_new_key: F,
-    ) -> Result<(KeyIdGuard, KeyEntry)>
-    where
-        F: Fn() -> Result<(Vec<u8>, BlobMetaData)>,
-    {
-        let _wp = wd::watch_millis("KeystoreDB::get_or_create_key_with", 500);
-
-        self.with_transaction(TransactionBehavior::Immediate, |tx| {
-            let id = {
-                let mut stmt = tx
-                    .prepare(
-                        "SELECT id FROM persistent.keyentry
-                    WHERE
-                    key_type = ?
-                    AND domain = ?
-                    AND namespace = ?
-                    AND alias = ?
-                    AND state = ?;",
-                    )
-                    .context(ks_err!("Failed to select from keyentry table."))?;
-                let mut rows = stmt
-                    .query(params![KeyType::Super, domain.0, namespace, alias, KeyLifeCycle::Live])
-                    .context(ks_err!("Failed to query from keyentry table."))?;
-
-                db_utils::with_rows_extract_one(&mut rows, |row| {
-                    Ok(match row {
-                        Some(r) => r.get(0).context("Failed to unpack id.")?,
-                        None => None,
-                    })
-                })
-                .context(ks_err!())?
-            };
-
-            let (id, entry) = match id {
-                Some(id) => (
-                    id,
-                    Self::load_key_components(tx, KeyEntryLoadBits::KM, id).context(ks_err!())?,
-                ),
-
-                None => {
-                    let id = Self::insert_with_retry(|id| {
-                        tx.execute(
-                            "INSERT into persistent.keyentry
-                        (id, key_type, domain, namespace, alias, state, km_uuid)
-                        VALUES(?, ?, ?, ?, ?, ?, ?);",
-                            params![
-                                id,
-                                KeyType::Super,
-                                domain.0,
-                                namespace,
-                                alias,
-                                KeyLifeCycle::Live,
-                                km_uuid,
-                            ],
-                        )
-                    })
-                    .context(ks_err!())?;
-
-                    let (blob, metadata) = create_new_key().context(ks_err!())?;
-                    Self::set_blob_internal(
-                        tx,
-                        id,
-                        SubComponentType::KEY_BLOB,
-                        Some(&blob),
-                        Some(&metadata),
-                    )
-                    .context(ks_err!())?;
-                    (
-                        id,
-                        KeyEntry {
-                            id,
-                            key_blob_info: Some((blob, metadata)),
-                            pure_cert: false,
-                            ..Default::default()
-                        },
-                    )
-                }
-            };
-            Ok((KEY_ID_LOCK.get(id), entry)).no_gc()
-        })
-        .context(ks_err!())
-    }
-
     /// Creates a transaction with the given behavior and executes f with the new transaction.
     /// The transaction is committed only if f returns Ok and retried if DatabaseBusy
     /// or DatabaseLocked is encountered.
@@ -1461,12 +1397,28 @@
     where
         F: Fn(&Transaction) -> Result<(bool, T)>,
     {
+        self.with_transaction_timeout(behavior, MAX_DB_BUSY_RETRY_PERIOD, f)
+    }
+    fn with_transaction_timeout<T, F>(
+        &mut self,
+        behavior: TransactionBehavior,
+        timeout: Duration,
+        f: F,
+    ) -> Result<T>
+    where
+        F: Fn(&Transaction) -> Result<(bool, T)>,
+    {
+        let start = std::time::Instant::now();
+        let name = behavior.name();
         loop {
             let result = self
                 .conn
-                .transaction_with_behavior(behavior)
+                .transaction_with_behavior(behavior.into())
                 .context(ks_err!())
-                .and_then(|tx| f(&tx).map(|result| (result, tx)))
+                .and_then(|tx| {
+                    let _wp = name.map(wd::watch);
+                    f(&tx).map(|result| (result, tx))
+                })
                 .and_then(|(result, tx)| {
                     tx.commit().context(ks_err!("Failed to commit transaction."))?;
                     Ok(result)
@@ -1475,7 +1427,8 @@
                 Ok(result) => break Ok(result),
                 Err(e) => {
                     if Self::is_locked_error(&e) {
-                        std::thread::sleep(std::time::Duration::from_micros(500));
+                        check_lock_timeout(&start, timeout)?;
+                        std::thread::sleep(DB_BUSY_RETRY_INTERVAL);
                         continue;
                     } else {
                         return Err(e).context(ks_err!());
@@ -1501,27 +1454,6 @@
         )
     }
 
-    /// Creates a new key entry and allocates a new randomized id for the new key.
-    /// The key id gets associated with a domain and namespace but not with an alias.
-    /// To complete key generation `rebind_alias` should be called after all of the
-    /// key artifacts, i.e., blobs and parameters have been associated with the new
-    /// key id. Finalizing with `rebind_alias` makes the creation of a new key entry
-    /// atomic even if key generation is not.
-    pub fn create_key_entry(
-        &mut self,
-        domain: &Domain,
-        namespace: &i64,
-        key_type: KeyType,
-        km_uuid: &Uuid,
-    ) -> Result<KeyIdGuard> {
-        let _wp = wd::watch_millis("KeystoreDB::create_key_entry", 500);
-
-        self.with_transaction(TransactionBehavior::Immediate, |tx| {
-            Self::create_key_entry_internal(tx, domain, namespace, key_type, km_uuid).no_gc()
-        })
-        .context(ks_err!())
-    }
-
     fn create_key_entry_internal(
         tx: &Transaction,
         domain: &Domain,
@@ -1570,9 +1502,9 @@
         blob: Option<&[u8]>,
         blob_metadata: Option<&BlobMetaData>,
     ) -> Result<()> {
-        let _wp = wd::watch_millis("KeystoreDB::set_blob", 500);
+        let _wp = wd::watch("KeystoreDB::set_blob");
 
-        self.with_transaction(TransactionBehavior::Immediate, |tx| {
+        self.with_transaction(Immediate("TX_set_blob"), |tx| {
             Self::set_blob_internal(tx, key_id.0, sc_type, blob, blob_metadata).need_gc()
         })
         .context(ks_err!())
@@ -1583,9 +1515,9 @@
     /// We use this to insert key blobs into the database which can then be garbage collected
     /// lazily by the key garbage collector.
     pub fn set_deleted_blob(&mut self, blob: &[u8], blob_metadata: &BlobMetaData) -> Result<()> {
-        let _wp = wd::watch_millis("KeystoreDB::set_deleted_blob", 500);
+        let _wp = wd::watch("KeystoreDB::set_deleted_blob");
 
-        self.with_transaction(TransactionBehavior::Immediate, |tx| {
+        self.with_transaction(Immediate("TX_set_deleted_blob"), |tx| {
             Self::set_blob_internal(
                 tx,
                 Self::UNASSIGNED_KEY_ID,
@@ -1644,7 +1576,7 @@
     /// and associates them with the given `key_id`.
     #[cfg(test)]
     fn insert_keyparameter(&mut self, key_id: &KeyIdGuard, params: &[KeyParameter]) -> Result<()> {
-        self.with_transaction(TransactionBehavior::Immediate, |tx| {
+        self.with_transaction(Immediate("TX_insert_keyparameter"), |tx| {
             Self::insert_keyparameter_internal(tx, key_id, params).no_gc()
         })
         .context(ks_err!())
@@ -1677,7 +1609,7 @@
     /// Insert a set of key entry specific metadata into the database.
     #[cfg(test)]
     fn insert_key_metadata(&mut self, key_id: &KeyIdGuard, metadata: &KeyMetaData) -> Result<()> {
-        self.with_transaction(TransactionBehavior::Immediate, |tx| {
+        self.with_transaction(Immediate("TX_insert_key_metadata"), |tx| {
             metadata.store_in_db(key_id.0, tx).no_gc()
         })
         .context(ks_err!())
@@ -1745,7 +1677,7 @@
         caller_uid: u32,
         check_permission: impl Fn(&KeyDescriptor) -> Result<()>,
     ) -> Result<()> {
-        let _wp = wd::watch_millis("KeystoreDB::migrate_key_namespace", 500);
+        let _wp = wd::watch("KeystoreDB::migrate_key_namespace");
 
         let destination = match destination.domain {
             Domain::APP => KeyDescriptor { nspace: caller_uid as i64, ..(*destination).clone() },
@@ -1765,7 +1697,7 @@
             .ok_or(KsError::Rc(ResponseCode::INVALID_ARGUMENT))
             .context(ks_err!("Alias must be specified."))?;
 
-        self.with_transaction(TransactionBehavior::Immediate, |tx| {
+        self.with_transaction(Immediate("TX_migrate_key_namespace"), |tx| {
             // Query the destination location. If there is a key, the migration request fails.
             if tx
                 .query_row(
@@ -1816,7 +1748,7 @@
         metadata: &KeyMetaData,
         km_uuid: &Uuid,
     ) -> Result<KeyIdGuard> {
-        let _wp = wd::watch_millis("KeystoreDB::store_new_key", 500);
+        let _wp = wd::watch("KeystoreDB::store_new_key");
 
         let (alias, domain, namespace) = match key {
             KeyDescriptor { alias: Some(alias), domain: Domain::APP, nspace, blob: None }
@@ -1828,7 +1760,7 @@
                     .context(ks_err!("Need alias and domain must be APP or SELINUX."));
             }
         };
-        self.with_transaction(TransactionBehavior::Immediate, |tx| {
+        self.with_transaction(Immediate("TX_store_new_key"), |tx| {
             let key_id = Self::create_key_entry_internal(tx, &domain, namespace, key_type, km_uuid)
                 .context("Trying to create new key entry.")?;
             let BlobInfo { blob, metadata: blob_metadata, superseded_blob } = *blob_info;
@@ -1895,7 +1827,7 @@
         cert: &[u8],
         km_uuid: &Uuid,
     ) -> Result<KeyIdGuard> {
-        let _wp = wd::watch_millis("KeystoreDB::store_new_certificate", 500);
+        let _wp = wd::watch("KeystoreDB::store_new_certificate");
 
         let (alias, domain, namespace) = match key {
             KeyDescriptor { alias: Some(alias), domain: Domain::APP, nspace, blob: None }
@@ -1907,7 +1839,7 @@
                     .context(ks_err!("Need alias and domain must be APP or SELINUX."));
             }
         };
-        self.with_transaction(TransactionBehavior::Immediate, |tx| {
+        self.with_transaction(Immediate("TX_store_new_certificate"), |tx| {
             let key_id = Self::create_key_entry_internal(tx, &domain, namespace, key_type, km_uuid)
                 .context("Trying to create new key entry.")?;
 
@@ -2175,9 +2107,9 @@
     /// zero, the key also gets marked unreferenced and scheduled for deletion.
     /// Returns Ok(true) if the key was marked unreferenced as a hint to the garbage collector.
     pub fn check_and_update_key_usage_count(&mut self, key_id: i64) -> Result<()> {
-        let _wp = wd::watch_millis("KeystoreDB::check_and_update_key_usage_count", 500);
+        let _wp = wd::watch("KeystoreDB::check_and_update_key_usage_count");
 
-        self.with_transaction(TransactionBehavior::Immediate, |tx| {
+        self.with_transaction(Immediate("TX_check_and_update_key_usage_count"), |tx| {
             let limit: Option<i32> = tx
                 .query_row(
                     "SELECT data FROM persistent.keyparameter WHERE keyentryid = ? AND tag = ?;",
@@ -2223,7 +2155,8 @@
         caller_uid: u32,
         check_permission: impl Fn(&KeyDescriptor, Option<KeyPermSet>) -> Result<()>,
     ) -> Result<(KeyIdGuard, KeyEntry)> {
-        let _wp = wd::watch_millis("KeystoreDB::load_key_entry", 500);
+        let _wp = wd::watch("KeystoreDB::load_key_entry");
+        let start = std::time::Instant::now();
 
         loop {
             match self.load_key_entry_internal(
@@ -2236,7 +2169,8 @@
                 Ok(result) => break Ok(result),
                 Err(e) => {
                     if Self::is_locked_error(&e) {
-                        std::thread::sleep(std::time::Duration::from_micros(500));
+                        check_lock_timeout(&start, MAX_DB_BUSY_RETRY_PERIOD)?;
+                        std::thread::sleep(DB_BUSY_RETRY_INTERVAL);
                         continue;
                     } else {
                         return Err(e).context(ks_err!());
@@ -2351,9 +2285,9 @@
         caller_uid: u32,
         check_permission: impl Fn(&KeyDescriptor, Option<KeyPermSet>) -> Result<()>,
     ) -> Result<()> {
-        let _wp = wd::watch_millis("KeystoreDB::unbind_key", 500);
+        let _wp = wd::watch("KeystoreDB::unbind_key");
 
-        self.with_transaction(TransactionBehavior::Immediate, |tx| {
+        self.with_transaction(Immediate("TX_unbind_key"), |tx| {
             let (key_id, access_key_descriptor, access_vector) =
                 Self::load_access_tuple(tx, key, key_type, caller_uid)
                     .context("Trying to get access tuple.")?;
@@ -2382,12 +2316,12 @@
     /// Delete all artifacts belonging to the namespace given by the domain-namespace tuple.
     /// This leaves all of the blob entries orphaned for subsequent garbage collection.
     pub fn unbind_keys_for_namespace(&mut self, domain: Domain, namespace: i64) -> Result<()> {
-        let _wp = wd::watch_millis("KeystoreDB::unbind_keys_for_namespace", 500);
+        let _wp = wd::watch("KeystoreDB::unbind_keys_for_namespace");
 
         if !(domain == Domain::APP || domain == Domain::SELINUX) {
             return Err(KsError::Rc(ResponseCode::INVALID_ARGUMENT)).context(ks_err!());
         }
-        self.with_transaction(TransactionBehavior::Immediate, |tx| {
+        self.with_transaction(Immediate("TX_unbind_keys_for_namespace"), |tx| {
             tx.execute(
                 "DELETE FROM persistent.keymetadata
                 WHERE keyentryid IN (
@@ -2427,7 +2361,7 @@
     }
 
     fn cleanup_unreferenced(tx: &Transaction) -> Result<()> {
-        let _wp = wd::watch_millis("KeystoreDB::cleanup_unreferenced", 500);
+        let _wp = wd::watch("KeystoreDB::cleanup_unreferenced");
         {
             tx.execute(
                 "DELETE FROM persistent.keymetadata
@@ -2476,9 +2410,9 @@
         user_id: u32,
         keep_non_super_encrypted_keys: bool,
     ) -> Result<()> {
-        let _wp = wd::watch_millis("KeystoreDB::unbind_keys_for_user", 500);
+        let _wp = wd::watch("KeystoreDB::unbind_keys_for_user");
 
-        self.with_transaction(TransactionBehavior::Immediate, |tx| {
+        self.with_transaction(Immediate("TX_unbind_keys_for_user"), |tx| {
             let mut stmt = tx
                 .prepare(&format!(
                     "SELECT id from persistent.keyentry
@@ -2553,9 +2487,9 @@
     /// be unlocked should remain usable when the lock screen is set to Swipe or None, as the device
     /// is always considered "unlocked" in that case.
     pub fn unbind_auth_bound_keys_for_user(&mut self, user_id: u32) -> Result<()> {
-        let _wp = wd::watch_millis("KeystoreDB::unbind_auth_bound_keys_for_user", 500);
+        let _wp = wd::watch("KeystoreDB::unbind_auth_bound_keys_for_user");
 
-        self.with_transaction(TransactionBehavior::Immediate, |tx| {
+        self.with_transaction(Immediate("TX_unbind_auth_bound_keys_for_user"), |tx| {
             let mut stmt = tx
                 .prepare(&format!(
                     "SELECT id from persistent.keyentry
@@ -2648,7 +2582,7 @@
         key_type: KeyType,
         start_past_alias: Option<&str>,
     ) -> Result<Vec<KeyDescriptor>> {
-        let _wp = wd::watch_millis("KeystoreDB::list_past_alias", 500);
+        let _wp = wd::watch("KeystoreDB::list_past_alias");
 
         let query = format!(
             "SELECT DISTINCT alias FROM persistent.keyentry
@@ -2703,7 +2637,7 @@
         namespace: i64,
         key_type: KeyType,
     ) -> Result<usize> {
-        let _wp = wd::watch_millis("KeystoreDB::countKeys", 500);
+        let _wp = wd::watch("KeystoreDB::countKeys");
 
         let num_keys = self.with_transaction(TransactionBehavior::Deferred, |tx| {
             tx.query_row(
@@ -2736,9 +2670,9 @@
         access_vector: KeyPermSet,
         check_permission: impl Fn(&KeyDescriptor, &KeyPermSet) -> Result<()>,
     ) -> Result<KeyDescriptor> {
-        let _wp = wd::watch_millis("KeystoreDB::grant", 500);
+        let _wp = wd::watch("KeystoreDB::grant");
 
-        self.with_transaction(TransactionBehavior::Immediate, |tx| {
+        self.with_transaction(Immediate("TX_grant"), |tx| {
             // Load the key_id and complete the access control tuple.
             // We ignore the access vector here because grants cannot be granted.
             // The access vector returned here expresses the permissions the
@@ -2802,9 +2736,9 @@
         grantee_uid: u32,
         check_permission: impl Fn(&KeyDescriptor) -> Result<()>,
     ) -> Result<()> {
-        let _wp = wd::watch_millis("KeystoreDB::ungrant", 500);
+        let _wp = wd::watch("KeystoreDB::ungrant");
 
-        self.with_transaction(TransactionBehavior::Immediate, |tx| {
+        self.with_transaction(Immediate("TX_ungrant"), |tx| {
             // Load the key_id and complete the access control tuple.
             // We ignore the access vector here because grants cannot be granted.
             let (key_id, access_key_descriptor, _) =
@@ -2868,7 +2802,7 @@
 
     /// Load descriptor of a key by key id
     pub fn load_key_descriptor(&mut self, key_id: i64) -> Result<Option<KeyDescriptor>> {
-        let _wp = wd::watch_millis("KeystoreDB::load_key_descriptor", 500);
+        let _wp = wd::watch("KeystoreDB::load_key_descriptor");
 
         self.with_transaction(TransactionBehavior::Deferred, |tx| {
             tx.query_row(
@@ -2899,9 +2833,9 @@
         user_id: i32,
         secure_user_id: i64,
     ) -> Result<Vec<i64>> {
-        let _wp = wd::watch_millis("KeystoreDB::get_app_uids_affected_by_sid", 500);
+        let _wp = wd::watch("KeystoreDB::get_app_uids_affected_by_sid");
 
-        let key_ids_and_app_uids = self.with_transaction(TransactionBehavior::Immediate, |tx| {
+        let ids = self.with_transaction(Immediate("TX_get_app_uids_affected_by_sid"), |tx| {
             let mut stmt = tx
                 .prepare(&format!(
                     "SELECT id, namespace from persistent.keyentry
@@ -2930,13 +2864,13 @@
             Ok(key_ids_and_app_uids).no_gc()
         })?;
         let mut app_uids_affected_by_sid: HashSet<i64> = Default::default();
-        for (key_id, app_uid) in key_ids_and_app_uids {
+        for (key_id, app_uid) in ids {
             // Read the key parameters for each key in its own transaction. It is OK to ignore
             // an error to get the properties of a particular key since it might have been deleted
             // under our feet after the previous transaction concluded. If the key was deleted
             // then it is no longer applicable if it was auth-bound or not.
             if let Ok(is_key_bound_to_sid) =
-                self.with_transaction(TransactionBehavior::Immediate, |tx| {
+                self.with_transaction(Immediate("TX_get_app_uids_affects_by_sid 2"), |tx| {
                     let params = Self::load_key_parameters(key_id, tx)
                         .context("Failed to load key parameters.")?;
                     // Check if the key is bound to this secure user ID.
@@ -2979,7 +2913,6 @@
     use android_hardware_security_secureclock::aidl::android::hardware::security::secureclock::{
         Timestamp::Timestamp,
     };
-    use rusqlite::TransactionBehavior;
     use std::cell::RefCell;
     use std::collections::BTreeMap;
     use std::fmt::Write;
@@ -2995,7 +2928,7 @@
         let conn = KeystoreDB::make_connection("file::memory:")?;
 
         let mut db = KeystoreDB { conn, gc: None, perboot: Arc::new(perboot::PerbootDB::new()) };
-        db.with_transaction(TransactionBehavior::Immediate, |tx| {
+        db.with_transaction(Immediate("TX_new_test_db"), |tx| {
             KeystoreDB::init_tables(tx).context("Failed to initialize tables.").no_gc()
         })?;
         Ok(db)
@@ -3008,7 +2941,7 @@
         domain: Domain,
         namespace: i64,
     ) -> Result<bool> {
-        db.with_transaction(TransactionBehavior::Immediate, |tx| {
+        db.with_transaction(Immediate("TX_rebind_alias"), |tx| {
             KeystoreDB::rebind_alias(tx, newid, alias, &domain, &namespace, KeyType::Client).no_gc()
         })
         .context(ks_err!())
@@ -3128,12 +3061,24 @@
         db.perboot.get_all_auth_token_entries()
     }
 
+    fn create_key_entry(
+        db: &mut KeystoreDB,
+        domain: &Domain,
+        namespace: &i64,
+        key_type: KeyType,
+        km_uuid: &Uuid,
+    ) -> Result<KeyIdGuard> {
+        db.with_transaction(Immediate("TX_create_key_entry"), |tx| {
+            KeystoreDB::create_key_entry_internal(tx, domain, namespace, key_type, km_uuid).no_gc()
+        })
+    }
+
     #[test]
     fn test_persistence_for_files() -> Result<()> {
         let temp_dir = TempDir::new("persistent_db_test")?;
         let mut db = KeystoreDB::new(temp_dir.path(), None)?;
 
-        db.create_key_entry(&Domain::APP, &100, KeyType::Client, &KEYSTORE_UUID)?;
+        create_key_entry(&mut db, &Domain::APP, &100, KeyType::Client, &KEYSTORE_UUID)?;
         let entries = get_keyentry(&db)?;
         assert_eq!(entries.len(), 1);
 
@@ -3152,8 +3097,8 @@
 
         let mut db = new_test_db()?;
 
-        db.create_key_entry(&Domain::APP, &100, KeyType::Client, &KEYSTORE_UUID)?;
-        db.create_key_entry(&Domain::SELINUX, &101, KeyType::Client, &KEYSTORE_UUID)?;
+        create_key_entry(&mut db, &Domain::APP, &100, KeyType::Client, &KEYSTORE_UUID)?;
+        create_key_entry(&mut db, &Domain::SELINUX, &101, KeyType::Client, &KEYSTORE_UUID)?;
 
         let entries = get_keyentry(&db)?;
         assert_eq!(entries.len(), 2);
@@ -3162,15 +3107,15 @@
 
         // Test that we must pass in a valid Domain.
         check_result_is_error_containing_string(
-            db.create_key_entry(&Domain::GRANT, &102, KeyType::Client, &KEYSTORE_UUID),
+            create_key_entry(&mut db, &Domain::GRANT, &102, KeyType::Client, &KEYSTORE_UUID),
             &format!("Domain {:?} must be either App or SELinux.", Domain::GRANT),
         );
         check_result_is_error_containing_string(
-            db.create_key_entry(&Domain::BLOB, &103, KeyType::Client, &KEYSTORE_UUID),
+            create_key_entry(&mut db, &Domain::BLOB, &103, KeyType::Client, &KEYSTORE_UUID),
             &format!("Domain {:?} must be either App or SELinux.", Domain::BLOB),
         );
         check_result_is_error_containing_string(
-            db.create_key_entry(&Domain::KEY_ID, &104, KeyType::Client, &KEYSTORE_UUID),
+            create_key_entry(&mut db, &Domain::KEY_ID, &104, KeyType::Client, &KEYSTORE_UUID),
             &format!("Domain {:?} must be either App or SELinux.", Domain::KEY_ID),
         );
 
@@ -3186,8 +3131,8 @@
         }
 
         let mut db = new_test_db()?;
-        db.create_key_entry(&Domain::APP, &42, KeyType::Client, &KEYSTORE_UUID)?;
-        db.create_key_entry(&Domain::APP, &42, KeyType::Client, &KEYSTORE_UUID)?;
+        create_key_entry(&mut db, &Domain::APP, &42, KeyType::Client, &KEYSTORE_UUID)?;
+        create_key_entry(&mut db, &Domain::APP, &42, KeyType::Client, &KEYSTORE_UUID)?;
         let entries = get_keyentry(&db)?;
         assert_eq!(entries.len(), 2);
         assert_eq!(
@@ -3438,7 +3383,7 @@
         drop(stmt);
 
         assert_eq!(
-            db.with_transaction(TransactionBehavior::Immediate, |tx| {
+            db.with_transaction(Immediate("TX_test"), |tx| {
                 BlobMetaData::load_from_db(id, tx).no_gc()
             })
             .expect("Should find blob metadata."),
@@ -4113,10 +4058,8 @@
             .unwrap();
         assert_eq!(key_entry, make_bootlevel_test_key_entry_test_vector(key_id_deleted, true));
 
-        db.with_transaction(TransactionBehavior::Immediate, |tx| {
-            KeystoreDB::from_0_to_1(tx).no_gc()
-        })
-        .unwrap();
+        db.with_transaction(Immediate("TX_test"), |tx| KeystoreDB::from_0_to_1(tx).no_gc())
+            .unwrap();
 
         let (_, key_entry) = db
             .load_key_entry(
@@ -4269,12 +4212,12 @@
 
         let _tx1 = db1
             .conn
-            .transaction_with_behavior(TransactionBehavior::Immediate)
+            .transaction_with_behavior(rusqlite::TransactionBehavior::Immediate)
             .expect("Failed to create first transaction.");
 
         let error = db2
             .conn
-            .transaction_with_behavior(TransactionBehavior::Immediate)
+            .transaction_with_behavior(rusqlite::TransactionBehavior::Immediate)
             .context("Transaction begin failed.")
             .expect_err("This should fail.");
         let root_cause = error.root_cause();
@@ -4804,7 +4747,7 @@
         max_usage_count: Option<i32>,
         sids: &[i64],
     ) -> Result<KeyIdGuard> {
-        let key_id = db.create_key_entry(&domain, &namespace, KeyType::Client, &KEYSTORE_UUID)?;
+        let key_id = create_key_entry(db, &domain, &namespace, KeyType::Client, &KEYSTORE_UUID)?;
         let mut blob_metadata = BlobMetaData::new();
         blob_metadata.add(BlobMetaEntry::EncryptedBy(EncryptedBy::Password));
         blob_metadata.add(BlobMetaEntry::Salt(vec![1, 2, 3]));
@@ -4863,7 +4806,7 @@
         alias: &str,
         logical_only: bool,
     ) -> Result<KeyIdGuard> {
-        let key_id = db.create_key_entry(&domain, &namespace, KeyType::Client, &KEYSTORE_UUID)?;
+        let key_id = create_key_entry(db, &domain, &namespace, KeyType::Client, &KEYSTORE_UUID)?;
         let mut blob_metadata = BlobMetaData::new();
         if !logical_only {
             blob_metadata.add(BlobMetaEntry::MaxBootLevel(3));
@@ -4903,7 +4846,7 @@
         super_key_id: i64,
     ) -> Result<KeyIdGuard> {
         let domain = Domain::APP;
-        let key_id = db.create_key_entry(&domain, &namespace, KeyType::Client, &KEYSTORE_UUID)?;
+        let key_id = create_key_entry(db, &domain, &namespace, KeyType::Client, &KEYSTORE_UUID)?;
 
         let mut blob_metadata = BlobMetaData::new();
         blob_metadata.add(BlobMetaEntry::KmUuid(KEYSTORE_UUID));
@@ -5329,7 +5272,7 @@
         let mut db = new_test_db()?;
         let mut working_stats = get_storage_stats_map(&mut db);
 
-        let key_id = db.create_key_entry(&Domain::APP, &42, KeyType::Client, &KEYSTORE_UUID)?;
+        let key_id = create_key_entry(&mut db, &Domain::APP, &42, KeyType::Client, &KEYSTORE_UUID)?;
         assert_storage_increased(
             &mut db,
             vec![
@@ -5580,4 +5523,81 @@
         assert_eq!(third_sid_apps, vec![second_app_id]);
         Ok(())
     }
+
+    #[test]
+    fn test_key_id_guard_immediate() -> Result<()> {
+        if !keystore2_flags::database_loop_timeout() {
+            eprintln!("Skipping test as loop timeout flag disabled");
+            return Ok(());
+        }
+        // Emit logging from test.
+        android_logger::init_once(
+            android_logger::Config::default()
+                .with_tag("keystore_database_tests")
+                .with_max_level(log::LevelFilter::Debug),
+        );
+
+        // Preparation: put a single entry into a test DB.
+        let temp_dir = Arc::new(TempDir::new("key_id_guard_immediate")?);
+        let temp_dir_clone_a = temp_dir.clone();
+        let temp_dir_clone_b = temp_dir.clone();
+        let mut db = KeystoreDB::new(temp_dir.path(), None)?;
+        let key_id = make_test_key_entry(&mut db, Domain::APP, 1, TEST_ALIAS, None)?.0;
+
+        let (a_sender, b_receiver) = std::sync::mpsc::channel();
+        let (b_sender, a_receiver) = std::sync::mpsc::channel();
+
+        // First thread starts an immediate transaction, then waits on a synchronization channel
+        // before trying to get the `KeyIdGuard`.
+        let handle_a = thread::spawn(move || {
+            let temp_dir = temp_dir_clone_a;
+            let mut db = KeystoreDB::new(temp_dir.path(), None).unwrap();
+
+            // Make sure the other thread has initialized its database access before we lock it out.
+            a_receiver.recv().unwrap();
+
+            let _result =
+                db.with_transaction_timeout(Immediate("TX_test"), Duration::from_secs(3), |_tx| {
+                    // Notify the other thread that we're inside the immediate transaction...
+                    a_sender.send(()).unwrap();
+                    // ...then wait to be sure that the other thread has the `KeyIdGuard` before
+                    // this thread also tries to get it.
+                    a_receiver.recv().unwrap();
+
+                    let _guard = KEY_ID_LOCK.get(key_id);
+                    Ok(()).no_gc()
+                });
+        });
+
+        // Second thread gets the `KeyIdGuard`, then waits before trying to perform an immediate
+        // transaction.
+        let handle_b = thread::spawn(move || {
+            let temp_dir = temp_dir_clone_b;
+            let mut db = KeystoreDB::new(temp_dir.path(), None).unwrap();
+            // Notify the other thread that we are initialized (so it can lock the immediate
+            // transaction).
+            b_sender.send(()).unwrap();
+
+            let _guard = KEY_ID_LOCK.get(key_id);
+            // Notify the other thread that we have the `KeyIdGuard`...
+            b_sender.send(()).unwrap();
+            // ...then wait to be sure that the other thread is in the immediate transaction before
+            // this thread also tries to do one.
+            b_receiver.recv().unwrap();
+
+            let result =
+                db.with_transaction_timeout(Immediate("TX_test"), Duration::from_secs(3), |_tx| {
+                    Ok(()).no_gc()
+                });
+            // Expect the attempt to get an immediate transaction to fail, and then this thread will
+            // exit and release the `KeyIdGuard`, allowing the other thread to complete.
+            assert!(result.is_err());
+            check_result_is_error_containing_string(result, "BACKEND_BUSY");
+        });
+
+        let _ = handle_a.join();
+        let _ = handle_b.join();
+
+        Ok(())
+    }
 }
diff --git a/keystore2/src/globals.rs b/keystore2/src/globals.rs
index 7ac1038..2d5c20a 100644
--- a/keystore2/src/globals.rs
+++ b/keystore2/src/globals.rs
@@ -165,7 +165,7 @@
         (
             Box::new(|uuid, blob| {
                 let km_dev = get_keymint_dev_by_uuid(uuid).map(|(dev, _)| dev)?;
-                let _wp = wd::watch_millis("In invalidate key closure: calling deleteKey", 500);
+                let _wp = wd::watch("In invalidate key closure: calling deleteKey");
                 map_km_error(km_dev.deleteKey(blob))
                     .context(ks_err!("Trying to invalidate key blob."))
             }),
@@ -306,7 +306,7 @@
         }
     };
 
-    let wp = wd::watch_millis("In connect_keymint: calling getHardwareInfo()", 500);
+    let wp = wd::watch("In connect_keymint: calling getHardwareInfo()");
     let mut hw_info =
         map_km_error(keymint.getHardwareInfo()).context(ks_err!("Failed to get hardware info."))?;
     drop(wp);
diff --git a/keystore2/src/legacy_importer.rs b/keystore2/src/legacy_importer.rs
index 7dcb98d..f64af0b 100644
--- a/keystore2/src/legacy_importer.rs
+++ b/keystore2/src/legacy_importer.rs
@@ -202,7 +202,7 @@
 
     /// List all aliases for uid in the legacy database.
     pub fn list_uid(&self, domain: Domain, namespace: i64) -> Result<Vec<KeyDescriptor>> {
-        let _wp = wd::watch_millis("LegacyImporter::list_uid", 500);
+        let _wp = wd::watch("LegacyImporter::list_uid");
 
         let uid = match (domain, namespace) {
             (Domain::APP, namespace) => namespace as u32,
@@ -299,7 +299,7 @@
     where
         F: Fn() -> Result<T>,
     {
-        let _wp = wd::watch_millis("LegacyImporter::with_try_import", 500);
+        let _wp = wd::watch("LegacyImporter::with_try_import");
 
         // Access the key and return on success.
         match key_accessor() {
@@ -355,7 +355,7 @@
     where
         F: FnMut() -> Result<Option<T>>,
     {
-        let _wp = wd::watch_millis("LegacyImporter::with_try_import_super_key", 500);
+        let _wp = wd::watch("LegacyImporter::with_try_import_super_key");
 
         match key_accessor() {
             Ok(Some(result)) => return Ok(Some(result)),
@@ -379,7 +379,7 @@
     /// Deletes all keys belonging to the given namespace, importing them into the database
     /// for subsequent garbage collection if necessary.
     pub fn bulk_delete_uid(&self, domain: Domain, nspace: i64) -> Result<()> {
-        let _wp = wd::watch_millis("LegacyImporter::bulk_delete_uid", 500);
+        let _wp = wd::watch("LegacyImporter::bulk_delete_uid");
 
         let uid = match (domain, nspace) {
             (Domain::APP, nspace) => nspace as u32,
@@ -402,7 +402,7 @@
         user_id: u32,
         keep_non_super_encrypted_keys: bool,
     ) -> Result<()> {
-        let _wp = wd::watch_millis("LegacyImporter::bulk_delete_user", 500);
+        let _wp = wd::watch("LegacyImporter::bulk_delete_user");
 
         let result = self.do_serialized(move |importer_state| {
             importer_state
@@ -923,7 +923,7 @@
         blob,
         &[],
         |blob| {
-            let _wd = wd::watch_millis("Calling GetKeyCharacteristics.", 500);
+            let _wd = wd::watch("Calling GetKeyCharacteristics.");
             map_km_error(km_dev.getKeyCharacteristics(blob, &[], &[]))
         },
         |_| Ok(()),
diff --git a/keystore2/src/maintenance.rs b/keystore2/src/maintenance.rs
index 8780e9e..644e7e5 100644
--- a/keystore2/src/maintenance.rs
+++ b/keystore2/src/maintenance.rs
@@ -302,13 +302,13 @@
             user_id,
             password.is_some()
         );
-        let _wp = wd::watch_millis("IKeystoreMaintenance::onUserPasswordChanged", 500);
+        let _wp = wd::watch("IKeystoreMaintenance::onUserPasswordChanged");
         map_or_log_err(Self::on_user_password_changed(user_id, password.map(|pw| pw.into())), Ok)
     }
 
     fn onUserAdded(&self, user_id: i32) -> BinderResult<()> {
         log::info!("onUserAdded(user={user_id})");
-        let _wp = wd::watch_millis("IKeystoreMaintenance::onUserAdded", 500);
+        let _wp = wd::watch("IKeystoreMaintenance::onUserAdded");
         map_or_log_err(self.add_or_remove_user(user_id), Ok)
     }
 
@@ -319,31 +319,31 @@
         allow_existing: bool,
     ) -> BinderResult<()> {
         log::info!("initUserSuperKeys(user={user_id}, allow_existing={allow_existing})");
-        let _wp = wd::watch_millis("IKeystoreMaintenance::initUserSuperKeys", 500);
+        let _wp = wd::watch("IKeystoreMaintenance::initUserSuperKeys");
         map_or_log_err(self.init_user_super_keys(user_id, password.into(), allow_existing), Ok)
     }
 
     fn onUserRemoved(&self, user_id: i32) -> BinderResult<()> {
         log::info!("onUserRemoved(user={user_id})");
-        let _wp = wd::watch_millis("IKeystoreMaintenance::onUserRemoved", 500);
+        let _wp = wd::watch("IKeystoreMaintenance::onUserRemoved");
         map_or_log_err(self.add_or_remove_user(user_id), Ok)
     }
 
     fn onUserLskfRemoved(&self, user_id: i32) -> BinderResult<()> {
         log::info!("onUserLskfRemoved(user={user_id})");
-        let _wp = wd::watch_millis("IKeystoreMaintenance::onUserLskfRemoved", 500);
+        let _wp = wd::watch("IKeystoreMaintenance::onUserLskfRemoved");
         map_or_log_err(Self::on_user_lskf_removed(user_id), Ok)
     }
 
     fn clearNamespace(&self, domain: Domain, nspace: i64) -> BinderResult<()> {
         log::info!("clearNamespace({domain:?}, nspace={nspace})");
-        let _wp = wd::watch_millis("IKeystoreMaintenance::clearNamespace", 500);
+        let _wp = wd::watch("IKeystoreMaintenance::clearNamespace");
         map_or_log_err(self.clear_namespace(domain, nspace), Ok)
     }
 
     fn earlyBootEnded(&self) -> BinderResult<()> {
         log::info!("earlyBootEnded()");
-        let _wp = wd::watch_millis("IKeystoreMaintenance::earlyBootEnded", 500);
+        let _wp = wd::watch("IKeystoreMaintenance::earlyBootEnded");
         map_or_log_err(Self::early_boot_ended(), Ok)
     }
 
@@ -353,13 +353,13 @@
         destination: &KeyDescriptor,
     ) -> BinderResult<()> {
         log::info!("migrateKeyNamespace(src={source:?}, dest={destination:?})");
-        let _wp = wd::watch_millis("IKeystoreMaintenance::migrateKeyNamespace", 500);
+        let _wp = wd::watch("IKeystoreMaintenance::migrateKeyNamespace");
         map_or_log_err(Self::migrate_key_namespace(source, destination), Ok)
     }
 
     fn deleteAllKeys(&self) -> BinderResult<()> {
         log::warn!("deleteAllKeys()");
-        let _wp = wd::watch_millis("IKeystoreMaintenance::deleteAllKeys", 500);
+        let _wp = wd::watch("IKeystoreMaintenance::deleteAllKeys");
         map_or_log_err(Self::delete_all_keys(), Ok)
     }
 
@@ -369,7 +369,7 @@
         secure_user_id: i64,
     ) -> BinderResult<std::vec::Vec<i64>> {
         log::info!("getAppUidsAffectedBySid(secure_user_id={secure_user_id:?})");
-        let _wp = wd::watch_millis("IKeystoreMaintenance::getAppUidsAffectedBySid", 500);
+        let _wp = wd::watch("IKeystoreMaintenance::getAppUidsAffectedBySid");
         map_or_log_err(Self::get_app_uids_affected_by_sid(user_id, secure_user_id), Ok)
     }
 }
diff --git a/keystore2/src/metrics.rs b/keystore2/src/metrics.rs
index cd1cd75..c114c8b 100644
--- a/keystore2/src/metrics.rs
+++ b/keystore2/src/metrics.rs
@@ -51,7 +51,7 @@
 
 impl IKeystoreMetrics for Metrics {
     fn pullMetrics(&self, atom_id: AtomID) -> BinderResult<Vec<KeystoreAtom>> {
-        let _wp = wd::watch_millis("IKeystoreMetrics::pullMetrics", 500);
+        let _wp = wd::watch("IKeystoreMetrics::pullMetrics");
         map_or_log_err(self.pull_metrics(atom_id), Ok)
     }
 }
diff --git a/keystore2/src/operation.rs b/keystore2/src/operation.rs
index 11eaf17..290bc74 100644
--- a/keystore2/src/operation.rs
+++ b/keystore2/src/operation.rs
@@ -286,7 +286,7 @@
         }
         *locked_outcome = Outcome::Pruned;
 
-        let _wp = wd::watch_millis("In Operation::prune: calling abort()", 500);
+        let _wp = wd::watch("In Operation::prune: calling abort()");
 
         // We abort the operation. If there was an error we log it but ignore it.
         if let Err(e) = map_km_error(self.km_op.abort()) {
@@ -362,7 +362,7 @@
             .context(ks_err!("Trying to get auth tokens."))?;
 
         self.update_outcome(&mut outcome, {
-            let _wp = wd::watch_millis("Operation::update_aad: calling updateAad", 500);
+            let _wp = wd::watch("Operation::update_aad: calling updateAad");
             map_km_error(self.km_op.updateAad(aad_input, hat.as_ref(), tst.as_ref()))
         })
         .context(ks_err!("Update failed."))?;
@@ -386,7 +386,7 @@
 
         let output = self
             .update_outcome(&mut outcome, {
-                let _wp = wd::watch_millis("Operation::update: calling update", 500);
+                let _wp = wd::watch("Operation::update: calling update");
                 map_km_error(self.km_op.update(input, hat.as_ref(), tst.as_ref()))
             })
             .context(ks_err!("Update failed."))?;
@@ -416,7 +416,7 @@
 
         let output = self
             .update_outcome(&mut outcome, {
-                let _wp = wd::watch_millis("Operation::finish: calling finish", 500);
+                let _wp = wd::watch("Operation::finish: calling finish");
                 map_km_error(self.km_op.finish(
                     input,
                     signature,
@@ -447,7 +447,7 @@
         *locked_outcome = outcome;
 
         {
-            let _wp = wd::watch_millis("Operation::abort: calling abort", 500);
+            let _wp = wd::watch("Operation::abort: calling abort");
             map_km_error(self.km_op.abort()).context(ks_err!("KeyMint::abort failed."))
         }
     }
@@ -821,7 +821,7 @@
 
 impl IKeystoreOperation for KeystoreOperation {
     fn updateAad(&self, aad_input: &[u8]) -> binder::Result<()> {
-        let _wp = wd::watch_millis("IKeystoreOperation::updateAad", 500);
+        let _wp = wd::watch("IKeystoreOperation::updateAad");
         map_or_log_err(
             self.with_locked_operation(
                 |op| op.update_aad(aad_input).context(ks_err!("KeystoreOperation::updateAad")),
@@ -832,7 +832,7 @@
     }
 
     fn update(&self, input: &[u8]) -> binder::Result<Option<Vec<u8>>> {
-        let _wp = wd::watch_millis("IKeystoreOperation::update", 500);
+        let _wp = wd::watch("IKeystoreOperation::update");
         map_or_log_err(
             self.with_locked_operation(
                 |op| op.update(input).context(ks_err!("KeystoreOperation::update")),
@@ -846,7 +846,7 @@
         input: Option<&[u8]>,
         signature: Option<&[u8]>,
     ) -> binder::Result<Option<Vec<u8>>> {
-        let _wp = wd::watch_millis("IKeystoreOperation::finish", 500);
+        let _wp = wd::watch("IKeystoreOperation::finish");
         map_or_log_err(
             self.with_locked_operation(
                 |op| op.finish(input, signature).context(ks_err!("KeystoreOperation::finish")),
@@ -857,7 +857,7 @@
     }
 
     fn abort(&self) -> binder::Result<()> {
-        let _wp = wd::watch_millis("IKeystoreOperation::abort", 500);
+        let _wp = wd::watch("IKeystoreOperation::abort");
         map_err_with(
             self.with_locked_operation(
                 |op| op.abort(Outcome::Abort).context(ks_err!("KeystoreOperation::abort")),
diff --git a/keystore2/src/raw_device.rs b/keystore2/src/raw_device.rs
index 44d805c..a8a88d2 100644
--- a/keystore2/src/raw_device.rs
+++ b/keystore2/src/raw_device.rs
@@ -211,13 +211,10 @@
                         KeyBlob::NonSensitive(key_blob_vec),
                         |key_blob| {
                             map_km_error({
-                                let _wp = wd::watch_millis(
-                                    concat!(
-                                        "In KeyMintDevice::lookup_or_generate_key: ",
-                                        "calling getKeyCharacteristics."
-                                    ),
-                                    500,
-                                );
+                                let _wp = wd::watch(concat!(
+                                    "In KeyMintDevice::lookup_or_generate_key: ",
+                                    "calling getKeyCharacteristics."
+                                ));
                                 self.km_dev.getKeyCharacteristics(key_blob, &[], &[])
                             })
                         },
@@ -308,7 +305,7 @@
         let (begin_result, _) = self
             .upgrade_keyblob_if_required_with(db, key_id_guard, key_blob, |blob| {
                 map_km_error({
-                    let _wp = wd::watch_millis("In use_key_in_one_step: calling: begin", 500);
+                    let _wp = wd::watch("In use_key_in_one_step: calling: begin");
                     self.km_dev.begin(purpose, blob, operation_parameters, auth_token)
                 })
             })
@@ -316,7 +313,7 @@
         let operation: Strong<dyn IKeyMintOperation> =
             begin_result.operation.ok_or_else(Error::sys).context(ks_err!("Operation missing"))?;
         map_km_error({
-            let _wp = wd::watch_millis("In use_key_in_one_step: calling: finish", 500);
+            let _wp = wd::watch("In use_key_in_one_step: calling: finish");
             operation.finish(Some(input), None, None, None, None)
         })
         .context(ks_err!("Failed to finish operation."))
diff --git a/keystore2/src/remote_provisioning.rs b/keystore2/src/remote_provisioning.rs
index 0ef8c95..cda93b3 100644
--- a/keystore2/src/remote_provisioning.rs
+++ b/keystore2/src/remote_provisioning.rs
@@ -31,7 +31,6 @@
 use anyhow::{Context, Result};
 use keystore2_crypto::parse_subject_from_certificate;
 
-use crate::database::Uuid;
 use crate::error::wrapped_rkpd_error_to_ks_error;
 use crate::globals::get_remotely_provisioned_component_name;
 use crate::ks_err;
@@ -44,18 +43,12 @@
 #[derive(Default)]
 pub struct RemProvState {
     security_level: SecurityLevel,
-    km_uuid: Uuid,
 }
 
 impl RemProvState {
     /// Creates a RemProvState struct.
-    pub fn new(security_level: SecurityLevel, km_uuid: Uuid) -> Self {
-        Self { security_level, km_uuid }
-    }
-
-    /// Returns the uuid for the KM instance attached to this RemProvState struct.
-    pub fn get_uuid(&self) -> Uuid {
-        self.km_uuid
+    pub fn new(security_level: SecurityLevel) -> Self {
+        Self { security_level }
     }
 
     fn is_rkp_only(&self) -> bool {
@@ -136,6 +129,6 @@
     // by the calling function and allow for natural fallback to the factory key.
     let rpc_name = get_remotely_provisioned_component_name(security_level)
         .context(ks_err!("Trying to get IRPC name."))?;
-    let _wd = wd::watch_millis("Calling get_rkpd_attestation_key()", 500);
+    let _wd = wd::watch("Calling get_rkpd_attestation_key()");
     rkpd_client::get_rkpd_attestation_key(&rpc_name, caller_uid)
 }
diff --git a/keystore2/src/security_level.rs b/keystore2/src/security_level.rs
index 5f9745f..59f2315 100644
--- a/keystore2/src/security_level.rs
+++ b/keystore2/src/security_level.rs
@@ -99,7 +99,7 @@
                 hw_info,
                 km_uuid,
                 operation_db: OperationDb::new(),
-                rem_prov_state: RemProvState::new(security_level, km_uuid),
+                rem_prov_state: RemProvState::new(security_level),
                 id_rotation_state,
             },
             BinderFeatures { set_requesting_sid: true, ..BinderFeatures::default() },
@@ -112,6 +112,13 @@
         wd::watch_millis_with(id, millis, move || format!("SecurityLevel {:?}", sec_level))
     }
 
+    fn watch(&self, id: &'static str) -> Option<wd::WatchPoint> {
+        let sec_level = self.security_level;
+        wd::watch_millis_with(id, wd::DEFAULT_TIMEOUT_MS, move || {
+            format!("SecurityLevel {:?}", sec_level)
+        })
+    }
+
     fn store_new_key(
         &self,
         key: KeyDescriptor,
@@ -323,10 +330,8 @@
                 operation_parameters,
                 |blob| loop {
                     match map_km_error({
-                        let _wp = self.watch_millis(
-                            "In KeystoreSecurityLevel::create_operation: calling begin",
-                            500,
-                        );
+                        let _wp =
+                            self.watch("In KeystoreSecurityLevel::create_operation: calling begin");
                         self.keymint.begin(
                             purpose,
                             blob,
@@ -440,10 +445,8 @@
         // If there is an attestation challenge we need to get an application id.
         if params.iter().any(|kp| kp.tag == Tag::ATTESTATION_CHALLENGE) {
             let aaid = {
-                let _wp = self.watch_millis(
-                    "In KeystoreSecurityLevel::add_required_parameters calling: get_aaid",
-                    500,
-                );
+                let _wp = self
+                    .watch("In KeystoreSecurityLevel::add_required_parameters calling: get_aaid");
                 keystore2_aaid::get_aaid(uid)
                     .map_err(|e| anyhow!(ks_err!("get_aaid returned status {}.", e)))
             }?;
@@ -675,8 +678,7 @@
 
         let km_dev = &self.keymint;
         let creation_result = map_km_error({
-            let _wp =
-                self.watch_millis("In KeystoreSecurityLevel::import_key: calling importKey.", 500);
+            let _wp = self.watch("In KeystoreSecurityLevel::import_key: calling importKey.");
             km_dev.importKey(&params, format, key_data, None /* attestKey */)
         })
         .context(ks_err!("Trying to call importKey"))?;
@@ -789,9 +791,8 @@
                 wrapping_blob_metadata.km_uuid().copied(),
                 &[],
                 |wrapping_blob| {
-                    let _wp = self.watch_millis(
+                    let _wp = self.watch(
                         "In KeystoreSecurityLevel::import_wrapped_key: calling importWrappedKey.",
-                        500,
                     );
                     let creation_result = map_km_error(self.keymint.importWrappedKey(
                         wrapped_data,
@@ -897,7 +898,7 @@
             params,
             f,
             |upgraded_blob| {
-                let _wp = wd::watch_millis("Calling store_rkpd_attestation_key()", 500);
+                let _wp = wd::watch("Calling store_rkpd_attestation_key()");
                 if let Err(e) = store_rkpd_attestation_key(&rpc_name, key_blob, upgraded_blob) {
                     Err(wrapped_rkpd_error_to_ks_error(&e)).context(format!("{e:?}"))
                 } else {
@@ -928,13 +929,10 @@
 
         let km_dev = &self.keymint;
         let res = {
-            let _wp = self.watch_millis(
-                concat!(
-                    "In IKeystoreSecurityLevel::convert_storage_key_to_ephemeral: ",
-                    "calling convertStorageKeyToEphemeral (1)"
-                ),
-                500,
-            );
+            let _wp = self.watch(concat!(
+                "In IKeystoreSecurityLevel::convert_storage_key_to_ephemeral: ",
+                "calling convertStorageKeyToEphemeral (1)"
+            ));
             map_km_error(km_dev.convertStorageKeyToEphemeral(key_blob))
         };
         match res {
@@ -943,17 +941,13 @@
             }
             Err(error::Error::Km(ErrorCode::KEY_REQUIRES_UPGRADE)) => {
                 let upgraded_blob = {
-                    let _wp = self.watch_millis(
-                        "In convert_storage_key_to_ephemeral: calling upgradeKey",
-                        500,
-                    );
+                    let _wp = self.watch("In convert_storage_key_to_ephemeral: calling upgradeKey");
                     map_km_error(km_dev.upgradeKey(key_blob, &[]))
                 }
                 .context(ks_err!("Failed to upgrade key blob."))?;
                 let ephemeral_key = {
-                    let _wp = self.watch_millis(
+                    let _wp = self.watch(
                         "In convert_storage_key_to_ephemeral: calling convertStorageKeyToEphemeral (2)",
-                        500,
                     );
                     map_km_error(km_dev.convertStorageKeyToEphemeral(&upgraded_blob))
                 }
@@ -986,8 +980,7 @@
 
         let km_dev = &self.keymint;
         {
-            let _wp =
-                self.watch_millis("In KeystoreSecuritylevel::delete_key: calling deleteKey", 500);
+            let _wp = self.watch("In KeystoreSecuritylevel::delete_key: calling deleteKey");
             map_km_error(km_dev.deleteKey(key_blob)).context(ks_err!("keymint device deleteKey"))
         }
     }
@@ -1002,7 +995,7 @@
         operation_parameters: &[KeyParameter],
         forced: bool,
     ) -> binder::Result<CreateOperationResponse> {
-        let _wp = self.watch_millis("IKeystoreSecurityLevel::createOperation", 500);
+        let _wp = self.watch("IKeystoreSecurityLevel::createOperation");
         map_or_log_err(self.create_operation(key, operation_parameters, forced), Ok)
     }
     fn generateKey(
@@ -1029,7 +1022,7 @@
         flags: i32,
         key_data: &[u8],
     ) -> binder::Result<KeyMetadata> {
-        let _wp = self.watch_millis("IKeystoreSecurityLevel::importKey", 500);
+        let _wp = self.watch("IKeystoreSecurityLevel::importKey");
         let result = self.import_key(key, attestation_key, params, flags, key_data);
         log_key_creation_event_stats(self.security_level, params, &result);
         log_key_imported(key, ThreadState::get_calling_uid(), result.is_ok());
@@ -1043,7 +1036,7 @@
         params: &[KeyParameter],
         authenticators: &[AuthenticatorSpec],
     ) -> binder::Result<KeyMetadata> {
-        let _wp = self.watch_millis("IKeystoreSecurityLevel::importWrappedKey", 500);
+        let _wp = self.watch("IKeystoreSecurityLevel::importWrappedKey");
         let result =
             self.import_wrapped_key(key, wrapping_key, masking_key, params, authenticators);
         log_key_creation_event_stats(self.security_level, params, &result);
@@ -1054,11 +1047,11 @@
         &self,
         storage_key: &KeyDescriptor,
     ) -> binder::Result<EphemeralStorageKeyResponse> {
-        let _wp = self.watch_millis("IKeystoreSecurityLevel::convertStorageKeyToEphemeral", 500);
+        let _wp = self.watch("IKeystoreSecurityLevel::convertStorageKeyToEphemeral");
         map_or_log_err(self.convert_storage_key_to_ephemeral(storage_key), Ok)
     }
     fn deleteKey(&self, key: &KeyDescriptor) -> binder::Result<()> {
-        let _wp = self.watch_millis("IKeystoreSecurityLevel::deleteKey", 500);
+        let _wp = self.watch("IKeystoreSecurityLevel::deleteKey");
         let result = self.delete_key(key);
         log_key_deleted(key, ThreadState::get_calling_uid(), result.is_ok());
         map_or_log_err(result, Ok)
@@ -1130,7 +1123,7 @@
             |new_blob| {
                 // This handler is only executed if a key upgrade was performed.
                 key_upgraded = true;
-                let _wp = wd::watch_millis("Calling store_rkpd_attestation_key()", 500);
+                let _wp = wd::watch("Calling store_rkpd_attestation_key()");
                 store_rkpd_attestation_key(&rpc_name, &key.keyBlob, new_blob).unwrap();
                 Ok(())
             },
diff --git a/keystore2/src/service.rs b/keystore2/src/service.rs
index 1459254..1eab8a8 100644
--- a/keystore2/src/service.rs
+++ b/keystore2/src/service.rs
@@ -387,7 +387,7 @@
         map_or_log_err(self.get_security_level(security_level), Ok)
     }
     fn getKeyEntry(&self, key: &KeyDescriptor) -> binder::Result<KeyEntryResponse> {
-        let _wp = wd::watch_millis("IKeystoreService::get_key_entry", 500);
+        let _wp = wd::watch("IKeystoreService::get_key_entry");
         map_or_log_err(self.get_key_entry(key), Ok)
     }
     fn updateSubcomponent(
@@ -396,15 +396,15 @@
         public_cert: Option<&[u8]>,
         certificate_chain: Option<&[u8]>,
     ) -> binder::Result<()> {
-        let _wp = wd::watch_millis("IKeystoreService::updateSubcomponent", 500);
+        let _wp = wd::watch("IKeystoreService::updateSubcomponent");
         map_or_log_err(self.update_subcomponent(key, public_cert, certificate_chain), Ok)
     }
     fn listEntries(&self, domain: Domain, namespace: i64) -> binder::Result<Vec<KeyDescriptor>> {
-        let _wp = wd::watch_millis("IKeystoreService::listEntries", 500);
+        let _wp = wd::watch("IKeystoreService::listEntries");
         map_or_log_err(self.list_entries(domain, namespace), Ok)
     }
     fn deleteKey(&self, key: &KeyDescriptor) -> binder::Result<()> {
-        let _wp = wd::watch_millis("IKeystoreService::deleteKey", 500);
+        let _wp = wd::watch("IKeystoreService::deleteKey");
         let result = self.delete_key(key);
         log_key_deleted(key, ThreadState::get_calling_uid(), result.is_ok());
         map_or_log_err(result, Ok)
@@ -415,11 +415,11 @@
         grantee_uid: i32,
         access_vector: i32,
     ) -> binder::Result<KeyDescriptor> {
-        let _wp = wd::watch_millis("IKeystoreService::grant", 500);
+        let _wp = wd::watch("IKeystoreService::grant");
         map_or_log_err(self.grant(key, grantee_uid, access_vector.into()), Ok)
     }
     fn ungrant(&self, key: &KeyDescriptor, grantee_uid: i32) -> binder::Result<()> {
-        let _wp = wd::watch_millis("IKeystoreService::ungrant", 500);
+        let _wp = wd::watch("IKeystoreService::ungrant");
         map_or_log_err(self.ungrant(key, grantee_uid), Ok)
     }
     fn listEntriesBatched(
@@ -428,12 +428,12 @@
         namespace: i64,
         start_past_alias: Option<&str>,
     ) -> binder::Result<Vec<KeyDescriptor>> {
-        let _wp = wd::watch_millis("IKeystoreService::listEntriesBatched", 500);
+        let _wp = wd::watch("IKeystoreService::listEntriesBatched");
         map_or_log_err(self.list_entries_batched(domain, namespace, start_past_alias), Ok)
     }
 
     fn getNumberOfEntries(&self, domain: Domain, namespace: i64) -> binder::Result<i32> {
-        let _wp = wd::watch_millis("IKeystoreService::getNumberOfEntries", 500);
+        let _wp = wd::watch("IKeystoreService::getNumberOfEntries");
         map_or_log_err(self.count_num_entries(domain, namespace), Ok)
     }
 }
diff --git a/keystore2/src/super_key.rs b/keystore2/src/super_key.rs
index 11ab734..1f9f5f8 100644
--- a/keystore2/src/super_key.rs
+++ b/keystore2/src/super_key.rs
@@ -919,10 +919,8 @@
                     &key_desc,
                     KeyType::Client, /* TODO Should be Super b/189470584 */
                     |dev| {
-                        let _wp = wd::watch_millis(
-                            "In lock_unlocked_device_required_keys: calling importKey.",
-                            500,
-                        );
+                        let _wp =
+                            wd::watch("In lock_unlocked_device_required_keys: calling importKey.");
                         dev.importKey(key_params.as_slice(), KeyFormat::RAW, &encrypting_key, None)
                     },
                 )?;
diff --git a/keystore2/src/utils.rs b/keystore2/src/utils.rs
index a3fd882..196cac5 100644
--- a/keystore2/src/utils.rs
+++ b/keystore2/src/utils.rs
@@ -143,10 +143,8 @@
         binder::get_interface("permission")?;
 
     let binder_result = {
-        let _wp = watchdog::watch_millis(
-            "In check_device_attestation_permissions: calling checkPermission.",
-            500,
-        );
+        let _wp =
+            watchdog::watch("In check_device_attestation_permissions: calling checkPermission.");
         permission_controller.checkPermission(
             permission,
             ThreadState::get_calling_pid(),
@@ -263,10 +261,7 @@
     log::debug!("import parameters={import_params:?}");
 
     let creation_result = {
-        let _wp = watchdog::watch_millis(
-            "In utils::import_keyblob_and_perform_op: calling importKey.",
-            500,
-        );
+        let _wp = watchdog::watch("In utils::import_keyblob_and_perform_op: calling importKey.");
         map_km_error(km_dev.importKey(&import_params, format, &key_material, None))
     }
     .context(ks_err!("Upgrade failed."))?;
@@ -306,10 +301,7 @@
     NewBlobHandler: FnOnce(&[u8]) -> Result<()>,
 {
     let upgraded_blob = {
-        let _wp = watchdog::watch_millis(
-            "In utils::upgrade_keyblob_and_perform_op: calling upgradeKey.",
-            500,
-        );
+        let _wp = watchdog::watch("In utils::upgrade_keyblob_and_perform_op: calling upgradeKey.");
         map_km_error(km_dev.upgradeKey(key_blob, upgrade_params))
     }
     .context(ks_err!("Upgrade failed."))?;
diff --git a/keystore2/src/watchdog_helper.rs b/keystore2/src/watchdog_helper.rs
index 92a0abc..03c7740 100644
--- a/keystore2/src/watchdog_helper.rs
+++ b/keystore2/src/watchdog_helper.rs
@@ -23,6 +23,11 @@
     pub use watchdog_rs::WatchPoint;
     use watchdog_rs::Watchdog;
 
+    /// Default timeout interval, in milliseconds.
+    pub const DEFAULT_TIMEOUT_MS: u64 = 500;
+
+    const DEFAULT_TIMEOUT: Duration = Duration::from_millis(DEFAULT_TIMEOUT_MS);
+
     lazy_static! {
         /// A Watchdog thread, that can be used to create watch points.
         static ref WD: Arc<Watchdog> = Watchdog::new(Duration::from_secs(10));
@@ -33,6 +38,11 @@
         Watchdog::watch(&WD, id, Duration::from_millis(millis))
     }
 
+    /// Sets a watch point with `id` and a default timeout of [`DEFAULT_TIMEOUT_MS`] milliseconds.
+    pub fn watch(id: &'static str) -> Option<WatchPoint> {
+        Watchdog::watch(&WD, id, DEFAULT_TIMEOUT)
+    }
+
     /// Like `watch_millis` but with a callback that is called every time a report
     /// is printed about this watch point.
     pub fn watch_millis_with(
@@ -53,6 +63,10 @@
     fn watch_millis(_: &'static str, _: u64) -> Option<WatchPoint> {
         None
     }
+    /// Sets a Noop watch point.
+    fn watch(_: &'static str) -> Option<WatchPoint> {
+        None
+    }
 
     pub fn watch_millis_with(
         _: &'static str,
diff --git a/provisioner/rkp_factory_extraction_tool.cpp b/provisioner/rkp_factory_extraction_tool.cpp
index 0a3a59a..62d62cf 100644
--- a/provisioner/rkp_factory_extraction_tool.cpp
+++ b/provisioner/rkp_factory_extraction_tool.cpp
@@ -24,6 +24,7 @@
 #include <remote_prov/remote_prov_utils.h>
 #include <sys/random.h>
 
+#include <future>
 #include <string>
 #include <vector>
 
@@ -91,7 +92,13 @@
 // for every IRemotelyProvisionedComponent.
 void getCsrForInstance(const char* name, void* /*context*/) {
     auto fullName = getFullServiceName(IRemotelyProvisionedComponent::descriptor, name);
-    AIBinder* rkpAiBinder = AServiceManager_getService(fullName.c_str());
+    std::future<AIBinder*> wait_for_service_func =
+        std::async(std::launch::async, AServiceManager_waitForService, fullName.c_str());
+    if (wait_for_service_func.wait_for(std::chrono::seconds(10)) == std::future_status::timeout) {
+        std::cerr << "Wait for service timed out after 10 seconds: " << fullName;
+        exit(-1);
+    }
+    AIBinder* rkpAiBinder = wait_for_service_func.get();
     ::ndk::SpAIBinder rkp_binder(rkpAiBinder);
     auto rkp_service = IRemotelyProvisionedComponent::fromBinder(rkp_binder);
     if (!rkp_service) {