From c46e2a8be93254188d4a3547d9a0eb076b9eec7b Mon Sep 17 00:00:00 2001 From: Oleksandr Deundiak Date: Wed, 18 Dec 2024 20:19:37 +0100 Subject: [PATCH] feat(rust): optimize sqlite queries --- ...thority_enrollment_token_repository_sql.rs | 62 +++++----------- .../rust/ockam/ockam_api/tests/perf.rs | 73 +++++++++++++++++++ .../src/identities/identities_attributes.rs | 8 +- .../storage/identity_attributes_repository.rs | 10 ++- .../identity_attributes_repository_sql.rs | 18 +++-- 5 files changed, 115 insertions(+), 56 deletions(-) create mode 100644 implementations/rust/ockam/ockam_api/tests/perf.rs 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_api/tests/perf.rs b/implementations/rust/ockam/ockam_api/tests/perf.rs new file mode 100644 index 00000000000..040a0d7fd92 --- /dev/null +++ b/implementations/rust/ockam/ockam_api/tests/perf.rs @@ -0,0 +1,73 @@ +use crate::common::common::{change_client_identifier, start_authority, AuthorityInfo}; +use ockam::identity::SecureChannels; +use ockam_api::authenticator::direct::{ + OCKAM_ROLE_ATTRIBUTE_ENROLLER_VALUE, OCKAM_ROLE_ATTRIBUTE_KEY, +}; +use ockam_api::authenticator::enrollment_tokens::TokenIssuer; +use ockam_core::Result; +use ockam_node::Context; +use std::collections::BTreeMap; +use std::time::Instant; + +mod common; + +#[ockam_macros::test] +async fn enrollment_token_loop(ctx: &mut Context) -> Result<()> { + let secure_channels = SecureChannels::builder().await?.build(); + + let AuthorityInfo { admins, .. } = start_authority(ctx, secure_channels.clone(), 1).await?; + let admin = &admins[0]; + + let mut attributes = BTreeMap::::default(); + attributes.insert( + OCKAM_ROLE_ATTRIBUTE_KEY.to_string(), + OCKAM_ROLE_ATTRIBUTE_ENROLLER_VALUE.to_string(), + ); + let otc = admin + .client + .create_token(ctx, attributes.clone(), None, None) + .await + .unwrap(); + + let enroller = secure_channels + .identities() + .identities_creation() + .create_identity() + .await?; + let enroller_client = change_client_identifier(&admin.client, &enroller, None); + + { + use ockam_api::authenticator::enrollment_tokens::TokenAcceptor; + enroller_client.present_token(ctx, otc).await.unwrap(); + } + + let mut attributes_member = BTreeMap::::default(); + attributes_member.insert("KEY".to_string(), "VALUE".to_string()); + + let otc = enroller_client + .create_token(ctx, attributes_member.clone(), None, Some(1024)) + .await + .unwrap(); + + let t1 = Instant::now(); + for _ in 0..1000 { + let member = secure_channels + .identities() + .identities_creation() + .create_identity() + .await?; + let member_client = change_client_identifier(&admin.client, &member, None); + + member_client.present_token(ctx, &otc).await.unwrap(); + + use ockam_api::enroll::enrollment::Enrollment; + + member_client.issue_credential(ctx).await.unwrap(); + } + + let t2 = Instant::now(); + + println!("TIME: {:?}", t2 - t1); + + Ok(()) +} 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,