From e8dcf5b3ac8aa1569987db5ef0d41476eda3de0c Mon Sep 17 00:00:00 2001
From: Jack Christensen <jack@jackchristensen.com>
Date: Sat, 13 Feb 2016 11:06:44 -0600
Subject: [PATCH] Add *Tx.AfterClose hook

ConnPool now implements its connection with Tx via this hook instead of
manipulating the internals of Tx.
---
 conn_pool.go |  6 +++++-
 tx.go        | 25 +++++++++++++++++++------
 tx_test.go   | 36 ++++++++++++++++++++++++++++++++++++
 3 files changed, 60 insertions(+), 7 deletions(-)

diff --git a/conn_pool.go b/conn_pool.go
index 6eb489e7..f6d5bf74 100644
--- a/conn_pool.go
+++ b/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())
+}
diff --git a/tx.go b/tx.go
index 9f148a4f..4e726d9c 100644
--- a/tx.go
+++ b/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)
+		}
+	}
+}
diff --git a/tx_test.go b/tx_test.go
index 6fb70719..3e76d77c 100644
--- a/tx_test.go
+++ b/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)
+	}
+}