From 17d0159c8f76d43d7f36b200c5da6573e5db9615 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Knut=20Olav=20L=C3=B8ite?= <koloite@gmail.com>
Date: Fri, 14 Aug 2020 11:57:22 +0200
Subject: [PATCH] fix(spanner): call ctx.cancel after stats have been recorded
 (#2728)

The cancel() method should be called once all code that is executing in the given
context have finished. In this case, that also means after the stats recording that
uses the same context has finished.

Fixes #2660
---
 spanner/session.go | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/spanner/session.go b/spanner/session.go
index f83c7ff26fc0..f89357259294 100644
--- a/spanner/session.go
+++ b/spanner/session.go
@@ -1501,7 +1501,6 @@ func (hc *healthChecker) worker(i int) {
 		if ws != nil {
 			ctx, cancel := context.WithTimeout(context.Background(), time.Minute)
 			err := ws.prepareForWrite(ctx)
-			cancel()
 			if err != nil {
 				// Skip handling prepare error, session can be prepared in next
 				// cycle.
@@ -1516,6 +1515,7 @@ func (hc *healthChecker) worker(i int) {
 			hc.pool.mu.Lock()
 			hc.pool.decNumBeingPreparedLocked(ctx)
 			hc.pool.mu.Unlock()
+			cancel()
 			hc.markDone(ws)
 		}
 		rs := getNextForPing()