Skip to content

Commit

Permalink
Problem: use of Option<Parent> in SubTransaction
Browse files Browse the repository at this point in the history 
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
  • Loading branch information
@yrashk
yrashk committed Dec 14, 2022
1 parent df654a1 commit d04aea8
Show file tree
Hide file tree
Showing 2 changed files with 142 additions and 94 deletions.
4 changes: 2 additions & 2 deletions pgx-tests/src/tests/subxact_tests.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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);
Expand All @@ -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
Expand Down
232 changes: 140 additions & 92 deletions pgx/src/subxact.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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
}
}

Expand Down

0 comments on commit d04aea8

Please sign in to comment.