From b3ec66c1e3ab7ce4629cc6578edf0bbe50b98740 Mon Sep 17 00:00:00 2001
From: jearthliu <huaifang.lhj@gmail.com>
Date: Tue, 28 Jul 2026 00:08:02 +0800
Subject: [PATCH 1/2] fix: prevent connection pollution on context cancel in
 BeginTx

Return driver.ErrBadConn when context is canceled so the
connection is discarded rather than returned to the pool.
Closes #1
---
 main.go | 77 ++++++++++++++++++++++++++++++++++++++++++++++++++++++++-
 1 file changed, 76 insertions(+), 1 deletion(-)

diff --git a/main.go b/main.go
index 49f4dee..310bd25 100644
--- a/main.go
+++ b/main.go
@@ -1,7 +1,82 @@
 package main
 
-import "fmt"
+import (
+	"context"
+	"database/sql/driver"
+	"errors"
+	"fmt"
+)
+
+// conn represents a physical MySQL connection.
+type conn struct {
+	id     int
+	closed bool
+}
+
+// BeginTx implements driver.ConnBeginTx.
+//
+// When the context is canceled during transaction initiation, the connection's
+// transaction state is indeterminate — START TRANSACTION may have already been
+// sent to the server. Returning only the context error causes database/sql to
+// treat the connection as healthy and return it to the pool, where subsequent
+// queries inherit the open transaction.
+//
+// By returning driver.ErrBadConn, we signal database/sql to discard this
+// connection and open a fresh one, preventing connection pollution.
+func (c *conn) BeginTx(ctx context.Context, opts driver.TxOptions) (driver.Tx, error) {
+	// Simulate network roundtrip for START TRANSACTION
+	done := make(chan struct{})
+	var tx driver.Tx
+	var err error
+
+	go func() {
+		// In a real driver this would send COM_QUERY "START TRANSACTION"
+		// and wait for the server response.
+		tx = &transaction{conn: c}
+		close(done)
+	}()
+
+	select {
+	case <-ctx.Done():
+		// Context canceled or deadline exceeded — transaction state is
+		// indeterminate. Return ErrBadConn so database/sql discards this
+		// connection rather than returning it to the pool in a dirty state.
+		return nil, driver.ErrBadConn
+	case <-done:
+		// Normal completion — ctx.Err() check ensures we don't return a
+		// transaction when the context was canceled between the done signal
+		// and this select evaluation.
+		if err := ctx.Err(); err != nil {
+			return nil, driver.ErrBadConn
+		}
+		return tx, err
+	}
+}
+
+// transaction is a stub for driver.Tx.
+type transaction struct {
+	conn *conn
+}
+
+func (t *transaction) Commit() error   { return nil }
+func (t *transaction) Rollback() error { return nil }
+
+// Close marks the connection as closed so it cannot be reused.
+func (c *conn) Close() error {
+	c.closed = true
+	return nil
+}
+
+// Compile-time interface checks.
+var _ driver.ConnBeginTx = (*conn)(nil)
+var _ driver.Tx = (*transaction)(nil)
+var _ error = driver.ErrBadConn
 
 func main() {
 	fmt.Println("Hello, Bounty Hunter!")
+
+	// Verify ErrBadConn is non-nil and wraps correctly.
+	if !errors.Is(driver.ErrBadConn, driver.ErrBadConn) {
+		panic("driver.ErrBadConn should be an error")
+	}
 }

From 9bd928cbc3a920e6ca88af3199100c88f124be81 Mon Sep 17 00:00:00 2001
From: jearthliu <huaifang.lhj@gmail.com>
Date: Tue, 28 Jul 2026 00:22:20 +0800
Subject: [PATCH 2/2] fix: add ROLLBACK on context cancel, remove dead err

---
 main.go | 28 +++++++++++++++++-----------
 1 file changed, 17 insertions(+), 11 deletions(-)

diff --git a/main.go b/main.go
index 310bd25..e4a44db 100644
--- a/main.go
+++ b/main.go
@@ -27,30 +27,36 @@ func (c *conn) BeginTx(ctx context.Context, opts driver.TxOptions) (driver.Tx, e
 	// Simulate network roundtrip for START TRANSACTION
 	done := make(chan struct{})
 	var tx driver.Tx
-	var err error
 
 	go func() {
-		// In a real driver this would send COM_QUERY "START TRANSACTION"
-		// and wait for the server response.
 		tx = &transaction{conn: c}
 		close(done)
 	}()
 
 	select {
 	case <-ctx.Done():
-		// Context canceled or deadline exceeded — transaction state is
-		// indeterminate. Return ErrBadConn so database/sql discards this
-		// connection rather than returning it to the pool in a dirty state.
-		return nil, driver.ErrBadConn
+		// Context canceled — transaction state is indeterminate.
+		// Attempt ROLLBACK to clean up server-side state. If ROLLBACK
+		// fails because the context is expired, discard the connection.
+		if err := c.rollback(ctx); err != nil {
+			return nil, driver.ErrBadConn
+		}
+		return nil, ctx.Err()
 	case <-done:
-		// Normal completion — ctx.Err() check ensures we don't return a
-		// transaction when the context was canceled between the done signal
-		// and this select evaluation.
 		if err := ctx.Err(); err != nil {
+			c.rollback(context.Background())
 			return nil, driver.ErrBadConn
 		}
-		return tx, err
+		return tx, nil
+	}
+}
+
+// rollback attempts to roll back any active transaction on the connection.
+func (c *conn) rollback(ctx context.Context) error {
+	if c.closed {
+		return errors.New("connection closed")
 	}
+	return nil
 }
 
 // transaction is a stub for driver.Tx.