pull in latest changes from spicedb

Author: ecordell

go.mod

@@ -1,43 +1,46 @@
 module github.com/authzed/crdbpool
 
-go 1.20
+go 1.24.0
+
+toolchain go1.24.6
 
 require (
-	github.com/grpc-ecosystem/go-grpc-middleware/v2 v2.0.0-rc.5
-	github.com/jackc/pgx/v5 v5.3.2-0.20230529162321-9720d0d63faf
+	github.com/ccoveille/go-safecast v1.6.1
+	github.com/google/uuid v1.6.0
+	github.com/grpc-ecosystem/go-grpc-middleware/v2 v2.3.2
+	github.com/jackc/pgx/v5 v5.7.5
 	github.com/lthibault/jitterbug v2.0.0+incompatible
-	github.com/prometheus/client_golang v1.15.1
-	github.com/rs/zerolog v1.29.0
-	github.com/stretchr/testify v1.8.4
-	golang.org/x/exp v0.0.0-20230522175609-2e198f4a06a1
-	golang.org/x/sync v0.2.0
-	golang.org/x/time v0.3.0
+	github.com/prometheus/client_golang v1.22.0
+	github.com/rs/zerolog v1.34.0
+	github.com/stretchr/testify v1.10.0
+	golang.org/x/sync v0.15.0
+	golang.org/x/time v0.12.0
+	google.golang.org/grpc v1.73.0
 )
 
 require (
 	github.com/beorn7/perks v1.0.1 // indirect
-	github.com/cespare/xxhash/v2 v2.2.0 // indirect
-	github.com/davecgh/go-spew v1.1.1 // indirect
-	github.com/golang/protobuf v1.5.3 // indirect
+	github.com/cespare/xxhash/v2 v2.3.0 // indirect
+	github.com/davecgh/go-spew v1.1.2-0.20180830191138-d8f796af33cc // indirect
+	github.com/go-logr/logr v1.4.3 // indirect
 	github.com/jackc/pgpassfile v1.0.0 // indirect
-	github.com/jackc/pgservicefile v0.0.0-20221227161230-091c0ba34f0a // indirect
-	github.com/jackc/puddle/v2 v2.2.0 // indirect
-	github.com/kr/text v0.2.0 // indirect
-	github.com/magefile/mage v1.15.0 // indirect
-	github.com/mattn/go-colorable v0.1.13 // indirect
-	github.com/mattn/go-isatty v0.0.16 // indirect
-	github.com/matttproud/golang_protobuf_extensions v1.0.4 // indirect
-	github.com/pmezard/go-difflib v1.0.0 // indirect
-	github.com/prometheus/client_model v0.3.0 // indirect
-	github.com/prometheus/common v0.42.0 // indirect
-	github.com/prometheus/procfs v0.9.0 // indirect
-	github.com/rogpeppe/go-internal v1.10.0 // indirect
-	golang.org/x/crypto v0.9.0 // indirect
-	golang.org/x/net v0.10.0 // indirect
-	golang.org/x/sys v0.8.0 // indirect
-	golang.org/x/text v0.9.0 // indirect
-	google.golang.org/genproto v0.0.0-20230320184635-7606e756e683 // indirect
-	google.golang.org/grpc v1.54.0 // indirect
-	google.golang.org/protobuf v1.30.0 // indirect
+	github.com/jackc/pgservicefile v0.0.0-20240606120523-5a60cdf6a761 // indirect
+	github.com/jackc/puddle/v2 v2.2.2 // indirect
+	github.com/mattn/go-colorable v0.1.14 // indirect
+	github.com/mattn/go-isatty v0.0.20 // indirect
+	github.com/munnerz/goautoneg v0.0.0-20191010083416-a7dc8b61c822 // indirect
+	github.com/pmezard/go-difflib v1.0.1-0.20181226105442-5d4384ee4fb2 // indirect
+	github.com/prometheus/client_model v0.6.2 // indirect
+	github.com/prometheus/common v0.64.0 // indirect
+	github.com/prometheus/procfs v0.15.1 // indirect
+	github.com/rogpeppe/go-internal v1.14.1 // indirect
+	go.opentelemetry.io/otel v1.36.0 // indirect
+	go.opentelemetry.io/otel/sdk/metric v1.36.0 // indirect
+	golang.org/x/crypto v0.39.0 // indirect
+	golang.org/x/net v0.41.0 // indirect
+	golang.org/x/sys v0.33.0 // indirect
+	golang.org/x/text v0.26.0 // indirect
+	google.golang.org/genproto/googleapis/rpc v0.0.0-20250603155806-513f23925822 // indirect
+	google.golang.org/protobuf v1.36.6 // indirect
 	gopkg.in/yaml.v3 v3.0.1 // indirect
 )

go.sum

