From 78754ac245c8862fb95157f850c5eeff38565b19 Mon Sep 17 00:00:00 2001 From: Duncan Fairbanks Date: Sun, 8 Sep 2024 17:07:02 -0700 Subject: [PATCH] feat(postgres): remove lifetime from `PgAdvisoryLockGuard` Unlike with CPU synchronization primitives, there is no semantic requirement that the guard borrows the lock object. We preserve the optimization of memoizing the release query by wrapping it in an Arc. This way all instances of `PgAdvisoryLockGuard` that originate from the same lock will share a single instance of the release query. --- sqlx-postgres/src/advisory_lock.rs | 35 ++++++++++++++---------------- 1 file changed, 16 insertions(+), 19 deletions(-) diff --git a/sqlx-postgres/src/advisory_lock.rs b/sqlx-postgres/src/advisory_lock.rs index d1aef176fb..48d510db82 100644 --- a/sqlx-postgres/src/advisory_lock.rs +++ b/sqlx-postgres/src/advisory_lock.rs @@ -5,6 +5,7 @@ use hkdf::Hkdf; use once_cell::sync::OnceCell; use sha2::Sha256; use std::ops::{Deref, DerefMut}; +use std::sync::Arc; /// A mutex-like type utilizing [Postgres advisory locks]. /// @@ -37,7 +38,7 @@ use std::ops::{Deref, DerefMut}; pub struct PgAdvisoryLock { key: PgAdvisoryLockKey, /// The query to execute to release this lock. - release_query: OnceCell, + release_query: Arc>, } /// A key type natively used by Postgres advisory locks. @@ -77,8 +78,8 @@ pub enum PgAdvisoryLockKey { /// /// This means the lock is not actually released as soon as the guard is dropped. To ensure the /// lock is eagerly released, you can call [`.release_now().await`][Self::release_now()]. -pub struct PgAdvisoryLockGuard<'lock, C: AsMut> { - lock: &'lock PgAdvisoryLock, +pub struct PgAdvisoryLockGuard> { + lock: PgAdvisoryLock, conn: Option, } @@ -163,7 +164,7 @@ impl PgAdvisoryLock { pub fn with_key(key: PgAdvisoryLockKey) -> Self { Self { key, - release_query: OnceCell::new(), + release_query: Arc::new(OnceCell::new()), } } @@ -201,7 +202,7 @@ impl PgAdvisoryLock { pub async fn acquire>( &self, mut conn: C, - ) -> Result> { + ) -> Result> { match &self.key { PgAdvisoryLockKey::BigInt(key) => { crate::query::query("SELECT pg_advisory_lock($1)") @@ -218,7 +219,7 @@ impl PgAdvisoryLock { } } - Ok(PgAdvisoryLockGuard::new(self, conn)) + Ok(PgAdvisoryLockGuard::new(self.clone(), conn)) } /// Acquires an exclusive lock using `pg_try_advisory_lock()`, returning immediately @@ -244,7 +245,7 @@ impl PgAdvisoryLock { pub async fn try_acquire>( &self, mut conn: C, - ) -> Result, C>> { + ) -> Result, C>> { let locked: bool = match &self.key { PgAdvisoryLockKey::BigInt(key) => { crate::query_scalar::query_scalar("SELECT pg_try_advisory_lock($1)") @@ -262,7 +263,7 @@ impl PgAdvisoryLock { }; if locked { - Ok(Either::Left(PgAdvisoryLockGuard::new(self, conn))) + Ok(Either::Left(PgAdvisoryLockGuard::new(self.clone(), conn))) } else { Ok(Either::Right(conn)) } @@ -322,8 +323,8 @@ impl PgAdvisoryLockKey { const NONE_ERR: &str = "BUG: PgAdvisoryLockGuard.conn taken"; -impl<'lock, C: AsMut> PgAdvisoryLockGuard<'lock, C> { - fn new(lock: &'lock PgAdvisoryLock, conn: C) -> Self { +impl> PgAdvisoryLockGuard { + fn new(lock: PgAdvisoryLock, conn: C) -> Self { PgAdvisoryLockGuard { lock, conn: Some(conn), @@ -362,7 +363,7 @@ impl<'lock, C: AsMut> PgAdvisoryLockGuard<'lock, C> { } } -impl<'lock, C: AsMut + AsRef> Deref for PgAdvisoryLockGuard<'lock, C> { +impl + AsRef> Deref for PgAdvisoryLockGuard { type Target = PgConnection; fn deref(&self) -> &Self::Target { @@ -376,17 +377,13 @@ impl<'lock, C: AsMut + AsRef> Deref for PgAdvisoryLo /// However, replacing the connection with a different one using, e.g. [`std::mem::replace()`] /// is a logic error and will cause a warning to be logged by the PostgreSQL server when this /// guard attempts to release the lock. -impl<'lock, C: AsMut + AsRef> DerefMut - for PgAdvisoryLockGuard<'lock, C> -{ +impl + AsRef> DerefMut for PgAdvisoryLockGuard { fn deref_mut(&mut self) -> &mut Self::Target { self.conn.as_mut().expect(NONE_ERR).as_mut() } } -impl<'lock, C: AsMut + AsRef> AsRef - for PgAdvisoryLockGuard<'lock, C> -{ +impl + AsRef> AsRef for PgAdvisoryLockGuard { fn as_ref(&self) -> &PgConnection { self.conn.as_ref().expect(NONE_ERR).as_ref() } @@ -398,7 +395,7 @@ impl<'lock, C: AsMut + AsRef> AsRef /// However, replacing the connection with a different one using, e.g. [`std::mem::replace()`] /// is a logic error and will cause a warning to be logged by the PostgreSQL server when this /// guard attempts to release the lock. -impl<'lock, C: AsMut> AsMut for PgAdvisoryLockGuard<'lock, C> { +impl> AsMut for PgAdvisoryLockGuard { fn as_mut(&mut self) -> &mut PgConnection { self.conn.as_mut().expect(NONE_ERR).as_mut() } @@ -407,7 +404,7 @@ impl<'lock, C: AsMut> AsMut for PgAdvisoryLockGuard< /// Queues a `pg_advisory_unlock()` call on the wrapped connection which will be flushed /// to the server the next time it is used, or when it is returned to [`PgPool`][crate::PgPool] /// in the case of [`PoolConnection`][crate::pool::PoolConnection]. -impl<'lock, C: AsMut> Drop for PgAdvisoryLockGuard<'lock, C> { +impl> Drop for PgAdvisoryLockGuard { fn drop(&mut self) { if let Some(mut conn) = self.conn.take() { // Queue a simple query message to execute next time the connection is used.