From 6ac98ef613f067950efdc8596c283d4d0f1bce1e Mon Sep 17 00:00:00 2001
From: Jack Christensen <jack@jackchristensen.com>
Date: Sat, 12 Sep 2015 19:10:41 -0500
Subject: [PATCH] Add test for ConnPool.Begin retry

Also fix bug where conn is not released on non ErrDeadConn error
returned from Conn.Begin
---
 conn_pool.go      |  2 ++
 conn_pool_test.go | 41 +++++++++++++++++++++++++++++++++++++++++
 2 files changed, 43 insertions(+)

diff --git a/conn_pool.go b/conn_pool.go
index 0c5f87f7..20323992 100644
--- a/conn_pool.go
+++ b/conn_pool.go
@@ -226,6 +226,7 @@ func (p *ConnPool) Begin() (*Tx, error) {
 			continue
 		}
 		if err != nil {
+			p.Release(c)
 			return nil, err
 		}
 
@@ -250,6 +251,7 @@ func (p *ConnPool) BeginIso(iso string) (*Tx, error) {
 			continue
 		}
 		if err != nil {
+			p.Release(c)
 			return nil, err
 		}
 
diff --git a/conn_pool_test.go b/conn_pool_test.go
index f6ce75da..b6b7d98a 100644
--- a/conn_pool_test.go
+++ b/conn_pool_test.go
@@ -355,6 +355,47 @@ func TestConnPoolTransactionIso(t *testing.T) {
 	}
 }
 
+func TestConnPoolBeginRetry(t *testing.T) {
+	t.Parallel()
+
+	pool := createConnPool(t, 2)
+	defer pool.Close()
+
+	killerConn, err := pool.Acquire()
+	if err != nil {
+		t.Fatal(err)
+	}
+	defer pool.Release(killerConn)
+
+	victimConn, err := pool.Acquire()
+	if err != nil {
+		t.Fatal(err)
+	}
+	pool.Release(victimConn)
+
+	// Terminate connection that was released to pool
+	if _, err = killerConn.Exec("select pg_terminate_backend($1)", victimConn.Pid); err != nil {
+		t.Fatalf("Unable to kill backend PostgreSQL process: %v", err)
+	}
+
+	// Since victimConn is the only available connection in the pool, pool.Begin should
+	// try to use it, fail, and allocate another connection
+	tx, err := pool.Begin()
+	if err != nil {
+		t.Fatalf("pool.Begin failed: %v", err)
+	}
+	defer tx.Rollback()
+
+	var txPid int32
+	err = tx.QueryRow("select pg_backend_pid()").Scan(&txPid)
+	if err != nil {
+		t.Fatalf("tx.QueryRow Scan failed: %v", err)
+	}
+	if txPid == victimConn.Pid {
+		t.Error("Expected txPid to defer from killed conn pid, but it didn't")
+	}
+}
+
 func TestConnPoolQuery(t *testing.T) {
 	t.Parallel()