Improved Connection Pooling (#14034)

Signed-off-by: Vicent Marti <vmg@strn.cat>

go.mod

@@ -109,6 +109,7 @@ require (
 	go.uber.org/goleak v1.2.1
 	golang.org/x/exp v0.0.0-20230817173708-d852ddb80c63
 	golang.org/x/sync v0.3.0
+	gonum.org/v1/gonum v0.14.0
 	modernc.org/sqlite v1.20.3
 )
 

go.sum

@@ -944,6 +944,8 @@ golang.org/x/xerrors v0.0.0-20191204190536-9bdfabe68543/go.mod h1:I/5z698sn9Ka8T
 golang.org/x/xerrors v0.0.0-20200804184101-5ec99f83aff1/go.mod h1:I/5z698sn9Ka8TeJc9MKroUUfqBBauWjQqLJ2OPfmY0=
 golang.org/x/xerrors v0.0.0-20220907171357-04be3eba64a2 h1:H2TDz8ibqkAF6YGhCdN3jS9O0/s90v0rJh3X/OLHEUk=
 golang.org/x/xerrors v0.0.0-20220907171357-04be3eba64a2/go.mod h1:K8+ghG5WaK9qNqU5K3HdILfMLy1f3aNYFI/wnl100a8=
+gonum.org/v1/gonum v0.14.0 h1:2NiG67LD1tEH0D7kM+ps2V+fXmsAnpUeec7n8tcr4S0=
+gonum.org/v1/gonum v0.14.0/go.mod h1:AoWeoz0becf9QMWtE8iWXNXc27fK4fNeHNf/oMejGfU=
 google.golang.org/api v0.4.0/go.mod h1:8k5glujaEP+g9n7WNsDg8QP6cUVNI86fCNMcbazEtwE=
 google.golang.org/api v0.7.0/go.mod h1:WtwebWUNSVBH/HAw79HIFXZNqEvBhG+Ra+ax0hx3E3M=
 google.golang.org/api v0.8.0/go.mod h1:o4eAsZoiT+ibD93RtjEohWalFOjRDx6CVaqeizhEnKg=

go/atomic2/atomic128.go

@@ -0,0 +1,63 @@
+//go:build amd64 || arm64
+
+/*
+Copyright 2023 The Vitess Authors.
+
+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 atomic2
+
+import (
+	"unsafe"
+)
+
+//go:linkname writeBarrier runtime.writeBarrier
+var writeBarrier struct {
+	enabled bool    // compiler emits a check of this before calling write barrier
+	pad     [3]byte // compiler uses 32-bit load for "enabled" field
+	needed  bool    // identical to enabled, for now (TODO: dedup)
+	alignme uint64  // guarantee alignment so that compiler can use a 32 or 64-bit load
+}
+
+//go:linkname atomicwb runtime.atomicwb
+//go:nosplit
+func atomicwb(ptr *unsafe.Pointer, new unsafe.Pointer)
+
+type PointerAndUint64[T any] struct {
+	p unsafe.Pointer
+	u uint64
+}
+
+//go:nosplit
+func loadUint128_(addr *unsafe.Pointer) (pp unsafe.Pointer, uu uint64)
+
+func (x *PointerAndUint64[T]) Load() (*T, uint64) {
+	p, u := loadUint128_(&x.p)
+	return (*T)(p), u
+}
+
+//go:nosplit
+func compareAndSwapUint128_(addr *unsafe.Pointer, oldp unsafe.Pointer, oldu uint64, newp unsafe.Pointer, newu uint64) (swapped bool)
+
+//go:nosplit
+func compareAndSwapUint128(addr *unsafe.Pointer, oldp unsafe.Pointer, oldu uint64, newp unsafe.Pointer, newu uint64) bool {
+	if writeBarrier.enabled {
+		atomicwb(addr, newp)
+	}
+	return compareAndSwapUint128_(addr, oldp, oldu, newp, newu)
+}
+
+func (x *PointerAndUint64[T]) CompareAndSwap(oldp *T, oldu uint64, newp *T, newu uint64) bool {
+	return compareAndSwapUint128(&x.p, unsafe.Pointer(oldp), oldu, unsafe.Pointer(newp), newu)
+}

go/atomic2/atomic128_amd64.s

@@ -0,0 +1,46 @@
+// Copyright 2023 The Vitess Authors.
+// Copyright (c) 2021, Carlo Alberto Ferraris
+// Copyright (c) 2017, Tom Thorogood
+//
+// 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.
+
+// Use of this source code is governed by a
+// Modified BSD License that can be found in
+// the LICENSE file.
+
+//+build !noasm,!appengine
+
+#include "textflag.h"
+
+TEXT ·compareAndSwapUint128_(SB), NOSPLIT, $0-41
+	MOVQ addr+0(FP), R8
+	MOVQ oldp+8(FP), AX
+	MOVQ oldu+16(FP), DX
+	MOVQ newp+24(FP), BX
+	MOVQ newu+32(FP), CX
+	LOCK
+	CMPXCHG16B (R8)
+	SETEQ swapped+40(FP)
+	RET
+
+TEXT ·loadUint128_(SB), NOSPLIT, $0-24
+	MOVQ addr+0(FP), R8
+	XORQ AX, AX
+	XORQ DX, DX
+	XORQ BX, BX
+	XORQ CX, CX
+	LOCK
+	CMPXCHG16B (R8)
+	MOVQ AX, pp+8(FP)
+	MOVQ DX, uu+16(FP)
+	RET

go/atomic2/atomic128_arm64.s

@@ -0,0 +1,39 @@
+// Copyright 2023 The Vitess Authors.
+//
+// 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.
+
+//+build !noasm,!appengine
+
+#include "textflag.h"
+
+TEXT ·compareAndSwapUint128_(SB), NOSPLIT, $0-41
+	MOVD addr+0(FP), R5
+	MOVD oldp+8(FP), R0
+	MOVD oldu+16(FP), R1
+	MOVD newp+24(FP), R2
+	MOVD newu+32(FP), R3
+	MOVD R0, R6
+	MOVD R1, R7
+	CASPD (R0, R1), (R5), (R2, R3)
+	CMP R0, R6
+	CCMP EQ, R1, R7, $0
+	CSET EQ, R0
+	MOVB R0, ret+40(FP)
+	RET
+
+TEXT ·loadUint128_(SB), NOSPLIT, $0-24
+	MOVD addr+0(FP), R3
+	LDAXP (R3), (R0, R1)
+	MOVD R0, val+8(FP)
+	MOVD R1, val+16(FP)
+	RET

go/atomic2/atomic128_spinlock.go

@@ -0,0 +1,61 @@
+//go:build !amd64 && !arm64
+
+/*
+Copyright 2023 The Vitess Authors.
+
+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 atomic2
+
+import (
+	"runtime"
+	"sync/atomic"
+)
+
+type PointerAndUint64[T any] struct {
+	spin atomic.Uint64
+	p    *T
+	u    uint64
+}
+
+func (x *PointerAndUint64[T]) Store(p *T, u uint64) {
+	for !x.spin.CompareAndSwap(0, 1) {
+		runtime.Gosched()
+	}
+	defer x.spin.Store(0)
+	x.p = p
+	x.u = u
+}
+
+func (x *PointerAndUint64[T]) Load() (*T, uint64) {
+	for !x.spin.CompareAndSwap(0, 1) {
+		runtime.Gosched()
+	}
+	defer x.spin.Store(0)
+	return x.p, x.u
+}
+
+func (x *PointerAndUint64[T]) CompareAndSwap(oldp *T, oldu uint64, newp *T, newu uint64) bool {
+	for !x.spin.CompareAndSwap(0, 1) {
+		runtime.Gosched()
+	}
+	defer x.spin.Store(0)
+
+	if x.p == oldp && x.u == oldu {
+		x.p = newp
+		x.u = newu
+		return true
+	}
+	return false
+}

go/atomic2/atomic128_test.go

@@ -0,0 +1,44 @@
+/*
+Copyright 2023 The Vitess Authors.
+
+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 atomic2
+
+import (
+	"testing"
+	"unsafe"
+
+	"github.com/stretchr/testify/require"
+)
+
+func TestCompareAndSwap(t *testing.T) {
+	i1 := new(int)
+	i2 := new(int)
+	n := &PointerAndUint64[int]{p: unsafe.Pointer(i1), u: 12345}
+
+	ok := n.CompareAndSwap(i1, 12345, i2, 67890)
+	require.Truef(t, ok, "unexpected CAS failure")
+
+	pp, uu := n.Load()
+	require.Equal(t, i2, pp)
+	require.Equal(t, uint64(67890), uu)
+
+	ok = n.CompareAndSwap(i1, 12345, nil, 0)
+	require.Falsef(t, ok, "unexpected CAS success")
+
+	pp, uu = n.Load()
+	require.Equal(t, pp, i2)
+	require.Equal(t, uu, uint64(67890))
+}

go/hack/runtime.go

@@ -54,3 +54,6 @@ func RuntimeAllocSize(size int64) int64 {
 
 //go:linkname ParseFloatPrefix strconv.parseFloatPrefix
 func ParseFloatPrefix(s string, bitSize int) (float64, int, error)
+
+//go:linkname FastRand runtime.fastrand
+func FastRand() uint32