crdb: add a connection-balancing retry-aware connection pool

e2e/go.mod

@@ -5,10 +5,10 @@ go 1.19
 require (
 	github.com/authzed/authzed-go v0.8.1-0.20230511215435-5ca552307f34
 	github.com/authzed/grpcutil v0.0.0-20220104222419-f813f77722e5
-	github.com/authzed/spicedb v1.5.0
+	github.com/authzed/spicedb v1.21.0
 	github.com/brianvoe/gofakeit/v6 v6.15.0
 	github.com/ecordell/optgen v0.0.9
-	github.com/jackc/pgx/v5 v5.3.1
+	github.com/jackc/pgx/v5 v5.3.2-0.20230516120255-eab316e200b1
 	github.com/stretchr/testify v1.8.2
 	golang.org/x/tools v0.9.1
 	google.golang.org/grpc v1.54.0

e2e/go.sum

@@ -88,8 +88,8 @@ github.com/jackc/pgpassfile v1.0.0 h1:/6Hmqy13Ss2zCq62VdNG8tM1wchn8zjSGOBJ6icpsI
 github.com/jackc/pgpassfile v1.0.0/go.mod h1:CEx0iS5ambNFdcRtxPj5JhEz+xB6uRky5eyVu/W2HEg=
 github.com/jackc/pgservicefile v0.0.0-20221227161230-091c0ba34f0a h1:bbPeKD0xmW/Y25WS6cokEszi5g+S0QxI/d45PkRi7Nk=
 github.com/jackc/pgservicefile v0.0.0-20221227161230-091c0ba34f0a/go.mod h1:5TJZWKEWniPve33vlWYSoGYefn3gLQRzjfDlhSJ9ZKM=
-github.com/jackc/pgx/v5 v5.3.1 h1:Fcr8QJ1ZeLi5zsPZqQeUZhNhxfkkKBOgJuYkJHoBOtU=
-github.com/jackc/pgx/v5 v5.3.1/go.mod h1:t3JDKnCBlYIc0ewLF0Q7B8MXmoIaBOZj/ic7iHozM/8=
+github.com/jackc/pgx/v5 v5.3.2-0.20230516120255-eab316e200b1 h1:D4pUuZ7lcoyZKket8etOvXOVNsasWJFpZ3CmO+pK5Sw=
+github.com/jackc/pgx/v5 v5.3.2-0.20230516120255-eab316e200b1/go.mod h1:sU+RaYl9qnhD3Ce+mwnFii6YEPx70mCYghBzKvqq4qo=
 github.com/jzelinskie/stringz v0.0.1 h1:IahR+y8ct2nyj7B6i8UtFsGFj4ex1SX27iKFYsAheLk=
 github.com/jzelinskie/stringz v0.0.1/go.mod h1:hHYbgxJuNLRw91CmpuFsYEOyQqpDVFg8pvEh23vy4P0=
 github.com/kisielk/errcheck v1.5.0/go.mod h1:pFxgyoBC7bSaBwPgfKdkLd5X25qrDl4LWUI2bnpBCr8=

go.mod

@@ -42,11 +42,12 @@ require (
 	github.com/influxdata/tdigest v0.0.1
 	github.com/jackc/pgio v1.0.0
 	github.com/jackc/pgx-zerolog v0.0.0-20230315001418-f978528409eb
-	github.com/jackc/pgx/v5 v5.3.1
+	github.com/jackc/pgx/v5 v5.3.2-0.20230516120255-eab316e200b1
 	github.com/johannesboyne/gofakes3 v0.0.0-20220314170512-33c13122505e
 	github.com/jzelinskie/cobrautil/v2 v2.0.0-20230403163312-3593f31d8fe1
 	github.com/jzelinskie/stringz v0.0.1
 	github.com/lib/pq v1.10.7
+	github.com/lthibault/jitterbug v2.0.0+incompatible
 	github.com/magefile/mage v1.15.0
 	github.com/mostynb/go-grpc-compression v1.1.17
 	github.com/ngrok/sqlmw v0.0.0-20211220175533-9d16fdc47b31
@@ -341,10 +342,12 @@ require (
 	golang.org/x/net v0.10.0 // indirect
 	golang.org/x/oauth2 v0.6.0 // indirect
 	golang.org/x/sys v0.8.0 // indirect
+	golang.org/x/term v0.8.0 // indirect
 	golang.org/x/text v0.9.0 // indirect
 	golang.org/x/time v0.3.0 // indirect
 	golang.org/x/xerrors v0.0.0-20220907171357-04be3eba64a2 // indirect
 	google.golang.org/appengine v1.6.7 // indirect
+	gopkg.in/inf.v0 v0.9.1 // indirect
 	gopkg.in/ini.v1 v1.67.0 // indirect
 	honnef.co/go/tools v0.4.3 // indirect
 	k8s.io/api v0.26.0 // indirect
@@ -360,8 +363,3 @@ require (
 	sigs.k8s.io/structured-merge-diff/v4 v4.2.3 // indirect
 	sigs.k8s.io/yaml v1.3.0 // indirect
 )
-
-require (
-	golang.org/x/term v0.8.0 // indirect
-	gopkg.in/inf.v0 v0.9.1 // indirect
-)

go.sum

@@ -911,8 +911,8 @@ github.com/jackc/pgservicefile v0.0.0-20221227161230-091c0ba34f0a h1:bbPeKD0xmW/
 github.com/jackc/pgservicefile v0.0.0-20221227161230-091c0ba34f0a/go.mod h1:5TJZWKEWniPve33vlWYSoGYefn3gLQRzjfDlhSJ9ZKM=
 github.com/jackc/pgx-zerolog v0.0.0-20230315001418-f978528409eb h1:pSv+zRVeAYjbXRFjyytFIMRBSKWVowCi7KbXSMR/+ug=
 github.com/jackc/pgx-zerolog v0.0.0-20230315001418-f978528409eb/go.mod h1:CRUuPsmIajLt3dZIlJ5+O8IDSib6y8yrst8DkCthTa4=
-github.com/jackc/pgx/v5 v5.3.1 h1:Fcr8QJ1ZeLi5zsPZqQeUZhNhxfkkKBOgJuYkJHoBOtU=
-github.com/jackc/pgx/v5 v5.3.1/go.mod h1:t3JDKnCBlYIc0ewLF0Q7B8MXmoIaBOZj/ic7iHozM/8=
+github.com/jackc/pgx/v5 v5.3.2-0.20230516120255-eab316e200b1 h1:D4pUuZ7lcoyZKket8etOvXOVNsasWJFpZ3CmO+pK5Sw=
+github.com/jackc/pgx/v5 v5.3.2-0.20230516120255-eab316e200b1/go.mod h1:sU+RaYl9qnhD3Ce+mwnFii6YEPx70mCYghBzKvqq4qo=
 github.com/jackc/puddle/v2 v2.2.0 h1:RdcDk92EJBuBS55nQMMYFXTxwstHug4jkhT5pq8VxPk=
 github.com/jackc/puddle/v2 v2.2.0/go.mod h1:vriiEXHvEE654aYKXXjOvZM39qJ0q+azkZFrfEOc3H4=
 github.com/jdxcode/netrc v0.0.0-20221124155335-4616370d1a84 h1:2uT3aivO7NVpUPGcQX7RbHijHMyWix/yCnIrCWc+5co=
@@ -998,6 +998,8 @@ github.com/leonklingele/grouper v1.1.1 h1:suWXRU57D4/Enn6pXR0QVqqWWrnJ9Osrz+5rjt
 github.com/leonklingele/grouper v1.1.1/go.mod h1:uk3I3uDfi9B6PeUjsCKi6ndcf63Uy7snXgR4yDYQVDY=
 github.com/lib/pq v1.10.7 h1:p7ZhMD+KsSRozJr34udlUrhboJwWAgCg34+/ZZNvZZw=
 github.com/lib/pq v1.10.7/go.mod h1:AlVN5x4E4T544tWzH6hKfbfQvm3HdbOxrmggDNAPY9o=
+github.com/lthibault/jitterbug v2.0.0+incompatible h1:qouq51IKzlMx25+15jbxhC/d79YyTj0q6XFoptNqaUw=
+github.com/lthibault/jitterbug v2.0.0+incompatible/go.mod h1:2l7akWd27PScEs6YkjyUVj/8hKgNhbbQ3KiJgJtlf6o=
 github.com/lufeee/execinquery v1.2.1 h1:hf0Ems4SHcUGBxpGN7Jz78z1ppVkP/837ZlETPCEtOM=
 github.com/lufeee/execinquery v1.2.1/go.mod h1:EC7DrEKView09ocscGHC+apXMIaorh4xqSxS/dy8SbM=
 github.com/lyft/protoc-gen-star v0.6.0/go.mod h1:TGAoBVkt8w7MPG72TrKIu85MIdXwDuzJYeZuUPFPNwA=

internal/datastore/crdb/caveat.go

@@ -9,7 +9,6 @@ import (
 	sq "github.com/Masterminds/squirrel"
 	"github.com/jackc/pgx/v5"
 
-	pgxcommon "github.com/authzed/spicedb/internal/datastore/postgres/common"
 	"github.com/authzed/spicedb/pkg/datastore"
 	core "github.com/authzed/spicedb/pkg/proto/core/v1"
 )
@@ -43,18 +42,17 @@ func (cr *crdbReader) ReadCaveatByName(ctx context.Context, name string) (*core.
 
 	var definitionBytes []byte
 	var timestamp time.Time
-	err = cr.executeWithTx(ctx, func(ctx context.Context, tx pgxcommon.DBReader) error {
-		if err := tx.QueryRow(ctx, sql, args...).Scan(&definitionBytes, &timestamp); err != nil {
-			if errors.Is(err, pgx.ErrNoRows) {
-				err = datastore.NewCaveatNameNotFoundErr(name)
-			}
-			return err
-		}
-		return nil
-	})
+
+	err = cr.query.QueryRowFunc(ctx, func(ctx context.Context, row pgx.Row) error {
+		return row.Scan(&definitionBytes, &timestamp)
+	}, sql, args...)
 	if err != nil {
+		if errors.Is(err, pgx.ErrNoRows) {
+			err = datastore.NewCaveatNameNotFoundErr(name)
+		}
 		return nil, datastore.NoRevision, fmt.Errorf(errReadCaveat, name, err)
 	}
+
 	loaded := &core.CaveatDefinition{}
 	if err := loaded.UnmarshalVT(definitionBytes); err != nil {
 		return nil, datastore.NoRevision, fmt.Errorf(errReadCaveat, name, err)
@@ -89,28 +87,23 @@ func (cr *crdbReader) lookupCaveats(ctx context.Context, caveatNames []string) (
 	if err != nil {
 		return nil, fmt.Errorf(errListCaveats, err)
 	}
-	var allDefinitionBytes []bytesAndTimestamp
 
-	err = cr.executeWithTx(ctx, func(ctx context.Context, tx pgxcommon.DBReader) error {
-		rows, err := tx.Query(ctx, sql, args...)
-		if err != nil {
-			return err
-		}
-		defer rows.Close()
+	var allDefinitionBytes []bytesAndTimestamp
 
+	err = cr.query.QueryFunc(ctx, func(ctx context.Context, rows pgx.Rows) error {
 		for rows.Next() {
 			var defBytes []byte
 			var name string
 			var timestamp time.Time
 			err = rows.Scan(&name, &defBytes, &timestamp)
 			if err != nil {
-				return err
+				return fmt.Errorf(errListCaveats, err)
 			}
 			allDefinitionBytes = append(allDefinitionBytes, bytesAndTimestamp{bytes: defBytes, timestamp: timestamp})
 			cr.addOverlapKey(name)
 		}
 		return nil
-	})
+	}, sql, args...)
 	if err != nil {
 		return nil, fmt.Errorf(errListCaveats, err)
 	}
@@ -155,12 +148,10 @@ func (rwt *crdbReadWriteTXN) WriteCaveats(ctx context.Context, caveats []*core.C
 	for _, val := range writtenCaveatNames {
 		rwt.addOverlapKey(val)
 	}
-	return rwt.executeWithTx(ctx, func(ctx context.Context, tx pgxcommon.DBReader) error {
-		if _, err := rwt.tx.Exec(ctx, sql, args...); err != nil {
-			return fmt.Errorf(errWriteCaveat, err)
-		}
-		return nil
-	})
+	if _, err := rwt.tx.Exec(ctx, sql, args...); err != nil {
+		return fmt.Errorf(errWriteCaveat, err)
+	}
+	return nil
 }
 
 func (rwt *crdbReadWriteTXN) DeleteCaveats(ctx context.Context, names []string) error {
@@ -172,22 +163,8 @@ func (rwt *crdbReadWriteTXN) DeleteCaveats(ctx context.Context, names []string)
 	for _, val := range names {
 		rwt.addOverlapKey(val)
 	}
-	return rwt.executeWithTx(ctx, func(ctx context.Context, tx pgxcommon.DBReader) error {
-		if _, err := tx.Exec(ctx, sql, args...); err != nil {
-			return fmt.Errorf(errDeleteCaveats, err)
-		}
-		return nil
-	})
-}
-
-func (cr *crdbReader) executeWithTx(ctx context.Context, f func(ctx context.Context, tx pgxcommon.DBReader) error) error {
-	return cr.execute(ctx, func(ctx context.Context) error {
-		tx, txCleanup, err := cr.txSource(ctx)
-		if err != nil {
-			return err
-		}
-		defer txCleanup(ctx)
-
-		return f(ctx, tx)
-	})
+	if _, err := rwt.tx.Exec(ctx, sql, args...); err != nil {
+		return fmt.Errorf(errDeleteCaveats, err)
+	}
+	return nil
 }