From 78a9f80ea72d2e487fc132e275e934b5d824f704 Mon Sep 17 00:00:00 2001
From: d10i <hello@hexago.tech>
Date: Fri, 2 Aug 2024 11:52:58 +0200
Subject: [PATCH] Fix sqlserver lock

When acquiring the `sqlserver` lock, it fails immediately if the lock
can't be acquired. When running multiple migrations concurrently, only
one of them will succeed and the other ones will fail.

Wait for up to 10 seconds when acquiring the lock. This way it lets the
other entity that has acquired the lock do its job first. This is
consistent with what this library does for `postgres` (though it waits
indefinitely) and for `mysql`.
---
 database/sqlserver/sqlserver.go      |  6 +-
 database/sqlserver/sqlserver_test.go | 86 ++++++++++++++++++++++++++++
 2 files changed, 89 insertions(+), 3 deletions(-)

diff --git a/database/sqlserver/sqlserver.go b/database/sqlserver/sqlserver.go
index 3cfa48bf9..6f3256fef 100644
--- a/database/sqlserver/sqlserver.go
+++ b/database/sqlserver/sqlserver.go
@@ -198,10 +198,10 @@ func (ss *SQLServer) Lock() error {
 			return err
 		}
 
-		// This will either obtain the lock immediately and return true,
-		// or return false if the lock cannot be acquired immediately.
+		// This will either obtain the lock within 10 seconds and return true,
+		// or return false if the lock cannot be acquired within 10 seconds.
 		// MS Docs: sp_getapplock: https://docs.microsoft.com/en-us/sql/relational-databases/system-stored-procedures/sp-getapplock-transact-sql?view=sql-server-2017
-		query := `EXEC sp_getapplock @Resource = @p1, @LockMode = 'Update', @LockOwner = 'Session', @LockTimeout = 0`
+		query := `EXEC sp_getapplock @Resource = @p1, @LockMode = 'Update', @LockOwner = 'Session', @LockTimeout = 10000`
 
 		var status mssql.ReturnStatus
 		if _, err = ss.conn.ExecContext(context.Background(), query, aid, &status); err == nil && status > -1 {
diff --git a/database/sqlserver/sqlserver_test.go b/database/sqlserver/sqlserver_test.go
index e7af5d7f9..96704ae36 100644
--- a/database/sqlserver/sqlserver_test.go
+++ b/database/sqlserver/sqlserver_test.go
@@ -8,6 +8,7 @@ import (
 	"log"
 	"runtime"
 	"strings"
+	"sync"
 	"testing"
 	"time"
 
@@ -104,6 +105,8 @@ func Test(t *testing.T) {
 	t.Run("testMsiTrue", testMsiTrue)
 	t.Run("testOpenWithPasswordAndMSI", testOpenWithPasswordAndMSI)
 	t.Run("testMsiFalse", testMsiFalse)
+	t.Run("testLock", testLock)
+	t.Run("testWithInstanceConcurrent", testWithInstanceConcurrent)
 
 	t.Cleanup(func() {
 		for _, spec := range specs {
@@ -339,3 +342,86 @@ func testMsiFalse(t *testing.T) {
 		}
 	})
 }
+
+func testLock(t *testing.T) {
+	dktesting.ParallelTest(t, specs, func(t *testing.T, c dktest.ContainerInfo) {
+		ip, port, err := c.FirstPort()
+		if err != nil {
+			t.Fatal(err)
+		}
+
+		addr := msConnectionString(ip, port)
+		p := &SQLServer{}
+		d, err := p.Open(addr)
+		if err != nil {
+			t.Fatal(err)
+		}
+
+		dt.Test(t, d, []byte("SELECT 1"))
+
+		ps := d.(*SQLServer)
+
+		err = ps.Lock()
+		if err != nil {
+			t.Fatal(err)
+		}
+
+		err = ps.Unlock()
+		if err != nil {
+			t.Fatal(err)
+		}
+
+		err = ps.Lock()
+		if err != nil {
+			t.Fatal(err)
+		}
+
+		err = ps.Unlock()
+		if err != nil {
+			t.Fatal(err)
+		}
+	})
+}
+
+func testWithInstanceConcurrent(t *testing.T) {
+	dktesting.ParallelTest(t, specs, func(t *testing.T, c dktest.ContainerInfo) {
+		ip, port, err := c.FirstPort()
+		if err != nil {
+			t.Fatal(err)
+		}
+
+		// The number of concurrent processes running WithInstance
+		const concurrency = 30
+
+		// We can instantiate a single database handle because it is
+		// actually a connection pool, and so, each of the below go
+		// routines will have a high probability of using a separate
+		// connection, which is something we want to exercise.
+		db, err := sql.Open("sqlserver", msConnectionString(ip, port))
+		if err != nil {
+			t.Fatal(err)
+		}
+		defer func() {
+			if err := db.Close(); err != nil {
+				t.Error(err)
+			}
+		}()
+
+		db.SetMaxIdleConns(concurrency)
+		db.SetMaxOpenConns(concurrency)
+
+		var wg sync.WaitGroup
+		defer wg.Wait()
+
+		wg.Add(concurrency)
+		for i := 0; i < concurrency; i++ {
+			go func(i int) {
+				defer wg.Done()
+				_, err := WithInstance(db, &Config{})
+				if err != nil {
+					t.Errorf("process %d error: %s", i, err)
+				}
+			}(i)
+		}
+	})
+}