From 7fe7f33557938739e5342d82d0720523c344eb71 Mon Sep 17 00:00:00 2001
From: "Andrew S. Brown" <andrew@launchdarkly.com>
Date: Sun, 4 Aug 2019 15:31:32 -0700
Subject: [PATCH] Terminate context prior to releasing when killing batch
 connection

---
 batch.go | 11 +++++++----
 1 file changed, 7 insertions(+), 4 deletions(-)

diff --git a/batch.go b/batch.go
index 4b624387..8c924e8d 100644
--- a/batch.go
+++ b/batch.go
@@ -135,7 +135,7 @@ func (b *Batch) Send(ctx context.Context, txOptions *TxOptions) error {
 
 	_, err = b.conn.conn.Write(buf)
 	if err != nil {
-		b.conn.die(err)
+		b.die(err)
 		return err
 	}
 
@@ -281,10 +281,13 @@ func (b *Batch) die(err error) {
 	}
 
 	b.err = err
-	b.conn.die(err)
+	if b.conn != nil {
+		err = b.conn.termContext(err)
+		b.conn.die(err)
 
-	if b.conn != nil && b.connPool != nil {
-		b.connPool.Release(b.conn)
+		if b.connPool != nil {
+			b.connPool.Release(b.conn)
+		}
 	}
 }