From 82485dcd054707d6e8fac379ac36099013ffb7cf Mon Sep 17 00:00:00 2001
From: Morgan Tocker <mtocker@squareup.com>
Date: Wed, 11 Jan 2023 13:22:30 -0700
Subject: [PATCH] Add lock_wait_timeout as retryable

---
 errors.go      | 44 ++++++++++++++++-------------
 errors_test.go | 75 ++++++++++++++++++++++++++++++++++++++++++++++++++
 2 files changed, 100 insertions(+), 19 deletions(-)

diff --git a/errors.go b/errors.go
index be7b3b8..8f63a87 100644
--- a/errors.go
+++ b/errors.go
@@ -19,25 +19,25 @@
 //
 // Error and CanRetry are the most important functions in this package. They are used like:
 //
-//   import my "github.com/go-mysql/errors"
+//	import my "github.com/go-mysql/errors"
 //
-//   SAVE_ITEM_LOOP:
-//   for tryNo := 1; tryNo <= maxTries; tryNo++ {
-//       if err := SaveItem(item); err != nil {
-//           if ok, myerr := my.Error(err); ok { // MySQL error
-//               if myerr == my.ErrDupeKey {
-//                   return http.StatusConflict
-//               }
-//               if my.CanRetry(myerr) {
-//                   time.Sleep(1 * time.Second)
-//                   continue SAVE_ITEM_LOOP // retry
-//               }
-//           }
-//           // Error not handled
-//           return http.StatusInternalServerError
-//       }
-//       break SAVE_ITEM_LOOP // success
-//   }
+//	SAVE_ITEM_LOOP:
+//	for tryNo := 1; tryNo <= maxTries; tryNo++ {
+//	    if err := SaveItem(item); err != nil {
+//	        if ok, myerr := my.Error(err); ok { // MySQL error
+//	            if myerr == my.ErrDupeKey {
+//	                return http.StatusConflict
+//	            }
+//	            if my.CanRetry(myerr) {
+//	                time.Sleep(1 * time.Second)
+//	                continue SAVE_ITEM_LOOP // retry
+//	            }
+//	        }
+//	        // Error not handled
+//	        return http.StatusInternalServerError
+//	    }
+//	    break SAVE_ITEM_LOOP // success
+//	}
 //
 // The example above tries N-many times to save an item into a database.
 // It handles MySQL errors explicitly. On ErrDupeKey it does not retry, it
@@ -83,6 +83,10 @@ var (
 	// ErrDupeKey is returned when a unique index prevents a value from being
 	// inserted or updated. CanRetry returns false on this error.
 	ErrDupeKey = errors.New("duplicate key value")
+
+	// ErrLockWaitTimeout is returned when innodb_lock_wait_timeout and a lock
+	// could not be acquired. CanRetry returns true on this error.
+	ErrLockWaitTimeout = errors.New("lock wait timeout")
 )
 
 // Error returns an error in this package if possible. The boolean return
@@ -113,6 +117,8 @@ func Error(err error) (bool, error) {
 		return true, ErrReadOnly
 	case 1062: // ER_DUP_ENTRY
 		return true, ErrDupeKey
+	case 1205: // ER_LOCK_WAIT_TIMEOUT
+		return true, ErrLockWaitTimeout
 	}
 
 	// A MySQL error, but not one we handle explicitly
@@ -160,7 +166,7 @@ func MySQLErrorCode(err error) uint16 {
 // It returns false for all other errors, including nil.
 func CanRetry(err error) bool {
 	switch err {
-	case ErrCannotConnect, ErrConnLost, ErrReadOnly, ErrQueryKilled:
+	case ErrCannotConnect, ErrConnLost, ErrReadOnly, ErrQueryKilled, ErrLockWaitTimeout:
 		return true
 	}
 	return false
diff --git a/errors_test.go b/errors_test.go
index 1b8880c..11f5adc 100644
--- a/errors_test.go
+++ b/errors_test.go
@@ -503,4 +503,79 @@ func TestCanRetry(t *testing.T) {
 	if retry == false {
 		t.Error("can try = false, expected true")
 	}
+	retry = my.CanRetry(my.ErrLockWaitTimeout)
+	if retry == false {
+		t.Error("can try = false, expected true")
+	}
+}
+
+func TestLockWaitTimeout(t *testing.T) {
+	setup(t)
+
+	// Make a connection as usual
+	db := newDB(t, defaultDSN)
+	ctx := context.TODO()
+	conn, err := db.Conn(ctx)
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	if conn == nil {
+		t.Fatal("got nil conn, expected it to be set")
+	}
+	defer conn.Close()
+
+	// Make sure to lower timeouts.
+	_, err = conn.ExecContext(ctx, "SET GLOBAL innodb_lock_wait_timeout = 1")
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	// Create a table with no data.
+	_, err = conn.ExecContext(ctx, "CREATE TABLE test.lock_wait_timeout (id INT NOT NULL PRIMARY KEY, col1 varchar(255))")
+	if err != nil {
+		t.Fatal(err)
+	}
+	_, err = conn.ExecContext(ctx, "INSERT INTO test.lock_wait_timeout VALUES (1, 'test')")
+	if err != nil {
+		t.Fatal(err)
+	}
+	defer func() {
+		_, err := conn.ExecContext(ctx, "DROP TABLE IF EXISTS test.lock_wait_timeout")
+		if err != nil {
+			t.Errorf("cannot drop table test.t: %s", err)
+		}
+	}()
+
+	// lock row 1
+	trx, err := conn.BeginTx(ctx, &sql.TxOptions{
+		Isolation: sql.LevelRepeatableRead,
+	})
+	if err != nil {
+		t.Fatal(err)
+	}
+	_, err = trx.Exec("SELECT * FROM test.lock_wait_timeout WHERE id = 1 FOR UPDATE")
+	if err != nil {
+		t.Fatal(err)
+	}
+	// In another session, try and lock ROW 1
+	// It should timeout.
+	conn2, err := db.Conn(ctx)
+	if err != nil {
+		t.Fatal(err)
+	}
+	defer conn2.Close()
+	_, err = conn2.ExecContext(ctx, "SELECT * FROM test.lock_wait_timeout WHERE id = 1 FOR UPDATE")
+	t.Logf("err: %v", err)
+	ok, myerr := my.Error(err)
+	if !ok {
+		t.Error("MySQL error = false, expected true)")
+	}
+	if myerr != my.ErrLockWaitTimeout {
+		t.Errorf("got error '%v', expected ErrLockWaitTimeout", err)
+	}
+	// Rollback the transaction
+	if err := trx.Rollback(); err != nil {
+		t.Fatal(err)
+	}
 }