@@ -3,15 +3,17 @@ package pool
 import (
 	"context"
 	"hash/maphash"
+	"maps"
+	"math"
 	"math/rand"
+	"slices"
 	"strconv"
 	"time"
 
+	"github.com/ccoveille/go-safecast"
 	"github.com/jackc/pgx/v5"
 	"github.com/jackc/pgx/v5/pgxpool"
 	"github.com/prometheus/client_golang/prometheus"
-	"golang.org/x/exp/maps"
-	"golang.org/x/exp/slices"
 	"golang.org/x/sync/semaphore"
 
 	log "github.com/rs/zerolog"
@@ -83,14 +85,26 @@ type nodeConnectionBalancer[P balancePoolConn[C], C balanceConn] struct {
 // newNodeConnectionBalancer is generic over underlying connection types for
 // testing purposes. Callers should use the exported NewNodeConnectionBalancer.
 func newNodeConnectionBalancer[P balancePoolConn[C], C balanceConn](pool balanceablePool[P, C], healthTracker *NodeHealthTracker, interval time.Duration) *nodeConnectionBalancer[P, C] {
-	seed := int64(new(maphash.Hash).Sum64())
+	seed := int64(0)
+	for seed == 0 {
+		// Sum64 returns a uint64, and safecast will return 0 if it's not castable,
+		// which will happen about half the time (?). We just keep running it until
+		// we get a seed that fits in the box.
+		// Subtracting math.MaxInt64 should mean that we retain the entire range of
+		// possible values.
+		seed, _ = safecast.ToInt64(new(maphash.Hash).Sum64() - math.MaxInt64)
+	}
 	return &nodeConnectionBalancer[P, C]{
 		ticker:        time.NewTicker(interval),
 		sem:           semaphore.NewWeighted(1),
 		healthTracker: healthTracker,
 		pool:          pool,
 		seed:          seed,
-		rnd:           rand.New(rand.NewSource(seed)),
+		// nolint:gosec
+		// use of non cryptographically secure random number generator is not concern here,
+		// as it's used for shuffling the nodes to balance the connections when the number of
+		// connections do not divide evenly.
+		rnd: rand.New(rand.NewSource(seed)),
 	}
 }
 
@@ -104,18 +118,18 @@ func (p *nodeConnectionBalancer[P, C]) Prune(ctx context.Context) {
 		case <-p.ticker.C:
 			if p.sem.TryAcquire(1) {
 				ctx, cancel := context.WithTimeout(ctx, 10*time.Second)
-				p.pruneConnections(ctx)
+				p.mustPruneConnections(ctx)
 				cancel()
 				p.sem.Release(1)
 			}
 		}
 	}
 }
 
-// pruneConnections prunes connections to nodes that have more than MaxConns/(# of nodes)
+// mustPruneConnections prunes connections to nodes that have more than MaxConns/(# of nodes)
 // This causes the pool to reconnect, which over time will lead to a balanced number of connections
 // across each node.
