From 697a7208b611f4133dacc004f3e040ece16a48ca Mon Sep 17 00:00:00 2001 From: Oleksandr Deundiak Date: Thu, 19 Dec 2024 15:15:32 +0100 Subject: [PATCH 1/2] feat(rust): optimize sqlite queries --- ...thority_enrollment_token_repository_sql.rs | 62 ++++++------------- .../src/identities/identities_attributes.rs | 8 ++- .../storage/identity_attributes_repository.rs | 10 ++- .../identity_attributes_repository_sql.rs | 18 +++--- 4 files changed, 42 insertions(+), 56 deletions(-) diff --git a/implementations/rust/ockam/ockam_api/src/authenticator/storage/authority_enrollment_token_repository_sql.rs b/implementations/rust/ockam/ockam_api/src/authenticator/storage/authority_enrollment_token_repository_sql.rs index 61d9d5bf78a..fc311069dc1 100644 --- a/implementations/rust/ockam/ockam_api/src/authenticator/storage/authority_enrollment_token_repository_sql.rs +++ b/implementations/rust/ockam/ockam_api/src/authenticator/storage/authority_enrollment_token_repository_sql.rs @@ -53,49 +53,25 @@ impl AuthorityEnrollmentTokenRepository for AuthorityEnrollmentTokenSqlxDatabase one_time_code: OneTimeCode, now: TimestampInSeconds, ) -> Result> { - // We need to delete expired tokens regularly - // Also makes sure we don't get expired tokens later inside this function - let query1 = query("DELETE FROM authority_enrollment_token WHERE expires_at <= $1") - .bind(now.0 as i64); - - let res = query1.execute(&*self.database.pool).await.into_core()?; - debug!("Deleted {} expired enrollment tokens", res.rows_affected()); - - let mut transaction = self.database.pool.begin().await.into_core()?; - - let query2 = query_as("SELECT one_time_code, reference, issued_by, created_at, expires_at, ttl_count, attributes FROM authority_enrollment_token WHERE one_time_code = $1") - .bind(one_time_code); - let row: Option = - query2.fetch_optional(&mut *transaction).await.into_core()?; - let token: Option = row.map(|r| r.try_into()).transpose()?; - - if let Some(token) = &token { - if token.ttl_count <= 1 { - let query3 = - query("DElETE FROM authority_enrollment_token WHERE one_time_code = $1") - .bind(one_time_code); - query3.execute(&mut *transaction).await.void()?; - debug!( - "Deleted enrollment token because it has been used. Reference: {}", - token.reference() - ); - } else { - let new_ttl_count = token.ttl_count - 1; - let query3 = query( - "UPDATE authority_enrollment_token SET ttl_count = $1 WHERE one_time_code = $2", - ) - .bind(new_ttl_count as i64) - .bind(one_time_code); - query3.execute(&mut *transaction).await.void()?; - debug!( - "Decreasing enrollment token usage count to {}. Reference: {}", - new_ttl_count, - token.reference() - ); - } - } - - transaction.commit().await.void()?; + // FIXME: We need to delete expired tokens regularly + // FIXME 2: Also now need to clear tokens with ttl_count = 0 + + let query = query_as("UPDATE authority_enrollment_token SET ttl_count = ttl_count - 1 WHERE one_time_code = $1 AND expires_at > $2 AND ttl_count > 0 RETURNING one_time_code, reference, issued_by, created_at, expires_at, ttl_count, attributes") + .bind(one_time_code) + .bind(now); + let row: Option = query + .fetch_optional(&*self.database.pool) + .await + .into_core()?; + let token = if let Some(row) = row { + let mut t = EnrollmentToken::try_from(row)?; + + t.ttl_count += 1; // We decremented it inside the query + + Some(t) + } else { + None + }; Ok(token) } diff --git a/implementations/rust/ockam/ockam_identity/src/identities/identities_attributes.rs b/implementations/rust/ockam/ockam_identity/src/identities/identities_attributes.rs index 9df0859face..7d28e4ebb4d 100644 --- a/implementations/rust/ockam/ockam_identity/src/identities/identities_attributes.rs +++ b/implementations/rust/ockam/ockam_identity/src/identities/identities_attributes.rs @@ -32,8 +32,12 @@ impl IdentitiesAttributes { subject: &Identifier, attested_by: &Identifier, ) -> Result> { - self.repository.delete_expired_attributes(now()?).await?; - self.repository.get_attributes(subject, attested_by).await + let now = now()?; + // FIXME: This should be run periodically in a separate task + // self.repository.delete_expired_attributes(now()?).await?; + self.repository + .get_non_expired_attributes(subject, attested_by, now) + .await } /// Set the attributes associated with the given identity identifier. diff --git a/implementations/rust/ockam/ockam_identity/src/identities/storage/identity_attributes_repository.rs b/implementations/rust/ockam/ockam_identity/src/identities/storage/identity_attributes_repository.rs index 614c87b851c..7afe957c53d 100644 --- a/implementations/rust/ockam/ockam_identity/src/identities/storage/identity_attributes_repository.rs +++ b/implementations/rust/ockam/ockam_identity/src/identities/storage/identity_attributes_repository.rs @@ -11,10 +11,11 @@ use ockam_node::retry; #[async_trait] pub trait IdentityAttributesRepository: Send + Sync + 'static { /// Get the attributes associated with the given identity identifier - async fn get_attributes( + async fn get_non_expired_attributes( &self, subject: &Identifier, attested_by: &Identifier, + now: TimestampInSeconds, ) -> Result>; /// Set the attributes associated with the given identity identifier. @@ -28,12 +29,15 @@ pub trait IdentityAttributesRepository: Send + Sync + 'static { #[cfg(feature = "std")] #[async_trait] impl IdentityAttributesRepository for AutoRetry { - async fn get_attributes( + async fn get_non_expired_attributes( &self, subject: &Identifier, attested_by: &Identifier, + now: TimestampInSeconds, ) -> Result> { - retry!(self.wrapped.get_attributes(subject, attested_by)) + retry!(self + .wrapped + .get_non_expired_attributes(subject, attested_by, now)) } async fn put_attributes(&self, subject: &Identifier, entry: AttributesEntry) -> Result<()> { diff --git a/implementations/rust/ockam/ockam_identity/src/identities/storage/identity_attributes_repository_sql.rs b/implementations/rust/ockam/ockam_identity/src/identities/storage/identity_attributes_repository_sql.rs index 79561bff4fd..c38c1fea855 100644 --- a/implementations/rust/ockam/ockam_identity/src/identities/storage/identity_attributes_repository_sql.rs +++ b/implementations/rust/ockam/ockam_identity/src/identities/storage/identity_attributes_repository_sql.rs @@ -54,16 +54,18 @@ impl IdentityAttributesSqlxDatabase { #[async_trait] impl IdentityAttributesRepository for IdentityAttributesSqlxDatabase { - async fn get_attributes( + async fn get_non_expired_attributes( &self, identity: &Identifier, attested_by: &Identifier, + now: TimestampInSeconds, ) -> Result> { let query = query_as( - "SELECT identifier, attributes, added, expires, attested_by FROM identity_attributes WHERE identifier = $1 AND attested_by = $2 AND node_name = $3" + "SELECT identifier, attributes, added, expires, attested_by FROM identity_attributes WHERE identifier = $1 AND attested_by = $2 AND (expires > $3 or expires IS NULL) AND node_name = $4" ) .bind(identity) .bind(attested_by) + .bind(now) .bind(&self.node_name); let identity_attributes: Option = query .fetch_optional(&*self.database.pool) @@ -186,12 +188,12 @@ mod tests { .await?; let result = repository - .get_attributes(&identifier1, &identifier1) + .get_non_expired_attributes(&identifier1, &identifier1, now) .await?; assert_eq!(result, Some(attributes1.clone())); let result = repository - .get_attributes(&identifier2, &identifier2) + .get_non_expired_attributes(&identifier2, &identifier2, now) .await?; assert_eq!(result, Some(attributes2.clone())); @@ -236,17 +238,17 @@ mod tests { repository.delete_expired_attributes(now.add(10)).await?; let result = repository - .get_attributes(&identifier1, &identifier1) + .get_non_expired_attributes(&identifier1, &identifier1, now) .await?; assert_eq!(result, None); let result = repository - .get_attributes(&identifier2, &identifier2) + .get_non_expired_attributes(&identifier2, &identifier2, now) .await?; assert_eq!(result, None); let result = repository - .get_attributes(&identifier3, &identifier3) + .get_non_expired_attributes(&identifier3, &identifier3, now) .await?; assert_eq!( result, @@ -255,7 +257,7 @@ mod tests { ); let result = repository - .get_attributes(&identifier4, &identifier4) + .get_non_expired_attributes(&identifier4, &identifier4, now) .await?; assert_eq!( result, From e1b1bddea52b48733f36c2f133a7e6df0e50dc85 Mon Sep 17 00:00:00 2001 From: Oleksandr Deundiak Date: Thu, 19 Dec 2024 16:44:14 +0100 Subject: [PATCH 2/2] chore(rust): change sqlite configuration --- .../ockam/ockam_node/src/storage/database/sqlx_database.rs | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/implementations/rust/ockam/ockam_node/src/storage/database/sqlx_database.rs b/implementations/rust/ockam/ockam_node/src/storage/database/sqlx_database.rs index 2fa7f525ff1..0f2c8368429 100644 --- a/implementations/rust/ockam/ockam_node/src/storage/database/sqlx_database.rs +++ b/implementations/rust/ockam/ockam_node/src/storage/database/sqlx_database.rs @@ -302,9 +302,11 @@ impl SqlxDatabase { let _ = connection .execute( r#" -PRAGMA synchronous = EXTRA; +PRAGMA synchronous = NORMAL; PRAGMA locking_mode = NORMAL; PRAGMA busy_timeout = 10000; +PRAGMA cache_size = -50000; -- 50 MB +PRAGMA mmap_size = 52428800; -- 50 MB "#, ) .await