From 7fc3eab5b711caaaa312f3ec93b6a0f2c2b2e297 Mon Sep 17 00:00:00 2001
From: 839 <8398a7@gmail.com>
Date: Mon, 14 Dec 2020 22:11:47 +0900
Subject: [PATCH] fix(spanner): fix session leak

Fixes #3460
---
 spanner/client.go      |  8 +++++---
 spanner/client_test.go | 31 +++++++++++++++++++++++++++++++
 2 files changed, 36 insertions(+), 3 deletions(-)

diff --git a/spanner/client.go b/spanner/client.go
index eb5ff9798d36..8f14968e04b4 100644
--- a/spanner/client.go
+++ b/spanner/client.go
@@ -451,6 +451,11 @@ func (c *Client) rwTransaction(ctx context.Context, f func(context.Context, *Rea
 	var (
 		sh *sessionHandle
 	)
+	defer func() {
+		if sh != nil {
+			sh.recycle()
+		}
+	}()
 	err = runWithRetryOnAbortedOrSessionNotFound(ctx, func(ctx context.Context) error {
 		var (
 			err error
@@ -480,9 +485,6 @@ func (c *Client) rwTransaction(ctx context.Context, f func(context.Context, *Rea
 		resp, err = t.runInTransaction(ctx, f)
 		return err
 	})
-	if sh != nil {
-		sh.recycle()
-	}
 	return resp, err
 }
 
diff --git a/spanner/client_test.go b/spanner/client_test.go
index 92cac5564bb8..aa0b32f24671 100644
--- a/spanner/client_test.go
+++ b/spanner/client_test.go
@@ -828,6 +828,37 @@ func TestClient_ReadWriteTransaction_Update_QueryOptions(t *testing.T) {
 	}
 }
 
+func TestClient_ReadWriteTransaction_DoNotLeakSessionOnPanic(t *testing.T) {
+	// Make sure that there is always only one session in the pool.
+	sc := SessionPoolConfig{
+		MinOpened: 1,
+		MaxOpened: 1,
+	}
+	_, client, teardown := setupMockedTestServerWithConfig(t, ClientConfig{SessionPoolConfig: sc})
+	defer teardown()
+	ctx := context.Background()
+
+	// If a panic occurs during a transaction, the session will not leak.
+	func() {
+		defer func() { recover() }()
+
+		_, err := client.ReadWriteTransaction(ctx, func(ctx context.Context, tx *ReadWriteTransaction) error {
+			panic("session leak")
+			return nil
+		})
+		if err != nil {
+			t.Fatalf("Unexpected error during transaction: %v", err)
+		}
+	}()
+
+	_, err := client.ReadWriteTransaction(ctx, func(ctx context.Context, tx *ReadWriteTransaction) error {
+		return nil
+	})
+	if err != nil {
+		t.Fatalf("Unexpected error during transaction: %v", err)
+	}
+}
+
 func TestClient_SessionNotFound(t *testing.T) {
 	// Ensure we always have at least one session in the pool.
 	sc := SessionPoolConfig{