chore: import upstream snapshot with attribution

integration-tests/go-sql-server-driver/sql_server_max_conns_test.go

@@ -0,0 +1,234 @@
+// Copyright 2023 Dolthub, Inc.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package main
+
+import (
+	"context"
+	"database/sql"
+	"testing"
+	"time"
+
+	"github.com/go-sql-driver/mysql"
+	"github.com/stretchr/testify/require"
+	"golang.org/x/sync/errgroup"
+
+	driver "github.com/dolthub/dolt/go/libraries/doltcore/dtestutils/sql_server_driver"
+)
+
+func TestSqlServerMaxConns(t *testing.T) {
+	t.Parallel()
+	t.Run("MaxConns 3", testMaxConns3)
+	t.Run("MaxConns 3 BackLog 0", testMaxConns3BackLog0)
+	t.Run("MaxConns 3 BackLog 1", testMaxConns3BackLog1)
+	t.Run("MaxConns 3 MaxConnectionsTimeout 10s", testMaxConns3Timeout10s)
+}
+
+// setupMaxConnsTest sets up a test "sql-server" with the provided arguments being cli arguments you would pass to
+// an actual server. Then it connects to the DB with 3 connections, simulating connections from 3 different clients.
+// Finally, it creates a two column table, "test_table", for tests to work with.
+//
+// DB and the 3 connections are returned, and the caller can close connections as needed in their test.
+func setupMaxConnsTest(t *testing.T, ctx context.Context, args ...string) (*sql.DB, []*sql.Conn) {
+	t.Parallel()
+	u, err := driver.NewDoltUser()
+	require.NoError(t, err)
+	t.Cleanup(func() {
+		u.Cleanup()
+	})
+	rs, err := u.MakeRepoStore()
+	require.NoError(t, err)
+	repo, err := rs.MakeRepo("max_conns_test")
+	require.NoError(t, err)
+	args = append(args, "--max-connections", "3")
+	args = append(args, "--port", `{{get_port "server"}}`)
+	srvSettings := &driver.Server{
+		Args:        args,
+		DynamicPort: "server",
+	}
+	var ports DynamicResources
+	ports.global = &GlobalPorts
+	ports.t = t
+	server := MakeServer(t, repo, srvSettings, &ports)
+	server.DBName = "max_conns_test"
+	db, err := server.DB(driver.Connection{User: "root"})
+	require.NoError(t, err)
+	t.Cleanup(func() {
+		db.Close()
+	})
+	db.SetMaxIdleConns(0)
+	
+	var conns []*sql.Conn
+	t.Cleanup(func() {
+		closeAll(conns)
+	})
+	for i := 0; i < 3; i++ {
+		conn, err := db.Conn(driver.WithConnectRetriesDisabled(ctx))
+		require.NoError(t, err)
+		conns = append(conns, conn)
+	}
+
+	_, err = conns[0].ExecContext(ctx, `
+CREATE TABLE test_table (
+  id INT AUTO_INCREMENT PRIMARY KEY,
+  str VARCHAR(20)
+);`)
+	require.NoError(t, err)
+
+	return db, conns
+}
+
+func closeAll(conns []*sql.Conn) {
+	for i, c := range conns {
+		if c != nil {
+			c.Close()
+		}
+		conns[i] = nil
+	}
+}
+
+func testMaxConns3BackLog0(t *testing.T) {
+	ctx := context.Background()
+	db, _ := setupMaxConnsTest(t, ctx, "--back-log", "0")
+	if t.Failed() {
+		return
+	}
+	_, err := db.Conn(driver.WithConnectRetriesDisabled(ctx))
+	require.ErrorIs(t, err, mysql.ErrInvalidConn)
+}
+
+func testMaxConns3Timeout10s(t *testing.T) {
+	ctx := context.Background()
+	db, _ := setupMaxConnsTest(t, ctx, "--max-connections-timeout", "10s")
+	if t.Failed() {
+		return
+	}
+	start := time.Now()
+	_, err := db.Conn(driver.WithConnectRetriesDisabled(ctx))
+	elapsed := time.Since(start)
+	require.ErrorIs(t, err, mysql.ErrInvalidConn)
+	require.True(t, elapsed > 9 * time.Second, "it took more than 9 seconds to fail")
+	require.True(t, elapsed < 12 * time.Second, "it took less than 12 seconds to fail")
+}
+
+func testMaxConns3(t *testing.T) {
+	ctx := context.Background()
+	db, conns := setupMaxConnsTest(t, ctx)
+	if t.Failed() {
+		return
+	}
+	eg, ctx := errgroup.WithContext(ctx)
+	eg.Go(func() error {
+		conn, err := db.Conn(driver.WithConnectRetriesDisabled(ctx))
+		if err != nil {
+			return err
+		}
+		defer conn.Close()
+		_, err = conn.ExecContext(ctx, "insert into test_table (str) values ('test4223')")
+		return err
+	})
+	eg.Go(func() error {
+		conn, err := db.Conn(driver.WithConnectRetriesDisabled(ctx))
+		if err != nil {
+			return err
+		}
+		defer conn.Close()
+		_, err = conn.ExecContext(ctx, "insert into test_table (str) values ('test9119')")
+		return err
+	})
+	conns[0].Close()
+	conns[0] = nil
+	require.NoError(t, eg.Wait())
+	ctx = context.Background()
+	rows, err := conns[1].QueryContext(ctx, `SELECT * FROM test_table ORDER BY str ASC`)
+	require.NoError(t, err)
+	defer rows.Close()
+	require.True(t, rows.Next())
+	var id int
+	var str string
+	require.NoError(t, rows.Scan(&id, &str))
+	require.Equal(t, "test4223", str)
+	require.True(t, rows.Next())
+	require.NoError(t, rows.Scan(&id, &str))
+	require.Equal(t, "test9119", str)
+	require.False(t, rows.Next())
+	require.NoError(t, rows.Err())
+}
+
+func testMaxConns3BackLog1(t *testing.T) {
+	ctx := context.Background()
+	db, conns := setupMaxConnsTest(t, ctx, "--back-log", "1")
+	if t.Failed() {
+		return
+	}
+
+	// 3 connections are open, filling all slots. backlog=1 means one
+	// additional connection can queue in the application-level backlog;
+	// the rest are fast-failed. Spawn 4 goroutines all trying to
+	// connect. Exactly 1 should enter the backlog; the other 3 should
+	// get ErrInvalidConn immediately.
+	const numAttempts = 4
+	type result struct {
+		conn *sql.Conn
+		err  error
+	}
+	results := make(chan result, numAttempts)
+	readyCh := make(chan struct{})
+	startCh := make(chan struct{})
+	for range numAttempts {
+		go func() {
+			readyCh <- struct{}{}
+			<-startCh
+			conn, err := db.Conn(driver.WithConnectRetriesDisabled(ctx))
+			results <- result{conn, err}
+		}()
+	}
+	for range numAttempts {
+		<-readyCh
+	}
+	close(startCh)
+
+	// 3 of 4 should be rejected fast (backlog full). Collect those
+	// before freeing a slot, so we don't open space for a second
+	// connection to sneak through.
+	var rejected int
+	for range numAttempts - 1 {
+		r := <-results
+		require.ErrorIs(t, r.err, mysql.ErrInvalidConn)
+		rejected++
+	}
+
+	// Now free a slot so the one goroutine waiting in the backlog can proceed.
+	conns[0].Close()
+	conns[0] = nil
+
+	r := <-results
+	require.NoError(t, r.err, "expected the backlogged connection to succeed")
+	require.Equal(t, numAttempts-1, rejected)
+
+	_, err := r.conn.ExecContext(ctx, "insert into test_table (str) values ('test4223')")
+	require.NoError(t, err)
+	r.conn.Close()
+
+	rows, err := conns[1].QueryContext(ctx, `SELECT * FROM test_table`)
+	require.NoError(t, err)
+	defer rows.Close()
+	require.True(t, rows.Next())
+	var id int
+	var str string
+	require.NoError(t, rows.Scan(&id, &str))
+	require.Equal(t, "test4223", str)
+	require.False(t, rows.Next())
+	require.NoError(t, rows.Err())
+}