panic when using pooled connection with a collection of async query

[fakeshadow]

Setup

Versions

• Rust: stable-x86_64-unknown-linux-gnu - rustc 1.81.0 (eeb90cda1 2024-09-04)
• Diesel: 2.2.4
• Diesel_async: 0.5.0
• Database: PostgreSQL 16.3
• Operating System Linux 6.10.6-10

Feature Flags

• diesel: []
• diesel_async: ["bb8", "postgres"]

Problem Description

This code would cause panic with message "Cannot access shared transaction state"
Cargo.toml:

diesel = "2"
diesel-async = { version = "0.5", features = ["bb8", "postgres"] }
tokio = { version = "1", features = ["rt-multi-thread", "macros"] }

main.rs:

use diesel::{sql_types::Integer, IntoSql};
use diesel_async::{
    pooled_connection::{bb8, AsyncDieselConnectionManager},
    AsyncPgConnection, RunQueryDsl,
};

#[tokio::main]
async fn main() {
    let pool = bb8::Pool::<AsyncPgConnection>::builder()
        .max_size(1)
        .test_on_check_out(false)
        .idle_timeout(None)
        .max_lifetime(None)
        .build(AsyncDieselConnectionManager::new(
            "postgres://postgres:postgres@localhost/test",
        ))
        .await
        .unwrap();

    // keep connection until task are finished and it's ok
    // let mut conn = pool.get().await.unwrap();

    let task = {
        // drop connection and return it to pool early cause panic
        let mut conn = pool.get().await.unwrap();

        core::iter::repeat_with(|| diesel::select(1_i32.into_sql::<Integer>()).execute(&mut conn))
            .take(20)
            .collect::<Vec<_>>()
    };

    for fut in task {
        fut.await.unwrap();
    }
}

Some additional context:

• Like the comments in example code if the pooled connection is kept around until Vec<impl Future> finished iteration it runs to complete sucessfully.
• I have tested deadpool feature and it acts similarly.
• I have tested early releases and diese-async = "0.2" seems to be the latest version that does not cause this panic.

What are you trying to accomplish?

Releasing pooled connection before collecting async responses

What is the expected output?

The code should execute successfully without panic

What is the actual output?

thread 'main' panicked at /home/fakeshadow/.cargo/registry/src/index.crates.io-6f17d22bba15001f/diesel-async-0.5.0/src/pg/mod.rs:216:13:
Cannot access shared transaction state
stack backtrace:
   0: rust_begin_unwind
             at /rustc/eeb90cda1969383f56a2637cbd3037bdf598841c/library/std/src/panicking.rs:665:5
   1: core::panicking::panic_fmt
             at /rustc/eeb90cda1969383f56a2637cbd3037bdf598841c/library/core/src/panicking.rs:74:14
   2: <diesel_async::pg::AsyncPgConnection as diesel_async::AsyncConnection>::transaction_state
             at /home/fakeshadow/.cargo/registry/src/index.crates.io-6f17d22bba15001f/diesel-async-0.5.0/src/pg/mod.rs:216:13
   3: <diesel_async::transaction_manager::AnsiTransactionManager as diesel_async::transaction_manager::TransactionManager<Conn>>::transaction_manager_status_mut
             at /home/fakeshadow/.cargo/registry/src/index.crates.io-6f17d22bba15001f/diesel-async-0.5.0/src/transaction_manager.rs:473:14
   4: diesel_async::transaction_manager::TransactionManager::is_broken_transaction_manager
             at /home/fakeshadow/.cargo/registry/src/index.crates.io-6f17d22bba15001f/diesel-async-0.5.0/src/transaction_manager.rs:91:15
   5: <diesel_async::pg::AsyncPgConnection as diesel_async::pooled_connection::PoolableConnection>::is_broken
             at /home/fakeshadow/.cargo/registry/src/index.crates.io-6f17d22bba15001f/diesel-async-0.5.0/src/pg/mod.rs:839:9
   6: diesel_async::pooled_connection::bb8::<impl bb8::api::ManageConnection for diesel_async::pooled_connection::AsyncDieselConnectionManager<C>>::has_broken
             at /home/fakeshadow/.cargo/registry/src/index.crates.io-6f17d22bba15001f/diesel-async-0.5.0/src/pooled_connection/bb8.rs:87:37
   7: bb8::inner::PoolInner<M>::put_back
             at /home/fakeshadow/.cargo/registry/src/index.crates.io-6f17d22bba15001f/bb8-0.8.5/src/inner.rs:151:23
   8: <bb8::api::PooledConnection<M> as core::ops::drop::Drop>::drop
             at /home/fakeshadow/.cargo/registry/src/index.crates.io-6f17d22bba15001f/bb8-0.8.5/src/api.rs:475:13
   9: core::ptr::drop_in_place<bb8::api::PooledConnection<diesel_async::pooled_connection::AsyncDieselConnectionManager<diesel_async::pg::AsyncPgConnection>>>
             at /rustc/eeb90cda1969383f56a2637cbd3037bdf598841c/library/core/src/ptr/mod.rs:542:1
  10: qa::main::{{closure}}
             at ./src/main.rs:67:5
  11: tokio::runtime::park::CachedParkThread::block_on::{{closure}}
             at /home/fakeshadow/.cargo/registry/src/index.crates.io-6f17d22bba15001f/tokio-1.40.0/src/runtime/park.rs:281:63
  12: tokio::runtime::coop::with_budget
             at /home/fakeshadow/.cargo/registry/src/index.crates.io-6f17d22bba15001f/tokio-1.40.0/src/runtime/coop.rs:107:5
  13: tokio::runtime::coop::budget
             at /home/fakeshadow/.cargo/registry/src/index.crates.io-6f17d22bba15001f/tokio-1.40.0/src/runtime/coop.rs:73:5
  14: tokio::runtime::park::CachedParkThread::block_on
             at /home/fakeshadow/.cargo/registry/src/index.crates.io-6f17d22bba15001f/tokio-1.40.0/src/runtime/park.rs:281:31
  15: tokio::runtime::context::blocking::BlockingRegionGuard::block_on
             at /home/fakeshadow/.cargo/registry/src/index.crates.io-6f17d22bba15001f/tokio-1.40.0/src/runtime/context/blocking.rs:66:9
  16: tokio::runtime::scheduler::multi_thread::MultiThread::block_on::{{closure}}
             at /home/fakeshadow/.cargo/registry/src/index.crates.io-6f17d22bba15001f/tokio-1.40.0/src/runtime/scheduler/multi_thread/mod.rs:87:13
  17: tokio::runtime::context::runtime::enter_runtime
             at /home/fakeshadow/.cargo/registry/src/index.crates.io-6f17d22bba15001f/tokio-1.40.0/src/runtime/context/runtime.rs:65:16
  18: tokio::runtime::scheduler::multi_thread::MultiThread::block_on
             at /home/fakeshadow/.cargo/registry/src/index.crates.io-6f17d22bba15001f/tokio-1.40.0/src/runtime/scheduler/multi_thread/mod.rs:86:9
  19: tokio::runtime::runtime::Runtime::block_on_inner
             at /home/fakeshadow/.cargo/registry/src/index.crates.io-6f17d22bba15001f/tokio-1.40.0/src/runtime/runtime.rs:363:45
  20: tokio::runtime::runtime::Runtime::block_on
             at /home/fakeshadow/.cargo/registry/src/index.crates.io-6f17d22bba15001f/tokio-1.40.0/src/runtime/runtime.rs:335:13
  21: qa::main
             at ./src/main.rs:69:5
  22: core::ops::function::FnOnce::call_once
             at /rustc/eeb90cda1969383f56a2637cbd3037bdf598841c/library/core/src/ops/function.rs:250:5
note: Some details are omitted, run with `RUST_BACKTRACE=full` for a verbose backtrace.

Are you seeing any additional errors?

No

Steps to reproduce

As the description

Checklist

• I have already looked over the issue tracker for similar possible closed issues.

• This issue can be reproduced on Rust's stable channel. (Your issue will be
  closed if this is not the case)
• This issue can be reproduced without requiring a third party crate

[fakeshadow]

I don't know the code base very well so I could be wrong but it looks like the cause of panic is the currently logic expecting connection is kept around until it's associated query future resolved.

Namely Arc<Mutex<AnsiTransactionManager>> is cloned into the query future while at the same time connection pool want to get &mut AnsiTransactionManager by calling Arc::get_mut on it.

I believe there could be at least two approaches to this issue:

1. let the query future hold a reference of connection until it's resolved:

let conn = pool.get().await?;
let query = diesel::select(1_i32.into_sql::<Integer>()).execute(&conn);
drop(conn); // this will cause a compile error as query future is still referencing conn

2. lock the mutex of transaction manager when pool running it's drop logic.

[weiznich]

Thanks for opening this issue. I think your analysis is mostly correct, but I fear there is just no easy solution to this as both "suggested" solutions won't work

let the query future hold a reference of connection until it's resolved:

That would break support for query pipelining, which is an important performance feature

Lock the mutex of transaction manager when pool running it's drop logic.

We don't have access to the inner working of the pool implementations to know when the drop logic is run, so we cannot lock there.

That all written: I likely won't have the capacity to work on a fix here for the next weeks, I'm happy to review a PR that fixes the problem as long as it does not break existing functionality like pipelining and contains a test case.

[fakeshadow]

Thanks for the correction. I'm working on integrating diesel_async to my db driver and I'll see what I can do after it's finished it.

[r0mi]

Would really appreciate a fix as I've run into the same issue

[weiznich]

@r0mi As written above: I'm happy to review a PR that fixes this.