Add *Tx.AfterClose hook

ConnPool now implements its connection with Tx via this hook instead of
manipulating the internals of Tx.

conn_pool.go

@@ -287,7 +287,11 @@ func (p *ConnPool) BeginIso(iso string) (*Tx, error) {
 			}
 		}
 
-		tx.pool = p
+		tx.AfterClose(p.txAfterClose)
 		return tx, nil
 	}
 }
+
+func (p *ConnPool) txAfterClose(tx *Tx) {
+	p.Release(tx.Conn())
+}

tx.go

@@ -54,9 +54,9 @@ func (c *Conn) begin(isoLevel string) (*Tx, error) {
 // All Tx methods return ErrTxClosed if Commit or Rollback has already been
 // called on the Tx.
 type Tx struct {
-	pool   *ConnPool
-	conn   *Conn
-	closed bool
+	conn       *Conn
+	afterClose func(*Tx)
+	closed     bool
 }
 
 // Commit commits the transaction
@@ -85,9 +85,8 @@ func (tx *Tx) Rollback() error {
 }
 
 func (tx *Tx) close() {
-	if tx.pool != nil {
-		tx.pool.Release(tx.conn)
-		tx.pool = nil
+	if tx.afterClose != nil {
+		tx.afterClose(tx)
 	}
 	tx.closed = true
 }
@@ -117,3 +116,17 @@ func (tx *Tx) QueryRow(sql string, args ...interface{}) *Row {
 	rows, _ := tx.Query(sql, args...)
 	return (*Row)(rows)
 }
+
+// AfterClose adds f to a LILO queue of functions that will be called when
+// the transaction is closed (either Commit or Rollback).
+func (tx *Tx) AfterClose(f func(*Tx)) {
+	if tx.afterClose == nil {
+		tx.afterClose = f
+	} else {
+		prevFn := tx.afterClose
+		tx.afterClose = func(tx *Tx) {
+			f(tx)
+			prevFn(tx)
+		}
+	}
+}

tx_test.go

@@ -3,6 +3,7 @@ package pgx_test
 import (
 	"github.com/jackc/pgx"
 	"testing"
+	"time"
 )
 
 func TestTransactionSuccessfulCommit(t *testing.T) {
@@ -114,3 +115,38 @@ func TestBeginIso(t *testing.T) {
 		}
 	}
 }
+
+func TestTxAfterClose(t *testing.T) {
+	t.Parallel()
+
+	conn := mustConnect(t, *defaultConnConfig)
+	defer closeConn(t, conn)
+
+	tx, err := conn.Begin()
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	var zeroTime, t1, t2 time.Time
+	tx.AfterClose(func(tx *pgx.Tx) {
+		t1 = time.Now()
+	})
+
+	tx.AfterClose(func(tx *pgx.Tx) {
+		t2 = time.Now()
+	})
+
+	tx.Rollback()
+
+	if t1 == zeroTime {
+		t.Error("First Tx.AfterClose callback not called")
+	}
+
+	if t2 == zeroTime {
+		t.Error("Second Tx.AfterClose callback not called")
+	}
+
+	if t1.After(t2) {
+		t.Error("AfterClose callbacks called out of order: %v, %v", t1, t2)
+	}
+}