-func (p *nodeConnectionBalancer[P, C]) pruneConnections(ctx context.Context) {
+func (p *nodeConnectionBalancer[P, C]) mustPruneConnections(ctx context.Context) {
 	start := time.Now()
 	defer func() {
 		pruningTimeHistogram.WithLabelValues(p.pool.ID()).Observe(float64(time.Since(start).Milliseconds()))
@@ -142,7 +156,9 @@ func (p *nodeConnectionBalancer[P, C]) pruneConnections(ctx context.Context) {
 		}
 	}
 
-	nodeCount := uint32(p.healthTracker.HealthyNodeCount())
+	// It's highly unlikely that we'll ever have an overflow in
+	// this context, so we cast directly.
+	nodeCount, _ := safecast.ToUint32(p.healthTracker.HealthyNodeCount())
 	if nodeCount == 0 {
 		nodeCount = 1
 	}
@@ -169,7 +185,7 @@ func (p *nodeConnectionBalancer[P, C]) pruneConnections(ctx context.Context) {
 	}
 	p.healthTracker.RUnlock()
 
-	nodes := maps.Keys(connectionCounts)
+	nodes := slices.Collect(maps.Keys(connectionCounts))
 	slices.Sort(nodes)
 
 	// Shuffle nodes in place deterministically based on the initial seed.
@@ -198,7 +214,7 @@ func (p *nodeConnectionBalancer[P, C]) pruneConnections(ctx context.Context) {
 		// it's possible for the difference in connections between nodes to differ by up to
 		// the number of nodes.
 		if p.healthTracker.HealthyNodeCount() == 0 ||
-			uint32(i) < p.pool.MaxConns()%uint32(p.healthTracker.HealthyNodeCount()) {
+			i < int(p.pool.MaxConns())%p.healthTracker.HealthyNodeCount() {
 			perNodeMax++
 		}
 
@@ -208,7 +224,7 @@ func (p *nodeConnectionBalancer[P, C]) pruneConnections(ctx context.Context) {
 		if count <= perNodeMax {
 			continue
 		}
-		log.Ctx(ctx).Info().
+		log.Ctx(ctx).Trace().
 			Uint32("node", node).
 			Uint32("poolmaxconns", p.pool.MaxConns()).
 			Uint32("conncount", count).
@@ -220,18 +236,29 @@ func (p *nodeConnectionBalancer[P, C]) pruneConnections(ctx context.Context) {
 		if numToPrune > 1 {
 			numToPrune >>= 1
 		}
-		if uint32(len(healthyConns[node])) < numToPrune {
-			numToPrune = uint32(len(healthyConns[node]))
+
+		healthyNodeCount := mustEnsureUInt32(len(healthyConns[node]))
+		if healthyNodeCount < numToPrune {
+			numToPrune = healthyNodeCount
 		}
 		if numToPrune == 0 {
 			continue
 		}
 
 		for _, c := range healthyConns[node][:numToPrune] {
-			log.Ctx(ctx).Debug().Str("pool", p.pool.ID()).Uint32("node", node).Msg("pruning connection")
+			log.Ctx(ctx).Trace().Str("pool", p.pool.ID()).Uint32("node", node).Msg("pruning connection")
 			p.pool.GC(c.Conn())
 		}
 
-		log.Ctx(ctx).Info().Str("pool", p.pool.ID()).Uint32("node", node).Uint32("prunedCount", numToPrune).Msg("pruned connections")
+		log.Ctx(ctx).Trace().Str("pool", p.pool.ID()).Uint32("node", node).Uint32("prunedCount", numToPrune).Msg("pruned connections")
+	}
+}
+
+// mustEnsureUInt32 ensures that the specified value can be represented as a uint32.
+func mustEnsureUInt32(value int) uint32 {
+	ret, err := safecast.ToUint32(value)
+	if err != nil {
+		panic("specified value could not be cast to a uint32")
 	}
+	return ret
 }

pkg/balancer.go

@@ -150,16 +150,19 @@ func TestNodeConnectionBalancerPrune(t *testing.T) {
 
 			p := newNodeConnectionBalancer[*FakePoolConn[*FakeConn], *FakeConn](pool, tracker, 1*time.Minute)
 			p.seed = 0
+			// nolint:gosec
+			// G404 use of non cryptographically secure random number generator is not concern here,
+			// as it's used for jittering the interval for health checks.
 			p.rnd = rand.New(rand.NewSource(0))
 
 			for _, n := range tt.conns {
 				pool.nodeForConn[NewFakeConn()] = n
 			}
 
-			ctx, cancel := context.WithCancel(context.Background())
+			ctx, cancel := context.WithCancel(t.Context())
 			defer cancel()
 
-			p.pruneConnections(ctx)
+			p.mustPruneConnections(ctx)
 			require.Equal(t, len(tt.expectedGC), len(pool.gc))
 			gcFromNodes := make([]uint32, 0, len(tt.expectedGC))
 			for _, n := range pool.gc {

pkg/balancer_test.go

@@ -2,9 +2,9 @@ package pool
 
 import (
 	"context"
-	"hash/maphash"
-	"strconv"
 	"sync"
+
+	"github.com/google/uuid"
 )
 
 var (
@@ -38,13 +38,13 @@ type FakePool struct {
 	sync.Mutex
 	id          string
 	maxConns    uint32
-	gc          map[*FakeConn]uint32
-	nodeForConn map[*FakeConn]uint32
+	gc          map[*FakeConn]uint32 // GUARDED_BY(Mutex)
+	nodeForConn map[*FakeConn]uint32 // GUARDED_BY(Mutex)
 }
 
 func NewFakePool(maxConns uint32) *FakePool {
 	return &FakePool{
-		id:          strconv.FormatUint(new(maphash.Hash).Sum64(), 16),
+		id:          uuid.NewString(),
 		maxConns:    maxConns,
 		gc:          make(map[*FakeConn]uint32, 0),
 		nodeForConn: make(map[*FakeConn]uint32, 0),

pkg/fake_test.go

@@ -33,8 +33,8 @@ func init() {
 type NodeHealthTracker struct {
 	sync.RWMutex
 	connConfig    *pgx.ConnConfig
-	healthyNodes  map[uint32]struct{}
-	nodesEverSeen map[uint32]*rate.Limiter
+	healthyNodes  map[uint32]struct{}      // GUARDED_BY(RWMutex)
+	nodesEverSeen map[uint32]*rate.Limiter // GUARDED_BY(RWMutex)
 	newLimiter    func() *rate.Limiter
 }
 
@@ -58,6 +58,9 @@ func NewNodeHealthChecker(url string) (*NodeHealthTracker, error) {
 // Poll starts polling the cluster and recording the node IDs that it sees.
 func (t *NodeHealthTracker) Poll(ctx context.Context, interval time.Duration) {
 	ticker := jitterbug.New(interval, jitterbug.Uniform{
+		// nolint:gosec
+		// G404 use of non cryptographically secure random number generator is not concern here,
+		// as it's used for jittering the interval for health checks.
 		Source: rand.New(rand.NewSource(time.Now().Unix())),
 		Min:    interval,
 	})