Problem: use of Option<Parent> in SubTransaction

It is only there so that we can move it out of manually copied type
that implements Drop

Solution: extract release logic into a separate trait

Author: yrashk

pgx-tests/src/tests/subxact_tests.rs

@@ -57,7 +57,7 @@ mod tests {
         Spi::execute(|c| {
             c.update("CREATE TABLE a (v INTEGER)", None, None);
             // The type below is explicit to ensure it's commit on drop by default
-            c.sub_transaction(|xact: SubTransaction<SpiClient, true>| {
+            c.sub_transaction(|xact: SubTransaction<SpiClient, CommitOnDrop>| {
                 xact.update("INSERT INTO a VALUES (0)", None, None);
                 // Dropped explicitly for illustration purposes
                 drop(xact);
@@ -81,7 +81,7 @@ mod tests {
         Spi::execute(|c| {
             c.update("CREATE TABLE a (v INTEGER)", None, None);
             // The type below is explicit to ensure it's commit on drop by default
-            c.sub_transaction(|xact: SubTransaction<SpiClient, true>| {
+            c.sub_transaction(|xact: SubTransaction<SpiClient, CommitOnDrop>| {
                 xact.update("INSERT INTO a VALUES (0)", None, None);
                 let xact = xact.rollback_on_drop();
                 // Dropped explicitly for illustration purposes

pgx/src/subxact.rs

@@ -2,152 +2,200 @@ use crate::{pg_sys, PgMemoryContexts, SpiClient};
 use std::fmt::Debug;
 use std::ops::Deref;
 
-/// Sub-transaction
-///
-/// Can be created by calling `SpiClient::sub_transaction`, `SubTransaction<Parent>::sub_transaction`
-/// or any other implementation of `SubTransactionExt` and obtaining it as an argument to the provided closure.
-///
-/// Unless rolled back or committed explicitly, it'll commit if `COMMIT` generic parameter is `true`
-/// (default) or roll back if it is `false`.
-#[derive(Debug)]
-pub struct SubTransaction<Parent: SubTransactionExt, const COMMIT: bool = true> {
+/// Releases a sub-transaction on Drop
+pub trait ReleaseOnDrop {}
+
+/// Sub-transaction's contextual information
+#[derive(Clone, Copy)]
+pub struct Context {
     memory_context: pg_sys::MemoryContext,
     // Resource ownership before the transaction
     //
     // Based on information from src/backend/utils/resowner/README
     // as well as practical use of it in src/pl/plpython/plpy_spi.c
     resource_owner: pg_sys::ResourceOwner,
-    // Should the transaction be released, or was it already committed or rolled back?
-    //
-    // The reason we are not calling this `released` as we're also using this flag when
-    // we convert between commit_on_drop and rollback_on_drop to ensure it doesn't get released
-    // on the drop of the original value.
-    should_release: bool,
-    parent: Option<Parent>,
 }
 
-impl<Parent: SubTransactionExt, const COMMIT: bool> SubTransaction<Parent, COMMIT> {
-    /// Create a new sub-transaction.
-    fn new(parent: Parent) -> Self {
+impl Context {
+    /// Captures the context
+    fn capture() -> Self {
         // Remember the memory context before starting the sub-transaction
-        let ctx = PgMemoryContexts::CurrentMemoryContext.value();
+        let memory_context = PgMemoryContexts::CurrentMemoryContext.value();
         // Remember resource owner before starting the sub-transaction
         let resource_owner = unsafe { pg_sys::CurrentResourceOwner };
-        unsafe {
-            pg_sys::BeginInternalSubTransaction(std::ptr::null() /* [no] transaction name */);
-        }
-        // Switch to the outer memory context so that all allocations remain
-        // there instead of the sub-transaction's context
-        PgMemoryContexts::For(ctx).set_as_current();
-        Self { memory_context: ctx, should_release: true, resource_owner, parent: Some(parent) }
+        Self { memory_context, resource_owner }
     }
+}
 
-    /// Commit the transaction, returning its parent
-    pub fn commit(mut self) -> Parent {
-        self.internal_commit();
-        self.should_release = false;
-        self.parent.take().unwrap()
+impl From<Context> for CommitOnDrop {
+    fn from(context: Context) -> Self {
+        CommitOnDrop(context)
     }
+}
 
-    /// Rollback the transaction, returning its parent
-    pub fn rollback(mut self) -> Parent {
-        self.internal_rollback();
-        self.should_release = false;
-        self.parent.take().unwrap()
+impl From<Context> for RollbackOnDrop {
+    fn from(context: Context) -> Self {
+        RollbackOnDrop(context)
     }
+}
 
-    fn internal_rollback(&self) {
+/// Commits a sub-transaction on Drop
+pub struct CommitOnDrop(Context);
+
+impl Drop for CommitOnDrop {
+    fn drop(&mut self) {
         unsafe {
-            pg_sys::RollbackAndReleaseCurrentSubTransaction();
-            pg_sys::CurrentResourceOwner = self.resource_owner;
+            pg_sys::ReleaseCurrentSubTransaction();
+            pg_sys::CurrentResourceOwner = self.0.resource_owner;
         }
-        PgMemoryContexts::For(self.memory_context).set_as_current();
+        PgMemoryContexts::For(self.0.memory_context).set_as_current();
     }
+}
 
-    fn internal_commit(&self) {
+impl ReleaseOnDrop for CommitOnDrop {}
+
+/// Rolls back a sub-transaction on Drop
+pub struct RollbackOnDrop(Context);
+
+impl Drop for RollbackOnDrop {
+    fn drop(&mut self) {
         unsafe {
-            pg_sys::ReleaseCurrentSubTransaction();
-            pg_sys::CurrentResourceOwner = self.resource_owner;
+            pg_sys::RollbackAndReleaseCurrentSubTransaction();
+            pg_sys::CurrentResourceOwner = self.0.resource_owner;
         }
-        PgMemoryContexts::For(self.memory_context).set_as_current();
+        PgMemoryContexts::For(self.0.memory_context).set_as_current();
     }
 }
 
-impl<Parent: SubTransactionExt> SubTransaction<Parent, true> {
-    /// Make this sub-transaction roll back on drop
-    pub fn rollback_on_drop(self) -> SubTransaction<Parent, false> {
-        self.into()
+impl ReleaseOnDrop for RollbackOnDrop {}
+
+impl Into<RollbackOnDrop> for CommitOnDrop {
+    fn into(self) -> RollbackOnDrop {
+        let result = RollbackOnDrop(self.0);
+        // IMPORTANT: avoid running Drop (that would commit)
+        std::mem::forget(self);
+        result
     }
 }
 
-impl<Parent: SubTransactionExt> SubTransaction<Parent, false> {
-    /// Make this sub-transaction commit on drop
-    pub fn commit_on_drop(self) -> SubTransaction<Parent, true> {
-        self.into()
+impl Into<CommitOnDrop> for RollbackOnDrop {
+    fn into(self) -> CommitOnDrop {
+        let result = CommitOnDrop(self.0);
+        // IMPORTANT: avoid running Drop (that would roll back)
+        std::mem::forget(self);
+        result
     }
 }
 
-impl<Parent: SubTransactionExt> Into<SubTransaction<Parent, false>>
-    for SubTransaction<Parent, true>
+struct NoOpOnDrop;
+
+impl ReleaseOnDrop for NoOpOnDrop {}
+
+/// Sub-transaction
+///
+/// Can be created by calling `SpiClient::sub_transaction`, `SubTransaction<Parent>::sub_transaction`
+/// or any other implementation of `SubTransactionExt` and obtaining it as an argument to the provided closure.
+///
+/// Unless rolled back or committed explicitly, it'll commit if `Release` generic parameter is `CommitOnDrop`
+/// (default) or roll back if it is `RollbackOnDrop`.
+#[derive(Debug)]
+pub struct SubTransaction<Parent: SubTransactionExt, Release: ReleaseOnDrop = CommitOnDrop> {
+    release: Release,
+    parent: Parent,
+}
+
+impl<Parent: SubTransactionExt, Release: ReleaseOnDrop> SubTransaction<Parent, Release>
+where
+    Release: From<Context>,
 {
-    fn into(mut self) -> SubTransaction<Parent, false> {
-        let result = SubTransaction {
-            memory_context: self.memory_context,
-            resource_owner: self.resource_owner,
-            should_release: self.should_release,
-            parent: self.parent.take(),
-        };
-        // Make sure original sub-transaction won't commit
-        self.should_release = false;
-        result
+    /// Create a new sub-transaction.
+    fn new(parent: Parent) -> Self {
+        let context = Context::capture();
+        let memory_context = context.memory_context;
+        let release = context.into();
+        unsafe {
+            pg_sys::BeginInternalSubTransaction(std::ptr::null() /* [no] transaction name */);
+        }
+        // Switch to the outer memory context so that all allocations remain
+        // there instead of the sub-transaction's context
+        PgMemoryContexts::For(memory_context).set_as_current();
+        Self { release, parent }
     }
 }
 
-impl<Parent: SubTransactionExt> Into<SubTransaction<Parent, true>>
-    for SubTransaction<Parent, false>
-{
-    fn into(mut self) -> SubTransaction<Parent, true> {
-        let result = SubTransaction {
-            memory_context: self.memory_context,
-            resource_owner: self.resource_owner,
-            should_release: self.should_release,
-            parent: self.parent.take(),
-        };
-        // Make sure original sub-transaction won't roll back
-        self.should_release = false;
-        result
+impl<Parent: SubTransactionExt> SubTransaction<Parent, CommitOnDrop> {
+    /// Commit the transaction, returning its parent
+    pub fn commit(self) -> Parent {
+        // `Self::do_nothing_on_drop()` will commit as `Release` is `CommitOnDrop`
+        self.do_nothing_on_drop().parent
     }
 }
 
-impl<Parent: SubTransactionExt, const COMMIT: bool> Drop for SubTransaction<Parent, COMMIT> {
-    fn drop(&mut self) {
-        if self.should_release {
-            if COMMIT {
-                self.internal_commit();
-            } else {
-                self.internal_rollback();
-            }
-        }
+impl<Parent: SubTransactionExt> SubTransaction<Parent, RollbackOnDrop> {
+    /// Commit the transaction, returning its parent
+    pub fn commit(self) -> Parent {
+        // Make sub-transaction commit on drop and then use `commit`
+        self.commit_on_drop().commit()
+    }
+}
+
+impl<Parent: SubTransactionExt> SubTransaction<Parent, RollbackOnDrop> {
+    /// Rollback the transaction, returning its parent
+    pub fn rollback(self) -> Parent {
+        // `Self::do_nothing_on_drop()` will roll back as `Release` is `RollbackOnDrop`
+        self.do_nothing_on_drop().parent
+    }
+}
+
+impl<Parent: SubTransactionExt> SubTransaction<Parent, CommitOnDrop> {
+    /// Rollback the transaction, returning its parent
+    pub fn rollback(self) -> Parent {
+        // Make sub-transaction roll back on drop and then use `rollback`
+        self.rollback_on_drop().rollback()
+    }
+}
+
+impl<Parent: SubTransactionExt> SubTransaction<Parent, CommitOnDrop> {
+    /// Make this sub-transaction roll back on drop
+    pub fn rollback_on_drop(self) -> SubTransaction<Parent, RollbackOnDrop> {
+        SubTransaction { parent: self.parent, release: self.release.into() }
+    }
+}
+
+impl<Parent: SubTransactionExt> SubTransaction<Parent, RollbackOnDrop> {
+    /// Make this sub-transaction commit on drop
+    pub fn commit_on_drop(self) -> SubTransaction<Parent, CommitOnDrop> {
+        SubTransaction { parent: self.parent, release: self.release.into() }
+    }
+}
+
+impl<Parent: SubTransactionExt, Release: ReleaseOnDrop> SubTransaction<Parent, Release> {
+    /// Make this sub-transaction do nothing on drop
+    ///
+    /// Releases the sub-transaction based on `Release` generic parameter. Further
+    /// dropping of the sub-transaction will not do anything.
+    fn do_nothing_on_drop(self) -> SubTransaction<Parent, NoOpOnDrop> {
+        SubTransaction { parent: self.parent, release: NoOpOnDrop }
     }
 }
 
 // This allows SubTransaction to be de-referenced to SpiClient
-impl<'conn, const COMMIT: bool> Deref for SubTransaction<SpiClient<'conn>, COMMIT> {
+impl<'conn, Release: ReleaseOnDrop> Deref for SubTransaction<SpiClient<'conn>, Release> {
     type Target = SpiClient<'conn>;
 
     fn deref(&self) -> &Self::Target {
-        self.parent.as_ref().unwrap()
+        &self.parent
     }
 }
 
 // This allows a SubTransaction of a SubTransaction to be de-referenced to SpiClient
-impl<Parent: SubTransactionExt, const COMMIT: bool> Deref
-    for SubTransaction<SubTransaction<Parent>, COMMIT>
+impl<Parent: SubTransactionExt, Release: ReleaseOnDrop> Deref
+    for SubTransaction<SubTransaction<Parent>, Release>
 {
     type Target = Parent;
 
     fn deref(&self) -> &Self::Target {
-        self.parent.as_ref().and_then(|p| p.parent.as_ref()).unwrap()
+        &self.parent.parent
     }
 }