Skip to content

authzed/crdbpool

Repository files navigation

crdbpool

Go Report Card GoDoc

crdbpool implements a node-aware connection pool for CockroachDB by wrapping pgx.

By introducing advanced health-checking, this library is able to improve:

  • performance by balancing connections evenly across nodes
  • resiliency by retrying failed queries against a different node

Have questions? Ask in our Discord.

Looking to contribute? See CONTRIBUTING.md.

You can find issues by priority: Urgent, High, Medium, Low, Maybe. There are also good first issues.

Example Usage

package main

import (
	"context"
	"fmt"
	"time"

	"github.com/authzed/crdbpool"
	"github.com/jackc/pgx/v5"
	"github.com/jackc/pgx/v5/pgxpool"
	"golang.org/x/sync/errgroup"
)

func main() {
	// Allocate a health checker that will track CockroachDB nodes
	healthChecker, err := crdbpool.NewNodeHealthChecker("postgres://username:password@localhost:5432/database_name")
	if err != nil {
		panic(err)
	}

	// Configure pgxpool as you normally would
	config, err := pgxpool.ParseConfig("postgres://username:password@localhost:5432/database_name")
	if err != nil {
		panic(err)
	}

	// Wrap the pgxpool with node-aware retry logic
	maxRetries := 10
	connectRate := 100 * time.Millisecond
	pool, err := crdbpool.NewRetryPool(context.TODO(), "pool", config, healthChecker, maxRetries, connectRate)
	if err != nil {
		panic(err)
	}

	ctx, cancel := context.WithCancel(context.TODO())
	defer cancel()

	// Start balancing connections and discovering nodes
	g, ctx := errgroup.WithContext(ctx)
	balancer := pool.NewNodeConnectionBalancer(pool, healthChecker, 5*time.Second)
	g.Go(func() error {
		balancer.Prune(ctx)
		return nil
	})
	g.Go(func() error {
		healthChecker.Poll(ctx, 5*time.Second)
		return nil
	})

	// Use the pgxpool as normal.
	pool.QueryFunc(ctx, func(ctx context.Context, rows pgx.Rows) error {
		docs := make([][]byte, 0)
		for rows.Next() {
			var doc []byte
			err = rows.Scan(&doc)
			if err != nil {
				return fmt.Errorf("couldn't scan document")
			}
			docs = append(docs, doc)
		}
		return rows.Err()
	}, "SELECT * FROM documents;")
}

Acknowledgements

This library was produced from AuthZed's findings along with the collaboration of Cockroach Labs. As a result, we'd like to thank a few notable contributors: