TOC/TOU bug on transaction rollback

[zepatrik]

Describe the bug

I have observed flaky behavior in a test that causes a transaction to be rolled back by canceling the context. I think it is caused by a toctou issue here:

pgx/tx.go

Lines 210 to 215 in c00cd65

	if tx.closed {
	return ErrTxClosed
	}
	_, err := tx.conn.Exec(ctx, "rollback")
	tx.closed = true

To Reproduce

It is a hard to reproduce issue without lots of setup, this is what is happening roughly:

func TestRollbackOnContextCanceled(t *testing.T) {
	dsn := dockertest.RunTestPostgreSQL(t) // or simply load from the env

	db, err := sql.Open("pgx/v5", dsn)
	require.NoError(t, err)
	t.Cleanup(func() { _ = db.Close() })
	require.EventuallyWithT(t, func(collect *assert.CollectT) {
		assert.NoError(collect, db.PingContext(t.Context()))
	}, 10*time.Second, 100*time.Millisecond)

	ctx, cancel := context.WithCancel(t.Context())
	t.Cleanup(cancel)

	tx, err := db.BeginTx(ctx, nil)
	require.NoError(t, err)

	_, err = tx.ExecContext(ctx, "SELECT 1")
	require.NoError(t, err)

	cancel()
	time.Sleep(1 * time.Nanosecond) // give the database/sql auto-rollback some time to kick in

	// this races now with the database/sql auto-rollback on context cancellation
	// https://github.com/golang/go/blob/07840ceeed4afd10324a552e8c87a8ee363aa24a/src/database/sql/sql.go#L2211-L2226
	// depending on the outcome, we either get sql.ErrTxDone or "conn closed"
	require.NoError(t, tx.Rollback())
}

Expected behavior

The close bool flag should be managed atomically, similar to how it is done in the database/sql package: https://github.com/golang/go/blob/07840ceeed4afd10324a552e8c87a8ee363aa24a/src/database/sql/sql.go#L2331

Actual behavior

Flaky behavior.

Version

• Go: $ go version -> go version go1.26.2 darwin/arm64
• PostgreSQL: $ psql --no-psqlrc --tuples-only -c 'select version()' -> PostgreSQL 16.13 (Debian 16.13-1.pgdg13+1) on aarch64-unknown-linux-gnu, compiled by gcc (Debian 14.2.0-19) 14.2.0, 64-bit
• pgx: $ grep 'github.com/jackc/pgx/v[0-9]' go.mod -> v5.9.2

Additional context

...

[jackc]

I'm not sure this is a pgx issue. The underlying pgx code is not concurrency safe nor intended to be. Both the rollback triggered by the context cancellation and the tx.Rollback call are database/sql concerns. If there is a synchronization issue I would expect it to be handled there.

That said, there is probably another underlying issue. In many cases, pgx will close a connection when there is a context cancellation. Depending on how database/sql plumbs that canceled ctx through the system it may be that pgx has closed the connection before database/sql has successfully issued the rollback. That might account for what your seeing.

In either case I don't think #2552 is the appropriate answer.

[zepatrik]

The underlying pgx code is not concurrency safe nor intended to be.

I understand this part, and it totally makes sense, but I'm not sure if there is any chance we can get this auto-rollback removed in the stdlib.

there is probably another underlying issue

What I see is that there are two rollbacks happening at the same time, and then I get the conn closed error from the wire protocol directly, instead of ErrTxClosed.

In either case I don't think #2552 is the appropriate answer.

The fix is just replicating what the stdlib tx does, so I think it is something that they already ran into and fixed this way. The code in question is unconditionally setting tx.closed = true, imo there should be no reason to not do this immediately instead of keeping that race window open.

I'm not sure this is a pgx issue.

I have one test which is reliably running into this flake, and it is only happening with postgres/pgx, other database/driver combinations are not running into any issue. While that is not evidence, it is still something that made me look closer at pgx and why I noticed the smelly-looking code.

[zepatrik]

After quite some time of late-night debugging, the issue seems to be something completely different from what I initially thought, not caused by the stdlib auto-rollback go routine. I opened #2557 for that, but I still think that it makes sense to use an atomic state for the transaction.
If someone would have a go routine that does the same thing as the stdlib's auto-rollback one, they would run into this TOC/TOU bug, and I don't think that this is an unreasonable use-case.
The stdlib transaction also does not support concurrency in general, but still uses the atomic closed bool to ensure the background rollback does not run into race-conditions.

[jackc]

If someone would have a go routine that does the same thing as the stdlib's auto-rollback one, they would run into this TOC/TOU bug, and I don't think that this is an unreasonable use-case.

pgx.Conn is not concurrency safe as a whole is not concurrency safe andRollback() from a different goroutine is not allowed. Making closed atomic would not change that. Rollback() would still not be concurrency safe.