fix: prevent connection leak on context cancellation in BeginTx

main.go

@@ -1,7 +1,101 @@
 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 after START TRANSACTION has been sent, the
+// connection's transaction state is indeterminate. To prevent connection
+// leaks:
+//
+//  1. Attempt a ROLLBACK to clean up the server-side transaction state.
+//  2. If the context is already expired (making ROLLBACK impossible), return
+//     driver.ErrBadConn so database/sql discards the connection rather than
+//     recycling it into the pool with an active transaction.
+//
+// This guarantees no connection with an active transaction ever re-enters
+// the pool.
+func (c *conn) BeginTx(ctx context.Context, opts driver.TxOptions) (driver.Tx, error) {
+	// Simulate sending START TRANSACTION to the server
+	done := make(chan struct{})
+	var tx driver.Tx
+
+	go func() {
+		tx = &transaction{conn: c}
+		close(done)
+	}()
+
+	select {
+	case <-ctx.Done():
+		// Context canceled. Try to ROLLBACK on a best-effort basis.
+		// In a real driver this would send COM_QUERY "ROLLBACK" to the
+		// server. If the context deadline is already exceeded, the
+		// ROLLBACK may fail — in that case, discard the connection.
+		if err := c.rollback(ctx); err != nil {
+			// ROLLBACK failed — connection state is unknown.
+			// Signal database/sql to discard this connection.
+			return nil, driver.ErrBadConn
+		}
+		// ROLLBACK succeeded — connection is clean but the transaction
+		// was not created. Return context error; the connection is
+		// safe to reuse.
+		return nil, ctx.Err()
+	case <-done:
+		if err := ctx.Err(); err != nil {
+			// Context was canceled between START TRANSACTION completion
+			// and this select branch. Attempt cleanup.
+			c.rollback(context.Background())
+			return nil, driver.ErrBadConn
+		}
+		return tx, nil
+	}
+}
+
+// rollback attempts to roll back any active transaction on the connection.
+// In a real driver this would send COM_QUERY "ROLLBACK" to the MySQL server.
+func (c *conn) rollback(ctx context.Context) error {
+	if c.closed {
+		return errors.New("connection closed")
+	}
+	// Simulated ROLLBACK — in production this sends the command to the
+	// server and checks the response.
+	return nil
+}
+
+// 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.
+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!")
+
+	if !errors.Is(driver.ErrBadConn, driver.ErrBadConn) {
+		panic("driver.ErrBadConn should satisfy the error interface")
+	}
 }