From 5713f0fecb86ddf9b8c85229e2ab104a2e40a308 Mon Sep 17 00:00:00 2001 From: mkbeh Date: Sat, 1 Aug 2026 04:39:23 +0300 Subject: [PATCH 01/41] feat: add core connection pool --- .gitignore | 2 + .golangci.yml | 45 ++- CHANGELOG.md | 4 +- LICENSE | 2 +- README.md | 309 +----------------- Taskfile.yml | 59 +++- doc.go | 2 + errors.go | 83 +---- examples/README.md | 1 + examples/sample/README.md | 46 --- examples/sample/docker-compose.yml | 30 -- examples/sample/main.go | 114 ------- .../sample/migrations/000001_init.down.sql | 1 - examples/sample/migrations/000001_init.up.sql | 5 - examples/sample/migrations/embed.go | 6 - go.mod | 31 +- go.sum | 122 +------ internal/pkg/pgxpoolcollector/v5/collector.go | 181 ---------- internal/pkg/pgxslog/adapter.go | 63 ---- internal/pkg/pgxtracer/tracer.go | 40 --- migrate.go | 52 --- options.go | 268 ++++----------- pool.go | 265 +++------------ tx.go | 33 -- utils.go | 17 - 25 files changed, 224 insertions(+), 1557 deletions(-) create mode 100644 doc.go create mode 100644 examples/README.md delete mode 100644 examples/sample/README.md delete mode 100644 examples/sample/docker-compose.yml delete mode 100644 examples/sample/main.go delete mode 100644 examples/sample/migrations/000001_init.down.sql delete mode 100644 examples/sample/migrations/000001_init.up.sql delete mode 100644 examples/sample/migrations/embed.go delete mode 100644 internal/pkg/pgxpoolcollector/v5/collector.go delete mode 100644 internal/pkg/pgxslog/adapter.go delete mode 100644 internal/pkg/pgxtracer/tracer.go delete mode 100644 migrate.go delete mode 100644 tx.go delete mode 100644 utils.go diff --git a/.gitignore b/.gitignore index 2149a16..15e302a 100644 --- a/.gitignore +++ b/.gitignore @@ -28,3 +28,5 @@ go.work.sum .vscode .idea +# other +/tmp/ diff --git a/.golangci.yml b/.golangci.yml index e419d53..dbfafd5 100755 --- a/.golangci.yml +++ b/.golangci.yml @@ -1,7 +1,10 @@ version: "2" + run: go: "1.26" + timeout: 5m allow-parallel-runners: true + linters: default: none enable: @@ -10,6 +13,8 @@ linters: - bodyclose - copyloopvar - durationcheck + - errcheck + - errorlint - exhaustive - gocritic - goprintffuncname @@ -24,17 +29,24 @@ linters: - rowserrcheck - sloglint - sqlclosecheck + - staticcheck - unconvert - unparam - unused - usetesting - wastedassign - whitespace + settings: + errcheck: + check-type-assertions: true + errorlint: errorf: false + exhaustive: default-signifies-exhaustive: true + gocritic: disabled-checks: - appendAssign @@ -61,15 +73,18 @@ linters: sizeThreshold: 256 rangeValCopy: sizeThreshold: 256 + gosec: excludes: - G104 - - G404 - G115 + - G404 + nolintlint: require-explanation: true require-specific: true allow-unused: false + revive: confidence: 0.8 severity: warning @@ -105,7 +120,6 @@ linters: - name: superfluous-else - name: time-equal - name: time-naming - - name: var-declaration - name: unconditional-recursion - name: unexported-naming - name: unexported-return @@ -113,7 +127,9 @@ linters: - name: unreachable-code - name: unused-parameter - name: useless-break + - name: var-declaration - name: waitgroup-by-value + sloglint: no-mixed-args: true no-global: default @@ -122,22 +138,32 @@ linters: no-raw-keys: false key-naming-case: snake args-on-sep-lines: true + staticcheck: checks: - all - -SA1012 - -SA1019 + - -ST1000 + exclusions: generated: lax paths: - - third_party$ - - builtin$ - - examples + - third_party/ + - builtin/ + - examples/ + - _tmp/ + rules: + - path: _test\.go + linters: + - revive + text: "dot-imports" + issues: max-same-issues: 0 + formatters: enable: - - gofmt - gofumpt - goimports settings: @@ -146,6 +172,7 @@ formatters: exclusions: generated: lax paths: - - third_party$ - - builtin$ - - examples$ + - third_party/ + - builtin/ + - examples/ + - _tmp/ \ No newline at end of file diff --git a/CHANGELOG.md b/CHANGELOG.md index b4bf2c2..c2c8c48 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1,3 +1 @@ -# 0.1.13 (Jun 22, 2026) - -* build: bump dependencies +[TODO] \ No newline at end of file diff --git a/LICENSE b/LICENSE index f7372df..d22da70 100644 --- a/LICENSE +++ b/LICENSE @@ -1,6 +1,6 @@ MIT License -Copyright (c) 2024 Nia +Copyright (c) 2024 mkbeh Permission is hereby granted, free of charge, to any person obtaining a copy of this software and associated documentation files (the "Software"), to deal diff --git a/README.md b/README.md index a0d2be9..9dd9622 100644 --- a/README.md +++ b/README.md @@ -15,318 +15,15 @@ building, normalized errors, and exposing PostgreSQL observability with OpenTele ## Features -* **Pools**: Separate writer and reader connection pools. -* **Queries**: PostgreSQL-friendly query builder support. -* **Transactions**: Transaction helpers with automatic rollback and panic recovery. -* **Migrations**: Embedded SQL migrations via [golang-migrate](https://github.com/golang-migrate/migrate). -* **Errors**: Normalized PostgreSQL error codes for common failure cases. -* **Observability**: OpenTelemetry tracing and Prometheus metrics out of the box. -* **Configuration**: Configure via Go structs or environment variables. +[TODOO] ## Installation -```bash -go get github.com/mkbeh/xpg -``` +[TODO] ## Quick start -The example below creates separate writer and reader pools and runs a simple query. - -```go -package main - -import ( - "context" - "fmt" - "log" - - postgres "github.com/mkbeh/xpg" -) - -func main() { - ctx := context.Background() - - cfg := &postgres.Config{ - ClusterHost: "127.0.0.1", - ClusterPort: "5432", // writer/master port - ClusterReplicaPort: "5432", // reader/replica port; can be different in production - User: "user", - Password: "pass", - DB: "postgres", - } - - writer, err := postgres.NewWriter( - postgres.WithConfig(cfg), - postgres.WithClientID("my-service"), // appended to application_name and metrics labels - ) - if err != nil { - log.Fatal("failed to init writer pool:", err) - } - defer writer.Close() - - reader, err := postgres.NewReader( - postgres.WithConfig(cfg), - postgres.WithClientID("my-service"), - ) - if err != nil { - log.Fatal("failed to init reader pool:", err) - } - defer reader.Close() - - // Use the writer pool for write-side operations. - if _, err := writer.Exec(ctx, "select 1"); err != nil { - log.Fatal("writer query failed:", err) - } - - var greeting string - - // Use the reader pool for read-only queries. - if err := reader.QueryRow(ctx, "select 'Hello, world!'").Scan(&greeting); err != nil { - log.Fatal("query failed:", err) - } - - fmt.Println(greeting) -} -``` - -More examples: [examples/](https://github.com/mkbeh/xpg/tree/main/examples) - -## Query Builder - -Each pool includes a preconfigured [squirrel](https://github.com/Masterminds/squirrel) statement builder with PostgreSQL -dollar placeholders out of the box. - - -```go -sql, args, err := writer.QueryBuilder(). - Insert("orders"). - Columns("id", "status"). - Values(orderID, "created"). - ToSql() -if err != nil { - log.Fatalf("failed to build query: %v", err) -} - -if _, err := writer.Exec(ctx, sql, args...); err != nil { - log.Fatalf("failed to execute insert: %v", err) -} -``` - - -## Transactions - -Use `RunInTxx` for transactions with default options. It acts as an alias for `RunInTx` using default `pgx.TxOptions`. - - -```go -err := writer.RunInTxx(ctx, func(ctx context.Context) error { - _, err := writer.Exec(ctx, "INSERT INTO orders (id) VALUES (\$1)", orderID) - return err -}) -if err != nil { - log.Fatalf("transaction failed: %v", err) -} -``` - - -For a custom isolation level or access mode, use `RunInTx`: - - -```go -err := writer.RunInTx(ctx, func(ctx context.Context) error { - _, err := writer.Exec(ctx, "INSERT INTO orders (id) VALUES (\$1)", orderID) - return err -}, pgx.TxOptions{ - IsoLevel: pgx.Serializable, -}) -if err != nil { - log.Fatalf("serializable transaction failed: %v", err) -} -``` - - -Rollback is handled automatically if an error occurs. The transaction is committed only if the function returns `nil`. -Any panics inside the block are recovered and returned as standard Go errors. - -## Migrations - -`xpg` supports embedded SQL migrations out of the box using [golang-migrate](https://github.com/golang-migrate/migrate). - -First, create an `embed.go` file inside your migrations directory: - - -```go -package migrations - -import "embed" - -//go:embed *.sql -var FS embed.FS -``` - - -Then, pass the embedded filesystem using `WithMigrations`: - - -```go -writer, err := postgres.NewWriter( - postgres.WithConfig(&postgres.Config{ - ClusterHost: "127.0.0.1", - ClusterPort: "5432", - ClusterReplicaPort: "5432", - User: "user", - Password: "pass", - DB: "postgres", - MigrateEnabled: true, - }), - postgres.WithMigrations(migrations.FS), -) -if err != nil { - log.Fatalf("failed to initialize writer and run migrations: %v", err) -} -defer writer.Close() -``` - - -Migrations will run automatically during `NewWriter` initialization if `MigrateEnabled` is set to `true`. - -Your SQL migration files must follow the standard `golang-migrate` naming convention: - -```text -000001_create_users.up.sql -000001_create_users.down.sql -``` - -## Observability - -`xpg` instruments PostgreSQL queries through native `pgx` tracing hooks and exposes pool metrics for Prometheus. - - -```go -writer, err := postgres.NewWriter( - postgres.WithConfig(cfg), - postgres.WithClientID("orders-service"), - postgres.WithTraceProvider(tracerProvider), - postgres.WithMetricsNamespace("orders"), -) -if err != nil { - log.Fatalf("failed to initialize observed writer pool: %v", err) -} -defer writer.Close() -``` - - - -The following Prometheus metric labels are added automatically: - -| Label | Description | -| :--- | :--- | -| `client_id` | Generated client identifier or configured ID with a unique suffix. | -| `client_kind` | `writer` for writer pools, `reader` for reader pools. | -| `db` | Database name from the configuration. | -| `shard_id` | Shard ID from the configuration. | - -## Error Handling - -`xpg` provides normalized PostgreSQL error codes through `ConvertError`, so application code does not need to deal with -raw `pgx` and `pgconn` error types directly. - - -```go -err := writer.QueryRow(ctx, "SELECT id FROM users WHERE id = \$1", userID).Scan(&id) -if err != nil { - pgErr := postgres.ConvertError(err) - - if pgErr.Code() == postgres.ErrNoRows { - // handle missing row - return nil - } - - if pgErr.Code() == postgres.ErrSerializable { - // retry transaction - return nil - } - - return pgErr -} -``` - - -Common PostgreSQL errors such as `ErrNoRows`, `ErrUniqViolation`, `ErrForeignKeyViolation`, and `ErrSerializable` are -mapped to stable `xpg` error codes. - -## Configuration - -The `Config` struct can be initialized directly in Go. It also includes `envconfig` tags, allowing you to seamlessly -populate it from environment variables using your preferred configuration library. - -### Config Struct - - -```go -cfg := &postgres.Config{ - ClusterHost: "127.0.0.1", // required - ClusterPort: "5432", // required, writer port - ClusterReplicaPort: "5433", // required, reader port - User: "user", // required - Password: "pass", // required - DB: "mydb", // required - - MaxRWConn: 16, - MaxROConn: 16, - MaxConnLifetime: 5 * time.Minute, - MaxConnIdleTime: 30 * time.Second, - - MigrateEnabled: true, -} -``` - - -The connection DSN is dynamically built from the `Config` fields using the following format: - -```text -postgres://user:pass@host:port/db?sslmode=disable&application_name=& -``` - -### Environment Variables - -| Variable | Required | Default | Description | -| :--- | :---: | :--- | :--- | -| `POSTGRES_CLUSTER_HOST` | ✓ | — | Database host. | -| `POSTGRES_CLUSTER_PORT` | ✓ | — | Writer pool port. | -| `POSTGRES_CLUSTER_REPLICA_PORT` | ✓ | — | Reader pool port. | -| `POSTGRES_USER` | ✓ | — | Database user. | -| `POSTGRES_PASSWORD` | ✓ | — | Database password. | -| `POSTGRES_DB` | ✓ | — | Database name. | -| `POSTGRES_SHARD_ID` | | `0` | Shard ID exposed in metrics. | -| `POSTGRES_MIN_RW_CONN` | | `1` | Minimum connections in the writer pool. | -| `POSTGRES_MIN_RO_CONN` | | `1` | Minimum connections in the reader pool. | -| `POSTGRES_MAX_RW_CONN` | | `max(4, NumCPU)`| Maximum connections in the writer pool. | -| `POSTGRES_MAX_RO_CONN` | | `max(4, NumCPU)`| Maximum connections in the reader pool. | -| `POSTGRES_MAX_CONN_LIFETIME` | | `1m` | Maximum connection lifetime. | -| `POSTGRES_MAX_CONN_IDLE_TIME` | | `30s` | Maximum idle connection lifetime. | -| `POSTGRES_QUERY_EXEC_MODE` | | `cache_statement` | Query execution mode. | -| `POSTGRES_STATEMENT_CACHE_CAPACITY` | | `128` | Statement cache size. | -| `POSTGRES_DESCRIPTION_CACHE_CAPACITY`| | `512` | Description cache size. | -| `POSTGRES_WRITER_ARGS` | | — | Extra DSN args for the writer connection. | -| `POSTGRES_REPLICA_ARGS` | | — | Extra DSN args for the reader connection. | -| `POSTGRES_MIGRATE_ENABLED` | | `false` | Run migrations on writer startup. | -| `POSTGRES_MIGRATE_PORT` | | `POSTGRES_CLUSTER_PORT` | Port used for migrations. | -| `POSTGRES_MIGRATE_ARGS` | | — | Extra DSN args for the migration connection. | - -### Query Execution Modes - -| Value | Protocol | Round Trips | Description | -| :--- | :--- | :--- | :--- | -| `cache_statement` | Extended | 1 after warm-up | Automatically prepares and caches statements. **Default.** May fail on first execution after schema changes. | -| `cache_describe` | Extended | 1 after warm-up | Caches argument and result type descriptions instead of prepared statements. Has the same schema-change caveat. | -| `describe_exec` | Extended | 2 | Fetches description on every execution, then executes. Safer with concurrent schema changes, but may break with connection poolers that switch connections between round trips. | -| `exec` | Extended | 1 | No prepare and no describe. Infers PostgreSQL types from Go types using text format. Register custom types with `pgtype.Map.RegisterDefaultPgType`. | -| `simple_protocol` | Simple | 1 | Uses the simple protocol. Useful with PgBouncer or proxies that do not support the extended protocol. Prefer `exec` when possible. | - -> 💡 **Tip for connection pooling:** For PgBouncer **transaction** pooling, use `simple_protocol`. For **session** -> pooling, `cache_statement` usually works fine. +[TODO] ## License diff --git a/Taskfile.yml b/Taskfile.yml index bc5cf57..e1d4358 100755 --- a/Taskfile.yml +++ b/Taskfile.yml @@ -1,28 +1,55 @@ -version: 3 +version: "3" vars: LINTER_VER: "v2.12.2" + GOMODCACHE: + sh: go env GOMODCACHE + + GOCACHE: + sh: go env GOCACHE + + PWD_DIR: + sh: pwd + +env: + GOWORK: "off" + tasks: run: - desc: "Full cycle" - deps: - - lint + desc: "Run all checks" + cmds: + - task: lint + - task: test-race lint: - desc: "lint" + desc: "Run golangci-lint" + cmds: + - task: _run-linter + + test: + desc: "Run tests" + cmds: + - go test -v -count=1 ./... + + test-race: + desc: "Run tests with the race detector" + cmds: + - go test -race -v -count=1 ./... + + _run-linter: + internal: true cmds: - - docker run --rm + - > + docker run --rm -u $(id -u):$(id -g) - -v {{.PWD}}:/app - -v {{.GOCACHE_ENV}}:/go/pkg/mod - -e GOCACHE=/go/pkg/mod - -e GOLANGCI_LINT_CACHE=/go/pkg/mod + -v {{.PWD_DIR}}:/app + -v {{.GOMODCACHE}}:/go/pkg/mod + -v {{.GOCACHE}}:/go/build-cache + -e GOMODCACHE=/go/pkg/mod + -e GOCACHE=/go/build-cache + -e GOLANGCI_LINT_CACHE=/go/build-cache/golangci-lint + -e GOWORK=off -w /app golangci/golangci-lint:{{.LINTER_VER}} - golangci-lint run --timeout 5m --fix --modules-download-mode vendor - vars: - GOCACHE_ENV: - sh: go env GOCACHE - PWD: - sh: pwd + golangci-lint run --timeout 5m \ No newline at end of file diff --git a/doc.go b/doc.go new file mode 100644 index 0000000..0b4c1db --- /dev/null +++ b/doc.go @@ -0,0 +1,2 @@ +// Package xpg provides pgx-first infrastructure primitives for PostgreSQL. +package xpg diff --git a/errors.go b/errors.go index b923dfc..d5c081c 100644 --- a/errors.go +++ b/errors.go @@ -1,82 +1 @@ -package postgres - -import ( - "context" - "errors" - "net" - - "github.com/jackc/pgerrcode" - "github.com/jackc/pgx/v5" - "github.com/jackc/pgx/v5/pgconn" -) - -type PgErrorCode int - -const ( - ErrContextDeadline PgErrorCode = iota - ErrNoRows - ErrUniqViolation - ErrForeignKeyViolation - ErrSerializable - ErrOther - ErrBeginTransaction - ErrCommitTransaction - ErrNoConnection -) - -type PgError struct { - code PgErrorCode - msg string -} - -func (e PgError) Error() string { - return e.msg -} - -func (e PgError) Code() PgErrorCode { - return e.code -} - -func NewPgError(code PgErrorCode, err error) *PgError { - return &PgError{code, err.Error()} -} - -func ConvertError(err error) *PgError { - if err == nil { - return nil - } - - if pgErr, ok := err.(*PgError); ok { - return pgErr - } - - if ne, ok := err.(net.Error); ok { - return NewPgError(ErrNoConnection, ne) - } - - switch { - case errors.Is(err, context.Canceled), pgconn.Timeout(err): - return NewPgError(ErrContextDeadline, err) - case errors.Is(err, pgx.ErrNoRows): - return NewPgError(ErrNoRows, err) - default: - var pgErr *pgconn.PgError - if !errors.As(err, &pgErr) { - return NewPgError(ErrOther, err) - } - return NewPgError(pgCodeToError(pgErr.Code), err) - } -} - -var pgCodeMap = map[string]PgErrorCode{ - pgerrcode.UniqueViolation: ErrUniqViolation, - pgerrcode.ForeignKeyViolation: ErrForeignKeyViolation, - pgerrcode.SerializationFailure: ErrSerializable, -} - -func pgCodeToError(code string) PgErrorCode { - if c, ok := pgCodeMap[code]; ok { - return c - } - return ErrOther -} +package xpg diff --git a/examples/README.md b/examples/README.md new file mode 100644 index 0000000..c2c8c48 --- /dev/null +++ b/examples/README.md @@ -0,0 +1 @@ +[TODO] \ No newline at end of file diff --git a/examples/sample/README.md b/examples/sample/README.md deleted file mode 100644 index 8b21258..0000000 --- a/examples/sample/README.md +++ /dev/null @@ -1,46 +0,0 @@ -# Description - -This is a sample REST API service implemented using pgx as the connector to a PostgreSQL data store. - -# Usage - -Create a PostgreSQL database. - -Configure the database connection with environment variables: - -```text -POSTGRES_CLUSTER_HOST=127.0.0.1 -POSTGRES_CLUSTER_PORT=54320 -POSTGRES_DB=db -POSTGRES_PASSWORD=pass -POSTGRES_USER=user -``` - -Set up docker-compose: -```shell -docker-compose up --build -d -``` - -Run main.go: - -``` -go run main.go -``` - -## Create tasks - -```shell -curl '127.0.0.1:8080/create' -``` - -## Get tasks - -```shell -curl '127.0.0.1:8080/get' -``` - -## Metrics - -```shell -curl 'http://localhost:8080/metrics' -``` \ No newline at end of file diff --git a/examples/sample/docker-compose.yml b/examples/sample/docker-compose.yml deleted file mode 100644 index be9059b..0000000 --- a/examples/sample/docker-compose.yml +++ /dev/null @@ -1,30 +0,0 @@ -version: '3.9' -services: - postgres: - image: postgres:14.2-alpine - container_name: sample-postgres - environment: - POSTGRES_DB: "db" - POSTGRES_PASSWORD: "pass" - POSTGRES_USER: "user" - ports: - - "54320:5432" - volumes: - - ./scripts/init_postgres.sh:/docker-entrypoint-initdb.d/init_postgres.sh - - sample-pgdata:/var/lib/postgresql/data - networks: - - sample-network - healthcheck: - test: [ "CMD-SHELL", "pg_isready -d $${POSTGRES_DB} -U $${POSTGRES_USER}" ] - interval: 10s - timeout: 5s - retries: 5 - start_period: 10s - - -networks: - sample-network: - driver: bridge - -volumes: - sample-pgdata: \ No newline at end of file diff --git a/examples/sample/main.go b/examples/sample/main.go deleted file mode 100644 index 2528e12..0000000 --- a/examples/sample/main.go +++ /dev/null @@ -1,114 +0,0 @@ -package main - -import ( - "encoding/json" - "log" - "net/http" - "os" - - postgres "github.com/mkbeh/xpg" - "github.com/mkbeh/xpg/examples/sample/migrations" - "github.com/prometheus/client_golang/prometheus/promhttp" -) - -var ( - writer *postgres.Pool - reader *postgres.Pool -) - -var ( - host string - port string - user string - pass string - db string -) - -func init() { - host = os.Getenv("POSTGRES_CLUSTER_HOST") - port = os.Getenv("POSTGRES_CLUSTER_PORT") - user = os.Getenv("POSTGRES_USER") - pass = os.Getenv("POSTGRES_PASSWORD") - db = os.Getenv("POSTGRES_DB") -} - -func getTasksHandler(w http.ResponseWriter, req *http.Request) { - type task struct { - Id int `json:"id"` - Description string `json:"description"` - } - - rows, err := reader.Query(req.Context(), "SELECT id, description FROM tasks;") - if err != nil { - http.Error(w, err.Error(), http.StatusInternalServerError) - return - } - defer rows.Close() - - tasks := make([]task, 0) - for rows.Next() { - var v task - if err := rows.Scan(&v.Id, &v.Description); err != nil { - http.Error(w, err.Error(), http.StatusInternalServerError) - return - } - tasks = append(tasks, v) - } - - data, _ := json.Marshal(tasks) - - w.Header().Set("Content-Type", "application/json") - w.Write(data) -} - -func createTasksHandler(w http.ResponseWriter, req *http.Request) { - _, err := writer.Exec(req.Context(), "INSERT INTO tasks VALUES (1, 'test1'), (2, 'test-2') ON CONFLICT DO NOTHING;") - if err != nil { - http.Error(w, err.Error(), http.StatusInternalServerError) - } else { - w.Header().Set("Content-Type", "text/plain") - w.Write([]byte("OK")) - } -} - -func main() { - var err error - - cfg := &postgres.Config{ - ClusterHost: host, - ClusterPort: port, - ClusterReplicaPort: port, - User: user, - Password: pass, - DB: db, - MigrateEnabled: true, - } - - writer, err = postgres.NewWriter( - postgres.WithConfig(cfg), - postgres.WithClientID("test-client"), - postgres.WithMigrations(migrations.FS), - ) - if err != nil { - log.Fatalln("failed init master pool", err) - } - defer writer.Close() - - reader, err = postgres.NewReader( - postgres.WithConfig(cfg), - postgres.WithClientID("test-client"), - ) - if err != nil { - log.Fatalln("failed init reader pool", err) - } - defer reader.Close() - - http.HandleFunc("/get", getTasksHandler) - http.HandleFunc("/create", createTasksHandler) - http.Handle("/metrics", promhttp.Handler()) - - err = http.ListenAndServe("localhost:8080", nil) - if err != nil { - log.Fatalln("Unable to start web server:", err) - } -} diff --git a/examples/sample/migrations/000001_init.down.sql b/examples/sample/migrations/000001_init.down.sql deleted file mode 100644 index 22c971f..0000000 --- a/examples/sample/migrations/000001_init.down.sql +++ /dev/null @@ -1 +0,0 @@ -drop table tasks; \ No newline at end of file diff --git a/examples/sample/migrations/000001_init.up.sql b/examples/sample/migrations/000001_init.up.sql deleted file mode 100644 index d7798b1..0000000 --- a/examples/sample/migrations/000001_init.up.sql +++ /dev/null @@ -1,5 +0,0 @@ -create table tasks -( - id serial primary key, - description text not null -); \ No newline at end of file diff --git a/examples/sample/migrations/embed.go b/examples/sample/migrations/embed.go deleted file mode 100644 index 91cca1c..0000000 --- a/examples/sample/migrations/embed.go +++ /dev/null @@ -1,6 +0,0 @@ -package migrations - -import "embed" - -//go:embed *.sql -var FS embed.FS diff --git a/go.mod b/go.mod index f94ec39..a572d11 100644 --- a/go.mod +++ b/go.mod @@ -2,37 +2,12 @@ module github.com/mkbeh/xpg go 1.26 -require ( - github.com/Masterminds/squirrel v1.5.4 - github.com/exaring/otelpgx v0.11.1 - github.com/golang-migrate/migrate/v4 v4.19.1 - github.com/google/uuid v1.6.0 - github.com/jackc/pgerrcode v0.0.0-20250907135507-afb5586c32a6 - github.com/jackc/pgx/v5 v5.10.0 - github.com/prometheus/client_golang v1.23.2 - go.opentelemetry.io/otel v1.44.0 - go.opentelemetry.io/otel/trace v1.44.0 -) +require github.com/jackc/pgx/v5 v5.10.0 require ( - github.com/beorn7/perks v1.0.1 // indirect - github.com/cespare/xxhash/v2 v2.3.0 // indirect - github.com/go-logr/logr v1.4.3 // indirect - github.com/go-logr/stdr v1.2.2 // indirect github.com/jackc/pgpassfile v1.0.0 // indirect github.com/jackc/pgservicefile v0.0.0-20240606120523-5a60cdf6a761 // indirect github.com/jackc/puddle/v2 v2.2.2 // indirect - github.com/lann/builder v0.0.0-20180802200727-47ae307949d0 // indirect - github.com/lann/ps v0.0.0-20150810152359-62de8c46ede0 // indirect - github.com/lib/pq v1.12.3 // indirect - github.com/munnerz/goautoneg v0.0.0-20191010083416-a7dc8b61c822 // indirect - github.com/prometheus/client_model v0.6.2 // indirect - github.com/prometheus/common v0.69.0 // indirect - github.com/prometheus/procfs v0.20.1 // indirect - go.opentelemetry.io/auto/sdk v1.2.1 // indirect - go.opentelemetry.io/otel/metric v1.44.0 // indirect - golang.org/x/sync v0.21.0 // indirect - golang.org/x/sys v0.46.0 // indirect - golang.org/x/text v0.38.0 // indirect - google.golang.org/protobuf v1.36.11 // indirect + golang.org/x/sync v0.22.0 // indirect + golang.org/x/text v0.40.0 // indirect ) diff --git a/go.sum b/go.sum index 6b66baa..583cb35 100644 --- a/go.sum +++ b/go.sum @@ -1,135 +1,25 @@ -github.com/Azure/go-ansiterm v0.0.0-20230124172434-306776ec8161 h1:L/gRVlceqvL25UVaW/CKtUDjefjrs0SPonmDGUVOYP0= -github.com/Azure/go-ansiterm v0.0.0-20230124172434-306776ec8161/go.mod h1:xomTg63KZ2rFqZQzSB4Vz2SUXa1BpHTVz9L5PTmPC4E= -github.com/Masterminds/squirrel v1.5.4 h1:uUcX/aBc8O7Fg9kaISIUsHXdKuqehiXAMQTYX8afzqM= -github.com/Masterminds/squirrel v1.5.4/go.mod h1:NNaOrjSoIDfDA40n7sr2tPNZRfjzjA400rg+riTZj10= -github.com/Microsoft/go-winio v0.6.2 h1:F2VQgta7ecxGYO8k3ZZz3RS8fVIXVxONVUPlNERoyfY= -github.com/Microsoft/go-winio v0.6.2/go.mod h1:yd8OoFMLzJbo9gZq8j5qaps8bJ9aShtEA8Ipt1oGCvU= -github.com/beorn7/perks v1.0.1 h1:VlbKKnNfV8bJzeqoa4cOKqO6bYr3WgKZxO8Z16+hsOM= -github.com/beorn7/perks v1.0.1/go.mod h1:G2ZrVWU2WbWT9wwq4/hrbKbnv/1ERSJQ0ibhJ6rlkpw= -github.com/cespare/xxhash/v2 v2.3.0 h1:UL815xU9SqsFlibzuggzjXhog7bL6oX9BbNZnL2UFvs= -github.com/cespare/xxhash/v2 v2.3.0/go.mod h1:VGX0DQ3Q6kWi7AoAeZDth3/j3BFtOZR5XLFGgcrjCOs= -github.com/containerd/errdefs v1.0.0 h1:tg5yIfIlQIrxYtu9ajqY42W3lpS19XqdxRQeEwYG8PI= -github.com/containerd/errdefs v1.0.0/go.mod h1:+YBYIdtsnF4Iw6nWZhJcqGSg/dwvV7tyJ/kCkyJ2k+M= -github.com/containerd/errdefs/pkg v0.3.0 h1:9IKJ06FvyNlexW690DXuQNx2KA2cUJXx151Xdx3ZPPE= -github.com/containerd/errdefs/pkg v0.3.0/go.mod h1:NJw6s9HwNuRhnjJhM7pylWwMyAkmCQvQ4GpJHEqRLVk= github.com/davecgh/go-spew v1.1.0/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38= +github.com/davecgh/go-spew v1.1.1 h1:vj9j/u1bqnvCEfJOwUhtlOARqs3+rkHYY13jYWTU97c= github.com/davecgh/go-spew v1.1.1/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38= -github.com/davecgh/go-spew v1.1.2-0.20180830191138-d8f796af33cc h1:U9qPSI2PIWSS1VwoXQT9A3Wy9MM3WgvqSxFWenqJduM= -github.com/davecgh/go-spew v1.1.2-0.20180830191138-d8f796af33cc/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38= -github.com/dhui/dktest v0.4.6 h1:+DPKyScKSEp3VLtbMDHcUq6V5Lm5zfZZVb0Sk7Ahom4= -github.com/dhui/dktest v0.4.6/go.mod h1:JHTSYDtKkvFNFHJKqCzVzqXecyv+tKt8EzceOmQOgbU= -github.com/distribution/reference v0.6.0 h1:0IXCQ5g4/QMHHkarYzh5l+u8T3t73zM5QvfrDyIgxBk= -github.com/distribution/reference v0.6.0/go.mod h1:BbU0aIcezP1/5jX/8MP0YiH4SdvB5Y4f/wlDRiLyi3E= -github.com/docker/docker v28.3.3+incompatible h1:Dypm25kh4rmk49v1eiVbsAtpAsYURjYkaKubwuBdxEI= -github.com/docker/docker v28.3.3+incompatible/go.mod h1:eEKB0N0r5NX/I1kEveEz05bcu8tLC/8azJZsviup8Sk= -github.com/docker/go-connections v0.5.0 h1:USnMq7hx7gwdVZq1L49hLXaFtUdTADjXGp+uj1Br63c= -github.com/docker/go-connections v0.5.0/go.mod h1:ov60Kzw0kKElRwhNs9UlUHAE/F9Fe6GLaXnqyDdmEXc= -github.com/docker/go-units v0.5.0 h1:69rxXcBk27SvSaaxTtLh/8llcHD8vYHT7WSdRZ/jvr4= -github.com/docker/go-units v0.5.0/go.mod h1:fgPhTUdO+D/Jk86RDLlptpiXQzgHJF7gydDDbaIK4Dk= -github.com/exaring/otelpgx v0.11.1 h1:pE79fIg/qh/Lpu00kvswFC5dKfqyJJhMJ4Y4N3w5Lj4= -github.com/exaring/otelpgx v0.11.1/go.mod h1:3OojrUKhhy3lTbYIMBijP3YjMey/jo14eHAW5cXcUdk= -github.com/felixge/httpsnoop v1.0.4 h1:NFTV2Zj1bL4mc9sqWACXbQFVBBg2W3GPvqp8/ESS2Wg= -github.com/felixge/httpsnoop v1.0.4/go.mod h1:m8KPJKqk1gH5J9DgRY2ASl2lWCfGKXixSwevea8zH2U= -github.com/go-logr/logr v1.2.2/go.mod h1:jdQByPbusPIv2/zmleS9BjJVeZ6kBagPoEUsqbVz/1A= -github.com/go-logr/logr v1.4.3 h1:CjnDlHq8ikf6E492q6eKboGOC0T8CDaOvkHCIg8idEI= -github.com/go-logr/logr v1.4.3/go.mod h1:9T104GzyrTigFIr8wt5mBrctHMim0Nb2HLGrmQ40KvY= -github.com/go-logr/stdr v1.2.2 h1:hSWxHoqTgW2S2qGc0LTAI563KZ5YKYRhT3MFKZMbjag= -github.com/go-logr/stdr v1.2.2/go.mod h1:mMo/vtBO5dYbehREoey6XUKy/eSumjCCveDpRre4VKE= -github.com/gogo/protobuf v1.3.2 h1:Ov1cvc58UF3b5XjBnZv7+opcTcQFZebYjWzi34vdm4Q= -github.com/gogo/protobuf v1.3.2/go.mod h1:P1XiOD3dCwIKUDQYPy72D8LYyHL2YPYrpS2s69NZV8Q= -github.com/golang-migrate/migrate/v4 v4.19.1 h1:OCyb44lFuQfYXYLx1SCxPZQGU7mcaZ7gH9yH4jSFbBA= -github.com/golang-migrate/migrate/v4 v4.19.1/go.mod h1:CTcgfjxhaUtsLipnLoQRWCrjYXycRz/g5+RWDuYgPrE= -github.com/google/go-cmp v0.7.0 h1:wk8382ETsv4JYUZwIsn6YpYiWiBsYLSJiTsyBybVuN8= -github.com/google/go-cmp v0.7.0/go.mod h1:pXiqmnSA92OHEEa9HXL2W4E7lf9JzCmGVUdgjX3N/iU= -github.com/google/uuid v1.6.0 h1:NIvaJDMOsjHA8n1jAhLSgzrAzy1Hgr+hNrb57e+94F0= -github.com/google/uuid v1.6.0/go.mod h1:TIyPZe4MgqvfeYDBFedMoGGpEw/LqOeaOT+nhxU+yHo= -github.com/jackc/pgerrcode v0.0.0-20250907135507-afb5586c32a6 h1:D/V0gu4zQ3cL2WKeVNVM4r2gLxGGf6McLwgXzRTo2RQ= -github.com/jackc/pgerrcode v0.0.0-20250907135507-afb5586c32a6/go.mod h1:a/s9Lp5W7n/DD0VrVoyJ00FbP2ytTPDVOivvn2bMlds= github.com/jackc/pgpassfile v1.0.0 h1:/6Hmqy13Ss2zCq62VdNG8tM1wchn8zjSGOBJ6icpsIM= github.com/jackc/pgpassfile v1.0.0/go.mod h1:CEx0iS5ambNFdcRtxPj5JhEz+xB6uRky5eyVu/W2HEg= github.com/jackc/pgservicefile v0.0.0-20240606120523-5a60cdf6a761 h1:iCEnooe7UlwOQYpKFhBabPMi4aNAfoODPEFNiAnClxo= github.com/jackc/pgservicefile v0.0.0-20240606120523-5a60cdf6a761/go.mod h1:5TJZWKEWniPve33vlWYSoGYefn3gLQRzjfDlhSJ9ZKM= -github.com/jackc/pgx/v5 v5.9.2 h1:3ZhOzMWnR4yJ+RW1XImIPsD1aNSz4T4fyP7zlQb56hw= -github.com/jackc/pgx/v5 v5.9.2/go.mod h1:mal1tBGAFfLHvZzaYh77YS/eC6IX9OWbRV1QIIM0Jn4= github.com/jackc/pgx/v5 v5.10.0 h1:VhSvgU2jSli8o3AqIEOTJr7rZwAEUVo4E4XhR94Zfr0= github.com/jackc/pgx/v5 v5.10.0/go.mod h1:mal1tBGAFfLHvZzaYh77YS/eC6IX9OWbRV1QIIM0Jn4= github.com/jackc/puddle/v2 v2.2.2 h1:PR8nw+E/1w0GLuRFSmiioY6UooMp6KJv0/61nB7icHo= github.com/jackc/puddle/v2 v2.2.2/go.mod h1:vriiEXHvEE654aYKXXjOvZM39qJ0q+azkZFrfEOc3H4= -github.com/klauspost/compress v1.18.0 h1:c/Cqfb0r+Yi+JtIEq73FWXVkRonBlf0CRNYc8Zttxdo= -github.com/klauspost/compress v1.18.0/go.mod h1:2Pp+KzxcywXVXMr50+X0Q/Lsb43OQHYWRCY2AiWywWQ= -github.com/kylelemons/godebug v1.1.0 h1:RPNrshWIDI6G2gRW9EHilWtl7Z6Sb1BR0xunSBf0SNc= -github.com/kylelemons/godebug v1.1.0/go.mod h1:9/0rRGxNHcop5bhtWyNeEfOS8JIWk580+fNqagV/RAw= -github.com/lann/builder v0.0.0-20180802200727-47ae307949d0 h1:SOEGU9fKiNWd/HOJuq6+3iTQz8KNCLtVX6idSoTLdUw= -github.com/lann/builder v0.0.0-20180802200727-47ae307949d0/go.mod h1:dXGbAdH5GtBTC4WfIxhKZfyBF/HBFgRZSWwZ9g/He9o= -github.com/lann/ps v0.0.0-20150810152359-62de8c46ede0 h1:P6pPBnrTSX3DEVR4fDembhRWSsG5rVo6hYhAB/ADZrk= -github.com/lann/ps v0.0.0-20150810152359-62de8c46ede0/go.mod h1:vmVJ0l/dxyfGW6FmdpVm2joNMFikkuWg0EoCKLGUMNw= -github.com/lib/pq v1.12.3 h1:tTWxr2YLKwIvK90ZXEw8GP7UFHtcbTtty8zsI+YjrfQ= -github.com/lib/pq v1.12.3/go.mod h1:/p+8NSbOcwzAEI7wiMXFlgydTwcgTr3OSKMsD2BitpA= -github.com/moby/docker-image-spec v1.3.1 h1:jMKff3w6PgbfSa69GfNg+zN/XLhfXJGnEx3Nl2EsFP0= -github.com/moby/docker-image-spec v1.3.1/go.mod h1:eKmb5VW8vQEh/BAr2yvVNvuiJuY6UIocYsFu/DxxRpo= -github.com/moby/term v0.5.0 h1:xt8Q1nalod/v7BqbG21f8mQPqH+xAaC9C3N3wfWbVP0= -github.com/moby/term v0.5.0/go.mod h1:8FzsFHVUBGZdbDsJw/ot+X+d5HLUbvklYLJ9uGfcI3Y= -github.com/morikuni/aec v1.0.0 h1:nP9CBfwrvYnBRgY6qfDQkygYDmYwOilePFkwzv4dU8A= -github.com/morikuni/aec v1.0.0/go.mod h1:BbKIizmSmc5MMPqRYbxO4ZU0S0+P200+tUnFx7PXmsc= -github.com/munnerz/goautoneg v0.0.0-20191010083416-a7dc8b61c822 h1:C3w9PqII01/Oq1c1nUAm88MOHcQC9l5mIlSMApZMrHA= -github.com/munnerz/goautoneg v0.0.0-20191010083416-a7dc8b61c822/go.mod h1:+n7T8mK8HuQTcFwEeznm/DIxMOiR9yIdICNftLE1DvQ= -github.com/opencontainers/go-digest v1.0.0 h1:apOUWs51W5PlhuyGyz9FCeeBIOUDA/6nW8Oi/yOhh5U= -github.com/opencontainers/go-digest v1.0.0/go.mod h1:0JzlMkj0TRzQZfJkVvzbP0HBR3IKzErnv2BNG4W4MAM= -github.com/opencontainers/image-spec v1.1.0 h1:8SG7/vwALn54lVB/0yZ/MMwhFrPYtpEHQb2IpWsCzug= -github.com/opencontainers/image-spec v1.1.0/go.mod h1:W4s4sFTMaBeK1BQLXbG4AdM2szdn85PY75RI83NrTrM= -github.com/pkg/errors v0.9.1 h1:FEBLx1zS214owpjy7qsBeixbURkuhQAwrK5UwLGTwt4= -github.com/pkg/errors v0.9.1/go.mod h1:bwawxfHBFNV+L2hUp1rHADufV3IMtnDRdf1r5NINEl0= +github.com/pmezard/go-difflib v1.0.0 h1:4DBwDE0NGyQoBHbLQYPwSUPoCMWR5BEzIk/f1lZbAQM= github.com/pmezard/go-difflib v1.0.0/go.mod h1:iKH77koFhYxTK1pcRnkKkqfTogsbg7gZNVY4sRDYZ/4= -github.com/pmezard/go-difflib v1.0.1-0.20181226105442-5d4384ee4fb2 h1:Jamvg5psRIccs7FGNTlIRMkT8wgtp5eCXdBlqhYGL6U= -github.com/pmezard/go-difflib v1.0.1-0.20181226105442-5d4384ee4fb2/go.mod h1:iKH77koFhYxTK1pcRnkKkqfTogsbg7gZNVY4sRDYZ/4= -github.com/prometheus/client_golang v1.23.2 h1:Je96obch5RDVy3FDMndoUsjAhG5Edi49h0RJWRi/o0o= -github.com/prometheus/client_golang v1.23.2/go.mod h1:Tb1a6LWHB3/SPIzCoaDXI4I8UHKeFTEQ1YCr+0Gyqmg= -github.com/prometheus/client_model v0.6.2 h1:oBsgwpGs7iVziMvrGhE53c/GrLUsZdHnqNwqPLxwZyk= -github.com/prometheus/client_model v0.6.2/go.mod h1:y3m2F6Gdpfy6Ut/GBsUqTWZqCUvMVzSfMLjcu6wAwpE= -github.com/prometheus/common v0.68.0 h1:8rQJvQmYltsR2L7h8Zw0Iyj8WYNNmpwikoQTZXwfVeA= -github.com/prometheus/common v0.68.0/go.mod h1:4soH+U8yJSROk7OJ//hmTiWKsxapv6zRGgTt3keN8gQ= -github.com/prometheus/common v0.69.0 h1:OA85nJQS/T/MaYh/Q2CcgDKSGWqNIgrBDvDH85CuiNk= -github.com/prometheus/common v0.69.0/go.mod h1:ZzL3f6u94qUxh9p+tJTrF+FvBS1XXbbRAZCQkytAL0Y= -github.com/prometheus/procfs v0.20.1 h1:XwbrGOIplXW/AU3YhIhLODXMJYyC1isLFfYCsTEycfc= -github.com/prometheus/procfs v0.20.1/go.mod h1:o9EMBZGRyvDrSPH1RqdxhojkuXstoe4UlK79eF5TGGo= github.com/stretchr/objx v0.1.0/go.mod h1:HFkY916IF+rwdDfMAkV7OtwuqBVzrE8GR6GFx+wExME= -github.com/stretchr/testify v1.2.2/go.mod h1:a8OnRcib4nhh0OaRAV+Yts87kKdq0PP7pXfy6kDkUVs= github.com/stretchr/testify v1.3.0/go.mod h1:M5WIy9Dh21IEIfnGCwXGc5bZfKNJtfHm1UVUgZn+9EI= github.com/stretchr/testify v1.7.0/go.mod h1:6Fq8oRcR53rry900zMqJjRRixrwX3KX962/h/Wwjteg= github.com/stretchr/testify v1.11.1 h1:7s2iGBzp5EwR7/aIZr8ao5+dra3wiQyKjjFuvgVKu7U= github.com/stretchr/testify v1.11.1/go.mod h1:wZwfW3scLgRK+23gO65QZefKpKQRnfz6sD981Nm4B6U= -go.opentelemetry.io/auto/sdk v1.2.1 h1:jXsnJ4Lmnqd11kwkBV2LgLoFMZKizbCi5fNZ/ipaZ64= -go.opentelemetry.io/auto/sdk v1.2.1/go.mod h1:KRTj+aOaElaLi+wW1kO/DZRXwkF4C5xPbEe3ZiIhN7Y= -go.opentelemetry.io/contrib/instrumentation/net/http/otelhttp v0.61.0 h1:F7Jx+6hwnZ41NSFTO5q4LYDtJRXBf2PD0rNBkeB/lus= -go.opentelemetry.io/contrib/instrumentation/net/http/otelhttp v0.61.0/go.mod h1:UHB22Z8QsdRDrnAtX4PntOl36ajSxcdUMt1sF7Y6E7Q= -go.opentelemetry.io/otel v1.44.0 h1:JjwHmHpA4iZ3wBxluu2fbbE7j4kqlE8jXyAyPXH7HqU= -go.opentelemetry.io/otel v1.44.0/go.mod h1:BMgjTHL9WPRlRjL2oZCBTL4whCGtXch2H4BhOPIAyYc= -go.opentelemetry.io/otel/metric v1.44.0 h1:1w0gILTcHdr3YI+ixLyjemwrVnsMURbTZFrSYCdDdmc= -go.opentelemetry.io/otel/metric v1.44.0/go.mod h1:8O7hanEPBNgEMmybD3s2VBKcgWOCsA6tzHBPODAiquo= -go.opentelemetry.io/otel/sdk v1.43.0 h1:pi5mE86i5rTeLXqoF/hhiBtUNcrAGHLKQdhg4h4V9Dg= -go.opentelemetry.io/otel/sdk v1.43.0/go.mod h1:P+IkVU3iWukmiit/Yf9AWvpyRDlUeBaRg6Y+C58QHzg= -go.opentelemetry.io/otel/sdk/metric v1.43.0 h1:S88dyqXjJkuBNLeMcVPRFXpRw2fuwdvfCGLEo89fDkw= -go.opentelemetry.io/otel/sdk/metric v1.43.0/go.mod h1:C/RJtwSEJ5hzTiUz5pXF1kILHStzb9zFlIEe85bhj6A= -go.opentelemetry.io/otel/trace v1.44.0 h1:jxF5CsGYCe74MCRx2X4g7WsY/VBKRqqpNvXlX/6gtIk= -go.opentelemetry.io/otel/trace v1.44.0/go.mod h1:oLl1jrMQAVo6v3GAggN+1VH9VIz9iUSvW53sW1Q8PIE= -go.uber.org/goleak v1.3.0 h1:2K3zAYmnTNqV73imy9J1T3WC+gmCePx2hEGkimedGto= -go.uber.org/goleak v1.3.0/go.mod h1:CoHD4mav9JJNrW/WLlf7HGZPjdw8EucARQHekz1X6bE= -go.yaml.in/yaml/v2 v2.4.4 h1:tuyd0P+2Ont/d6e2rl3be67goVK4R6deVxCUX5vyPaQ= -go.yaml.in/yaml/v2 v2.4.4/go.mod h1:gMZqIpDtDqOfM0uNfy0SkpRhvUryYH0Z6wdMYcacYXQ= -golang.org/x/sync v0.20.0 h1:e0PTpb7pjO8GAtTs2dQ6jYa5BWYlMuX047Dco/pItO4= -golang.org/x/sync v0.20.0/go.mod h1:9xrNwdLfx4jkKbNva9FpL6vEN7evnE43NNNJQ2LF3+0= -golang.org/x/sync v0.21.0 h1:HLII4xRRTtCRkxYp4HNFF0Js/Og6q2i++KXbg0gHCwM= -golang.org/x/sync v0.21.0/go.mod h1:9xrNwdLfx4jkKbNva9FpL6vEN7evnE43NNNJQ2LF3+0= -golang.org/x/sys v0.45.0 h1:dO4czNzziLiiXplLQgBCEpCvXQ3dnkn0SdaZSYdQ+FY= -golang.org/x/sys v0.45.0/go.mod h1:4GL1E5IUh+htKOUEOaiffhrAeqysfVGipDYzABqnCmw= -golang.org/x/sys v0.46.0 h1:noSf2Fq6F8DBgS+LysIkx7rIExoNHJsxOAtPp4rthXw= -golang.org/x/sys v0.46.0/go.mod h1:4GL1E5IUh+htKOUEOaiffhrAeqysfVGipDYzABqnCmw= -golang.org/x/text v0.37.0 h1:Cqjiwd9eSg8e0QAkyCaQTNHFIIzWtidPahFWR83rTrc= -golang.org/x/text v0.37.0/go.mod h1:a5sjxXGs9hsn/AJVwuElvCAo9v8QYLzvavO5z2PiM38= -golang.org/x/text v0.38.0 h1:sXmwo9DwP3OK9EZ7PqAdaooSGozfl/3a6/xJcbzPRhE= -golang.org/x/text v0.38.0/go.mod h1:YXZt3QhHUKYT53r2lLKFIVi6Ao1jdzrTR/KQ09qyxF4= -google.golang.org/protobuf v1.36.11 h1:fV6ZwhNocDyBLK0dj+fg8ektcVegBBuEolpbTQyBNVE= -google.golang.org/protobuf v1.36.11/go.mod h1:HTf+CrKn2C3g5S8VImy6tdcUvCska2kB7j23XfzDpco= +golang.org/x/sync v0.22.0 h1:SZjpbeLmrCk4xhRSZFNZW5gFUeCeFgjekvI/+gfScek= +golang.org/x/sync v0.22.0/go.mod h1:9xrNwdLfx4jkKbNva9FpL6vEN7evnE43NNNJQ2LF3+0= +golang.org/x/text v0.40.0 h1:Ub2Z6/xjgF1WrYQz2nuITOEegKFtiIy+rieRJ5lHZKs= +golang.org/x/text v0.40.0/go.mod h1:hpnzDAfGV753zIKo+wk3u1bVKCGPbrnF7+7LBF/UHVY= gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0= gopkg.in/yaml.v3 v3.0.0-20200313102051-9f266ea9e77c/go.mod h1:K4uyk7z7BCEPqu6E+C64Yfv1cQ7kz7rIZviUmN+EgEM= gopkg.in/yaml.v3 v3.0.1 h1:fxVm/GzAzEWqLHuvctI91KS9hhNmmWOoWu0XTYJS7CA= diff --git a/internal/pkg/pgxpoolcollector/v5/collector.go b/internal/pkg/pgxpoolcollector/v5/collector.go deleted file mode 100644 index f9cd7ed..0000000 --- a/internal/pkg/pgxpoolcollector/v5/collector.go +++ /dev/null @@ -1,181 +0,0 @@ -// Package v5 provides prometheus plug-in metrics for a pgx client. -// -// This package tracks the following metrics under the following names: -// -// #ns_postgres_acquire_count{} -// #ns_postgres_acquire_duration{} -// #ns_postgres_acquired_conns{} -// #ns_postgres_canceled_acquire_count{} -// #ns_postgres_constructing_conns{} -// #ns_postgres_empty_acquire_count{} -// #ns_postgres_idle_conns{} -// #ns_postgres_max_conns{} -// #ns_postgres_total_conns{} -// -// Labels list: -// client_id=#{client_id} -// client_kind=#{master/replica} -// db=#{db} -// shard_id=#{shard_id} - -package v5 - -import ( - "github.com/jackc/pgx/v5/pgxpool" - "github.com/prometheus/client_golang/prometheus" -) - -// StatsGetter is an interface that gets sql.DBStats. -// It's implemented by e.g. *sql.DB or *sqlx.DB. -type StatsGetter interface { - Stat() *pgxpool.Stat -} - -// StatsCollector implements the prometheus.Collector interface. -type StatsCollector struct { - sg StatsGetter - - // descriptions of exported metrics - acquireCount *prometheus.Desc - acquireDuration *prometheus.Desc - acquiredConns *prometheus.Desc - canceledAcquireCount *prometheus.Desc - constructingConns *prometheus.Desc - emptyAcquireCount *prometheus.Desc - idleConns *prometheus.Desc - maxConns *prometheus.Desc - totalConns *prometheus.Desc -} - -// NewStatsCollector creates a new StatsCollector. -func NewStatsCollector(namespace, subsystem string, constLabels prometheus.Labels, sg StatsGetter) *StatsCollector { - return &StatsCollector{ - sg: sg, - acquireCount: prometheus.NewDesc( - prometheus.BuildFQName(namespace, subsystem, "acquire_count"), - "Cumulative count of successful acquires from the pool.", - nil, - constLabels, - ), - acquireDuration: prometheus.NewDesc( - prometheus.BuildFQName(namespace, subsystem, "acquire_duration"), - "Total duration of all successful acquires from the pool.", - nil, - constLabels, - ), - acquiredConns: prometheus.NewDesc( - prometheus.BuildFQName(namespace, subsystem, "acquired_conns"), - "Number of currently acquired connections in the pool.", - nil, - constLabels, - ), - canceledAcquireCount: prometheus.NewDesc( - prometheus.BuildFQName(namespace, subsystem, "canceled_acquire_count"), - "Cumulative count of acquires from the pool that were canceled by a context.", - nil, - constLabels, - ), - constructingConns: prometheus.NewDesc( - prometheus.BuildFQName(namespace, subsystem, "constructing_conns"), - "Number of conns with construction in progress in the pool.", - nil, - constLabels, - ), - emptyAcquireCount: prometheus.NewDesc( - prometheus.BuildFQName(namespace, subsystem, "empty_acquire_count"), - "Cumulative count of successful acquires from the pool that waited for a resource to be released or constructed because the pool was empty.", - nil, - constLabels, - ), - idleConns: prometheus.NewDesc( - prometheus.BuildFQName(namespace, subsystem, "idle_conns"), - "Number of currently idle conns in the pool.", - nil, - constLabels, - ), - maxConns: prometheus.NewDesc( - prometheus.BuildFQName(namespace, subsystem, "max_conns"), - "Maximum size of the pool.", - nil, - constLabels, - ), - totalConns: prometheus.NewDesc( - prometheus.BuildFQName(namespace, subsystem, "total_conns"), - "Total number of resources currently in the pool. The value is the sum of ConstructingConns, AcquiredConns, and IdleConns.", - nil, - constLabels, - ), - } -} - -// Describe implements the prometheus.Collector interface. -func (c StatsCollector) Describe(ch chan<- *prometheus.Desc) { - ch <- c.acquireCount - ch <- c.acquireDuration - ch <- c.acquiredConns - ch <- c.canceledAcquireCount - ch <- c.constructingConns - ch <- c.emptyAcquireCount - ch <- c.idleConns - ch <- c.maxConns - ch <- c.totalConns -} - -// Collect implements the prometheus.Collector interface. -func (c StatsCollector) Collect(ch chan<- prometheus.Metric) { - stats := c.sg.Stat() - - ch <- prometheus.MustNewConstMetric( - c.acquireCount, - prometheus.GaugeValue, - float64(stats.AcquireCount()), - ) - - ch <- prometheus.MustNewConstMetric( - c.acquireDuration, - prometheus.GaugeValue, - float64(stats.AcquireDuration()), - ) - - ch <- prometheus.MustNewConstMetric( - c.acquiredConns, - prometheus.GaugeValue, - float64(stats.AcquiredConns()), - ) - - ch <- prometheus.MustNewConstMetric( - c.canceledAcquireCount, - prometheus.GaugeValue, - float64(stats.CanceledAcquireCount()), - ) - - ch <- prometheus.MustNewConstMetric( - c.constructingConns, - prometheus.GaugeValue, - float64(stats.ConstructingConns()), - ) - - ch <- prometheus.MustNewConstMetric( - c.emptyAcquireCount, - prometheus.GaugeValue, - float64(stats.EmptyAcquireCount()), - ) - - ch <- prometheus.MustNewConstMetric( - c.idleConns, - prometheus.GaugeValue, - float64(stats.IdleConns()), - ) - - ch <- prometheus.MustNewConstMetric( - c.maxConns, - prometheus.GaugeValue, - float64(stats.MaxConns()), - ) - - ch <- prometheus.MustNewConstMetric( - c.totalConns, - prometheus.GaugeValue, - float64(stats.TotalConns()), - ) -} diff --git a/internal/pkg/pgxslog/adapter.go b/internal/pkg/pgxslog/adapter.go deleted file mode 100644 index 9651cf4..0000000 --- a/internal/pkg/pgxslog/adapter.go +++ /dev/null @@ -1,63 +0,0 @@ -package pgxslog - -import ( - "context" - "log/slog" - - "github.com/jackc/pgx/v5/tracelog" -) - -type Logger struct { - l *slog.Logger -} - -func NewLogger(l *slog.Logger) *Logger { - return &Logger{l} -} - -func (l *Logger) Log(ctx context.Context, level tracelog.LogLevel, msg string, data map[string]interface{}) { - attrs := make([]slog.Attr, 0, len(data)) - for k, v := range data { - attrs = append(attrs, slog.Any(k, v)) - } - - var lvl slog.Level - switch level { - case tracelog.LogLevelTrace: - lvl = slog.LevelDebug - 1 - attrs = append(attrs, slog.Any("pgx_log_level", level)) - case tracelog.LogLevelDebug: - lvl = slog.LevelDebug - case tracelog.LogLevelInfo: - lvl = slog.LevelInfo - case tracelog.LogLevelWarn: - lvl = slog.LevelWarn - case tracelog.LogLevelError: - lvl = slog.LevelError - default: - lvl = slog.LevelError - attrs = append(attrs, slog.Any("invalid_pgx_log_level", level)) - } - //nolint:sloglint // there is no other option to pass the message - l.l.LogAttrs(ctx, lvl, msg, attrs...) -} - -// all key constants must be defined with the "Key" suffix. - -// Constants and constructors to standardize the keys used in logs. - -const ( - // errorKey - to pass error to log. - errorKey string = "error" - // componentKey - to identify the component that is logging (e.g.: "kafka-consumer") - componentKey = "component" -) - -// Error for passing error to log. -func Error(err error) slog.Attr { return slog.Any(errorKey, err) } - -// Component to identify the component that is logging (e.g.: "kafka-consumer"). -// here any distinct level of your application abstraction can be used. -// it's not mandatory associated with kind of an external connection, -// e.g.: "location-event-processor" is also applicable for the field. -func Component(component string) slog.Attr { return slog.String(componentKey, component) } diff --git a/internal/pkg/pgxtracer/tracer.go b/internal/pkg/pgxtracer/tracer.go deleted file mode 100644 index b92fec0..0000000 --- a/internal/pkg/pgxtracer/tracer.go +++ /dev/null @@ -1,40 +0,0 @@ -// Package pgxtracer is a chaining QueryTracer for pgx. -package pgxtracer - -import ( - "context" - - "github.com/jackc/pgx/v5" -) - -// QueryTracer traces Query, QueryRow, and Exec. -type QueryTracer interface { - // TraceQueryStart is called at the beginning of Query, QueryRow, and Exec calls. The returned context is used for the - // rest of the call and will be passed to TraceQueryEnd. - TraceQueryStart(ctx context.Context, conn *pgx.Conn, data pgx.TraceQueryStartData) context.Context - - TraceQueryEnd(ctx context.Context, conn *pgx.Conn, data pgx.TraceQueryEndData) -} - -type tracer struct { - tracers []QueryTracer -} - -func (t *tracer) TraceQueryStart(ctx context.Context, conn *pgx.Conn, data pgx.TraceQueryStartData) context.Context { - for i := range t.tracers { - ctx = t.tracers[i].TraceQueryStart(ctx, conn, data) - } - return ctx -} - -func (t *tracer) TraceQueryEnd(ctx context.Context, conn *pgx.Conn, data pgx.TraceQueryEndData) { - for i := range t.tracers { - t.tracers[i].TraceQueryEnd(ctx, conn, data) - } -} - -func New(tracers ...QueryTracer) QueryTracer { - return &tracer{ - tracers: tracers, - } -} diff --git a/migrate.go b/migrate.go deleted file mode 100644 index a23e882..0000000 --- a/migrate.go +++ /dev/null @@ -1,52 +0,0 @@ -package postgres - -import ( - "context" - "embed" - "errors" - "log/slog" - - "github.com/golang-migrate/migrate/v4" - _ "github.com/golang-migrate/migrate/v4/database/postgres" // init - "github.com/golang-migrate/migrate/v4/source" - "github.com/golang-migrate/migrate/v4/source/iofs" -) - -func applyMigrations(fs embed.FS, dsn string, l *slog.Logger) (err error) { - var src source.Driver - ctx := context.Background() - - src, err = iofs.New(fs, ".") - if err != nil { - err = errors.Join(errors.New("embed.FS init failed"), err) - return err - } - defer func() { - if ce := src.Close(); ce != nil { - err = errors.Join(err, ce) - } - }() - - var instance *migrate.Migrate - instance, err = migrate.NewWithSourceInstance("iofs", src, dsn) - if err != nil { - err = errors.Join(errors.New("db instance init failed"), err) - return err - } - defer func() { - if sErr, ie := instance.Close(); sErr != nil || ie != nil { - err = errors.Join(err, sErr, ie) - } - }() - - if mErr := instance.Up(); mErr != nil && !errors.Is(mErr, migrate.ErrNoChange) { - err = errors.Join(errors.New("migrate-up failed"), mErr) - } else { - ver, dirty, _ := instance.Version() - l.InfoContext(ctx, "migrate-up done", - slog.Any("version", ver), - slog.Any("dirty", dirty)) - } - - return err -} diff --git a/options.go b/options.go index 2ba07d3..561f696 100644 --- a/options.go +++ b/options.go @@ -1,239 +1,113 @@ -package postgres +package xpg import ( - "embed" + "errors" "fmt" - "log/slog" - "runtime" - "time" - - "github.com/jackc/pgx/v5" - "go.opentelemetry.io/otel/trace" + "strings" ) -// An Option lets you add opts to pberrors interceptors using With* funcs. +// Option configures a Pool. +// +// The interface is sealed so options can only be created by this package. type Option interface { - apply(p *Pool) + apply(*settings) error } -type optionFunc func(p *Pool) - -func (f optionFunc) apply(p *Pool) { - f(p) -} +type optionFunc func(*settings) error -func WithLogger(l *slog.Logger) Option { - return optionFunc(func(p *Pool) { - if l != nil { - p.logger = l - } - }) +func (option optionFunc) apply(settings *settings) error { + return option(settings) } -func WithConfig(config *Config) Option { - return optionFunc(func(p *Pool) { - if config != nil { - p.cfg = config - } - }) +type settings struct { + name string + labels map[string]string } -func WithClientID(id string) Option { - return optionFunc(func(p *Pool) { - if id != "" { - p.id = fmt.Sprintf("%s-%s", id, GenerateUUID()) - } - }) -} - -func WithTraceProvider(provider trace.TracerProvider) Option { - return optionFunc(func(p *Pool) { - p.traceProvider = provider - }) +func defaultSettings() *settings { + return &settings{ + labels: make(map[string]string), + } } -func WithMigrations(migrations ...embed.FS) Option { - return optionFunc(func(p *Pool) { - if len(migrations) > 0 { - p.migrations = migrations +func applyOptions(settings *settings, opts ...Option) error { + for _, opt := range opts { + if opt == nil { + return errors.New("xpg: option is nil") } - }) -} -func WithMetricsNamespace(ns string) Option { - return optionFunc(func(p *Pool) { - if ns != "" { - p.namespace = ns + if err := opt.apply(settings); err != nil { + return fmt.Errorf("xpg: apply option: %w", err) } - }) -} - -type Config struct { - ShardID int `envconfig:"POSTGRES_SHARD_ID"` - ClusterHost string `envconfig:"POSTGRES_CLUSTER_HOST" required:"true"` - ClusterPort string `envconfig:"POSTGRES_CLUSTER_PORT" required:"true"` - ClusterReplicaPort string `envconfig:"POSTGRES_CLUSTER_REPLICA_PORT" required:"true"` - User string `envconfig:"POSTGRES_USER" required:"true"` - Password string `envconfig:"POSTGRES_PASSWORD" required:"true"` - DB string `envconfig:"POSTGRES_DB" required:"true"` - - MinRWConn int32 `envconfig:"POSTGRES_MIN_RW_CONN"` - MinROConn int32 `envconfig:"POSTGRES_MIN_RO_CONN"` - MaxRWConn int32 `envconfig:"POSTGRES_MAX_RW_CONN"` - MaxROConn int32 `envconfig:"POSTGRES_MAX_RO_CONN"` - MaxConnLifetime time.Duration `envconfig:"POSTGRES_MAX_CONN_LIFETIME"` - MaxConnIdleTime time.Duration `envconfig:"POSTGRES_MAX_CONN_IDLE_TIME"` - QueryExecMode string `envconfig:"POSTGRES_QUERY_EXEC_MODE"` - StatementCacheCapacity int `envconfig:"POSTGRES_STATEMENT_CACHE_CAPACITY"` - DescriptionCacheCapacity int `envconfig:"POSTGRES_DESCRIPTION_CACHE_CAPACITY"` - - MasterArgs string `envconfig:"POSTGRES_MASTER_ARGS"` - ReplicaArgs string `envconfig:"POSTGRES_REPLICA_ARGS"` - - MigrateEnabled bool `envconfig:"POSTGRES_MIGRATE_ENABLED"` - MigrateArgs string `envconfig:"POSTGRES_MIGRATE_ARGS"` - MigratePort string `envconfig:"POSTGRES_MIGRATE_PORT"` - - writer bool - appName string -} - -// DSN postgres://username:password@host:port/db?sslmode=disable& -func (c *Config) getDSN() string { - return formatDSN(c.User, c.Password, c.ClusterHost, c.getPort(), c.DB, c.appName, c.getArgs()) -} - -func (c *Config) getMigrateDSN() string { - return formatDSN(c.User, c.Password, c.ClusterHost, c.getMigratePort(), c.DB, c.appName, c.MigrateArgs) -} - -func (c *Config) getPort() string { - if c.writer { - return c.ClusterPort } - return c.ClusterReplicaPort -} -func (c *Config) getMigratePort() string { - if c.MigratePort != "" { - return c.MigratePort - } - return c.ClusterPort + return nil } -func (c *Config) getArgs() string { - if c.writer { - return c.MasterArgs - } - return c.ReplicaArgs -} +// WithName assigns a stable logical name to the pool. +// +// Name is metadata for diagnostics and observability. It does not change the +// PostgreSQL application_name runtime parameter. +func WithName(name string) Option { + name = strings.TrimSpace(name) -func (c *Config) getMinConns() int32 { - if c.writer { - return c.MinRWConn - } - return c.MinROConn -} - -func (c *Config) getMaxConns() int32 { - if c.writer { - return c.MaxRWConn - } - return c.MaxROConn -} - -func (c *Config) resolvedMinConns() int32 { - if v := c.getMinConns(); v > 0 { - return v - } - return 1 -} + return optionFunc(func(settings *settings) error { + if name == "" { + return errors.New("pool name must not be blank") + } -func (c *Config) resolvedMaxConns() int32 { - v := c.getMaxConns() - if v <= 0 { - v = 4 - } - if numCPU := int32(runtime.NumCPU()); numCPU > v { - return numCPU - } - return v -} + settings.name = name -func formatDSN(user, pass, host, port, db, appName, args string) string { - return fmt.Sprintf("postgres://%s:%s@%s:%s/%s?sslmode=disable&application_name=%s&%s", - user, pass, host, port, db, appName, args, - ) + return nil + }) } -const ( - QueryExecModeCacheStatement = "cache_statement" - QueryExecModeCacheDescribe = "cache_describe" - QueryExecModeDescribeExec = "describe_exec" - QueryExecModeExec = "exec" - QueryExecModeSimpleProtocol = "simple_protocol" -) - -func getQueryExecMode(mode string) pgx.QueryExecMode { - switch mode { - case QueryExecModeCacheStatement: - return pgx.QueryExecModeCacheStatement +// WithLabels merges labels into the pool metadata. +// +// Labels are defensively copied. When the same key is configured more than +// once, the last value wins. +func WithLabels(labels map[string]string) Option { + labels = cloneLabels(labels) - case QueryExecModeCacheDescribe: - return pgx.QueryExecModeCacheDescribe + return optionFunc(func(settings *settings) error { + for key, value := range labels { + key = strings.TrimSpace(key) + if key == "" { + return errors.New("label key must not be blank") + } - case QueryExecModeDescribeExec: - return pgx.QueryExecModeDescribeExec - - case QueryExecModeExec: - return pgx.QueryExecModeExec - - case QueryExecModeSimpleProtocol: - return pgx.QueryExecModeSimpleProtocol + settings.labels[key] = value + } - default: - return pgx.QueryExecModeCacheStatement - } + return nil + }) } -func parseConfig(cfg *Config) *options { - o := &options{ - dsn: cfg.getDSN(), - minConns: cfg.resolvedMinConns(), - maxConns: cfg.resolvedMaxConns(), - maxConnLifetime: time.Minute * 1, - maxConnIdleTime: time.Second * 30, - defaultQueryExecMode: getQueryExecMode(cfg.QueryExecMode), - statementCacheCapacity: 128, - descriptionCacheCapacity: 512, - } - - if cfg.getMinConns() > 0 { - o.minConns = cfg.getMinConns() - } +// WithLabel adds or replaces one pool label. +func WithLabel(key, value string) Option { + key = strings.TrimSpace(key) - if cfg.getMaxConns() > 0 { - o.maxConns = cfg.getMaxConns() - if numCPU := int32(runtime.NumCPU()); numCPU > cfg.getMaxConns() { - o.maxConns = numCPU + return optionFunc(func(settings *settings) error { + if key == "" { + return errors.New("label key must not be blank") } - } - if cfg.MaxConnLifetime > 0 { - o.maxConnLifetime = cfg.MaxConnLifetime - } + settings.labels[key] = value - if cfg.MaxConnIdleTime > 0 { - o.maxConnIdleTime = cfg.MaxConnIdleTime - } + return nil + }) +} - if cfg.StatementCacheCapacity > 0 { - o.statementCacheCapacity = cfg.StatementCacheCapacity +func cloneLabels(labels map[string]string) map[string]string { + if len(labels) == 0 { + return nil } - if cfg.DescriptionCacheCapacity > 0 { - o.descriptionCacheCapacity = cfg.DescriptionCacheCapacity + cloned := make(map[string]string, len(labels)) + + for key, value := range labels { + cloned[key] = value } - return o + return cloned } diff --git a/pool.go b/pool.go index 6ad6756..3ff64a1 100644 --- a/pool.go +++ b/pool.go @@ -1,260 +1,103 @@ -package postgres +package xpg import ( "context" - "embed" "errors" "fmt" - "log/slog" - "strconv" - "time" + "sync" - "github.com/Masterminds/squirrel" - "github.com/exaring/otelpgx" "github.com/jackc/pgx/v5" "github.com/jackc/pgx/v5/pgconn" "github.com/jackc/pgx/v5/pgxpool" - "github.com/jackc/pgx/v5/tracelog" - poolcollector "github.com/mkbeh/xpg/internal/pkg/pgxpoolcollector/v5" - "github.com/mkbeh/xpg/internal/pkg/pgxslog" - "github.com/mkbeh/xpg/internal/pkg/pgxtracer" - "github.com/prometheus/client_golang/prometheus" - "go.opentelemetry.io/otel" - "go.opentelemetry.io/otel/trace" ) +// Pool is a concurrency-safe PostgreSQL connection pool backed by pgxpool. type Pool struct { - *pgxpool.Pool + pool *pgxpool.Pool - id string - cfg *Config - logger *slog.Logger - traceProvider trace.TracerProvider - qBuilder squirrel.StatementBuilderType - migrations []embed.FS - namespace string - labels prometheus.Labels -} - -type options struct { - dsn string - minConns int32 - maxConns int32 - maxConnLifetime time.Duration - maxConnIdleTime time.Duration - statementCacheCapacity int - descriptionCacheCapacity int - defaultQueryExecMode pgx.QueryExecMode - logger *slog.Logger - traceProvider trace.TracerProvider - tracers []pgxtracer.QueryTracer -} + name string + labels map[string]string -func NewWriter(opts ...Option) (*Pool, error) { - return newPool(true, opts) + closeOnce sync.Once } -func NewReader(opts ...Option) (*Pool, error) { - return newPool(false, opts) -} - -func newPool(writer bool, opts []Option) (*Pool, error) { - p := &Pool{ - cfg: &Config{}, - logger: slog.Default(), - qBuilder: squirrel.StatementBuilder.PlaceholderFormat(squirrel.Dollar), - } - - for _, opt := range opts { - opt.apply(p) +// Open parses a DSN and creates a Pool. +func Open(ctx context.Context, dsn string, options ...Option) (*Pool, error) { + config, err := pgxpool.ParseConfig(dsn) + if err != nil { + return nil, fmt.Errorf("xpg: parse pool config: %w", err) } - p.cfg.writer = writer - p.cfg.appName = p.getID() - - if p.traceProvider == nil { - p.traceProvider = otel.GetTracerProvider() - } + return New(ctx, config, options...) +} - if writer { - p.logger = p.logger.With(pgxslog.Component("postgres_master")) - } else { - p.logger = p.logger.With(pgxslog.Component("postgres_replica")) +// New creates a Pool from config. +// +// Config must have been created by pgxpool.ParseConfig. New passes a defensive +// copy to pgxpool, so subsequent changes to the original config do not affect +// the created Pool. +// +// As with pgxpool.Config.Copy, the referenced tls.Config remains shared and +// must not be modified after it has been used to create connections. +func New(ctx context.Context, config *pgxpool.Config, options ...Option) (*Pool, error) { + if config == nil || config.ConnConfig == nil { + return nil, errors.New("xpg: pool config is nil") } - connOpts := parseConfig(p.cfg) - connOpts.logger = p.logger - connOpts.traceProvider = p.traceProvider + settings := defaultSettings() - conn, err := connect(connOpts) - if err != nil { + if err := applyOptions(settings, options...); err != nil { return nil, err } - p.Pool = conn - p.exposeMetrics(writer) - - collector := poolcollector.NewStatsCollector(p.namespace, "postgres", p.labels, p.Pool) - prometheus.MustRegister(collector) - - if p.cfg.writer && p.cfg.MigrateEnabled { - for _, fs := range p.migrations { - if err := applyMigrations(fs, p.cfg.getMigrateDSN(), p.logger); err != nil { - return nil, err - } - } + pgxPool, err := pgxpool.NewWithConfig(ctx, config.Copy()) + if err != nil { + return nil, fmt.Errorf("xpg: create pool: %w", err) } - return p, err + return &Pool{ + pool: pgxPool, + name: settings.name, + labels: cloneLabels(settings.labels), + }, nil } -func (p *Pool) QueryBuilder() squirrel.StatementBuilderType { - return p.qBuilder +// Name returns the logical pool name configured with WithName. +func (p *Pool) Name() string { + return p.name } -func (p *Pool) Logger() *slog.Logger { - return p.logger +// Raw returns the underlying pgxpool.Pool. +func (p *Pool) Raw() *pgxpool.Pool { + return p.pool } -func (p *Pool) Close() error { - p.Pool.Close() - return nil +func (p *Pool) Ping(ctx context.Context) error { + return p.pool.Ping(ctx) } -func (p *Pool) SendBatch(ctx context.Context, b *pgx.Batch) pgx.BatchResults { - if tx := extractTx(ctx); tx != nil { - return tx.SendBatch(ctx, b) - } - return p.Pool.SendBatch(ctx, b) +// Close closes the underlying pool and waits for acquired connections to be +// returned. Close is safe to call multiple times. +func (p *Pool) Close() { + p.closeOnce.Do(p.pool.Close) } func (p *Pool) Exec(ctx context.Context, sql string, arguments ...any) (pgconn.CommandTag, error) { - if tx := extractTx(ctx); tx != nil { - return tx.Exec(ctx, sql, arguments...) - } - return p.Pool.Exec(ctx, sql, arguments...) + return p.pool.Exec(ctx, sql, arguments...) } func (p *Pool) Query(ctx context.Context, sql string, args ...any) (pgx.Rows, error) { - if tx := extractTx(ctx); tx != nil { - return tx.Query(ctx, sql, args...) - } - return p.Pool.Query(ctx, sql, args...) + return p.pool.Query(ctx, sql, args...) } func (p *Pool) QueryRow(ctx context.Context, sql string, args ...any) pgx.Row { - if tx := extractTx(ctx); tx != nil { - return tx.QueryRow(ctx, sql, args...) - } - return p.Pool.QueryRow(ctx, sql, args...) -} - -// RunInTxx alias for RunInTx. -func (p *Pool) RunInTxx(ctx context.Context, fn func(ctx context.Context) error) error { - return p.RunInTx(ctx, fn, pgx.TxOptions{}) -} - -func (p *Pool) RunInTx(ctx context.Context, fn func(ctx context.Context) error, txOptions pgx.TxOptions) (err error) { - tx, err := p.Pool.BeginTx(ctx, txOptions) - if err != nil { - p.Logger().ErrorContext(ctx, "failed to begin transaction", pgxslog.Error(err)) - return NewPgError(ErrBeginTransaction, err) - } - - defer func() { - if r := recover(); r != nil { - p.Logger().ErrorContext(ctx, "panic recovered", slog.Any("error", r)) - err = NewPgError(ErrOther, fmt.Errorf("%v", r)) - } - - if rErr := tx.Rollback(ctx); rErr != nil { - if !errors.Is(rErr, pgx.ErrTxClosed) { - p.Logger().ErrorContext(ctx, "failed to rollback transaction", pgxslog.Error(rErr)) - } - } - }() - - if err = fn(injectTx(ctx, tx)); err != nil { - return err - } - - if err = tx.Commit(ctx); err != nil { - p.Logger().ErrorContext(ctx, "failed to commit transaction", pgxslog.Error(err)) - return NewPgError(ErrCommitTransaction, err) - } - - return nil -} - -func (p *Pool) AcquireTxLock(ctx context.Context, key string, durationSeconds float64) (isLocked bool, err error) { - row := p.QueryRow(ctx, ` - SELECT CASE - WHEN pg_try_advisory_xact_lock($1) THEN (SELECT concat(pg_sleep($2), 'false'))::bool - ELSE true - END AS is_locked;`, - int64(StringAsHash64(key)), - durationSeconds) - err = row.Scan(&isLocked) - return isLocked, err -} - -func (p *Pool) getID() string { - if p.id == "" { - return GenerateUUID() - } - return p.id + return p.pool.QueryRow(ctx, sql, args...) } -func (p *Pool) exposeMetrics(writer bool) { - if p.labels == nil { - p.labels = make(prometheus.Labels) - } - - p.labels["client_id"] = p.getID() - p.labels["db"] = p.cfg.DB - p.labels["shard_id"] = strconv.Itoa(p.cfg.ShardID) - - if writer { - p.labels["client_kind"] = "master" - } else { - p.labels["client_kind"] = "replica" - } +func (p *Pool) SendBatch(ctx context.Context, batch *pgx.Batch) pgx.BatchResults { + return p.pool.SendBatch(ctx, batch) } -func connect(opts *options) (*pgxpool.Pool, error) { - poolCfg, err := pgxpool.ParseConfig(opts.dsn) - if err != nil { - return nil, err - } - - opts.tracers = append(opts.tracers, - &tracelog.TraceLog{Logger: pgxslog.NewLogger(opts.logger), LogLevel: tracelog.LogLevelError}, - otelpgx.NewTracer( - otelpgx.WithTrimSQLInSpanName(), - otelpgx.WithTracerProvider(opts.traceProvider), - ), - ) - - poolCfg.MinConns = opts.minConns - poolCfg.MaxConns = opts.maxConns - poolCfg.MaxConnLifetime = opts.maxConnLifetime - poolCfg.MaxConnIdleTime = opts.maxConnIdleTime - - poolCfg.ConnConfig.StatementCacheCapacity = opts.statementCacheCapacity - poolCfg.ConnConfig.DescriptionCacheCapacity = opts.descriptionCacheCapacity - poolCfg.ConnConfig.DefaultQueryExecMode = opts.defaultQueryExecMode - poolCfg.ConnConfig.Tracer = pgxtracer.New(opts.tracers...) - - ctx := context.Background() - - pool, err := pgxpool.NewWithConfig(ctx, poolCfg) - if err != nil { - return nil, err - } - - if err := pool.Ping(ctx); err != nil { - return nil, err - } - - return pool, nil +func (p *Pool) CopyFrom(ctx context.Context, tableName pgx.Identifier, columnNames []string, rowSrc pgx.CopyFromSource) (int64, error) { + return p.pool.CopyFrom(ctx, tableName, columnNames, rowSrc) } diff --git a/tx.go b/tx.go deleted file mode 100644 index ed977b4..0000000 --- a/tx.go +++ /dev/null @@ -1,33 +0,0 @@ -package postgres - -import ( - "context" - - "github.com/jackc/pgx/v5" -) - -type ctxTx struct { - tx pgx.Tx -} - -type txKey struct{} - -var ( - txMarkerKey = &txKey{} - nullTx = &ctxTx{} -) - -func injectTx(ctx context.Context, tx pgx.Tx) context.Context { - t := &ctxTx{ - tx: tx, - } - return context.WithValue(ctx, txMarkerKey, t) -} - -func extractTx(ctx context.Context) pgx.Tx { - t, ok := ctx.Value(txMarkerKey).(*ctxTx) - if !ok || t == nil { - return nullTx.tx - } - return t.tx -} diff --git a/utils.go b/utils.go deleted file mode 100644 index 455cb02..0000000 --- a/utils.go +++ /dev/null @@ -1,17 +0,0 @@ -package postgres - -import ( - "hash/fnv" - - "github.com/google/uuid" -) - -func GenerateUUID() string { - return uuid.New().String() -} - -func StringAsHash64(s string) uint64 { - hash := fnv.New64() - _, _ = hash.Write([]byte(s)) - return hash.Sum64() -} From 090f02a3e1f02f4e1ba0a3fef3774f67e9138cfb Mon Sep 17 00:00:00 2001 From: mkbeh Date: Mon, 3 Aug 2026 17:40:31 +0300 Subject: [PATCH 02/41] feat: add transaction helpers and savepoints --- examples/README.md | 30 ++++- examples/docker-compose.yml | 32 ++++++ examples/transactions/README.md | 122 ++++++++++++++++++++ examples/transactions/go.mod | 7 ++ examples/transactions/main.go | 193 ++++++++++++++++++++++++++++++++ examples/transactions/setup.sql | 24 ++++ tx.go | 77 +++++++++++++ tx_test.go | 86 ++++++++++++++ 8 files changed, 570 insertions(+), 1 deletion(-) create mode 100644 examples/docker-compose.yml create mode 100644 examples/transactions/README.md create mode 100644 examples/transactions/go.mod create mode 100644 examples/transactions/main.go create mode 100644 examples/transactions/setup.sql create mode 100644 tx.go create mode 100644 tx_test.go diff --git a/examples/README.md b/examples/README.md index c2c8c48..6fe2f93 100644 --- a/examples/README.md +++ b/examples/README.md @@ -1 +1,29 @@ -[TODO] \ No newline at end of file +# Examples + +This directory contains runnable examples demonstrating the main features and usage patterns of `xpg`. + +| Example | Demonstrates | +|:-------------------------------|:------------------------------------------------------------------------------------------| +| [`basic`](basic) | Pool lifecycle and common query methods | +| [`transactions`](transactions) | Committing an outer transaction after an optional operation is rolled back to a savepoint | + +## Running the examples + +The examples use Docker Compose to start PostgreSQL and any required supporting services. + +From the `examples` directory, start PostgreSQL and Adminer: + +```shell +docker compose --profile tools up -d +``` + +Then run the example from its directory: + +```shell +cd transactions +go run . +``` + +> [!NOTE] +> Some examples may require different services or configuration. Refer to the README in the corresponding example +> directory for the exact startup command, connection settings, and expected output. diff --git a/examples/docker-compose.yml b/examples/docker-compose.yml new file mode 100644 index 0000000..4cb6461 --- /dev/null +++ b/examples/docker-compose.yml @@ -0,0 +1,32 @@ +services: + postgres: + image: postgres:18-alpine + environment: + POSTGRES_DB: postgres + POSTGRES_USER: postgres + POSTGRES_PASSWORD: postgres + ports: + - "5432:5432" + volumes: + - postgres-data:/var/lib/postgresql + healthcheck: + test: "pg_isready -U $${POSTGRES_USER} -d $${POSTGRES_DB}" + interval: 2s + timeout: 5s + retries: 10 + start_period: 5s + + adminer: + image: adminer:5.4.2-standalone + profiles: + - tools + environment: + ADMINER_DEFAULT_SERVER: postgres + ports: + - "8080:8080" + depends_on: + postgres: + condition: service_healthy + +volumes: + postgres-data: \ No newline at end of file diff --git a/examples/transactions/README.md b/examples/transactions/README.md new file mode 100644 index 0000000..8b666c8 --- /dev/null +++ b/examples/transactions/README.md @@ -0,0 +1,122 @@ +# Transactions and savepoints + +This example creates an order in a transaction and applies an optional promo code inside a PostgreSQL savepoint. +The promo code is already used, so only the savepoint is rolled back while the outer transaction commits the order. + +**This example demonstrates:** + +* Creating an `xpg.Pool` and explicitly checking PostgreSQL connectivity +* Executing a callback with an explicit `pgx.Tx` +* Isolating an optional operation with `xpg.InSavepoint` +* Inspecting a wrapped `pgconn.PgError` +* Continuing and committing the outer transaction after a savepoint rollback + +## Configuration + +The example connects to PostgreSQL using: + +```text +postgres://postgres:postgres@localhost:5432/postgres?sslmode=disable +``` + +Set `XPG_DATABASE_URL` to use another connection string: + +```shell +export XPG_DATABASE_URL='postgres://user:password@localhost:5432/database?sslmode=disable' +``` + +## Local PostgreSQL setup + +The example can use the local PostgreSQL setup from `examples/docker-compose.yml`. + +From the repository root: + +```shell +docker compose -f examples/docker-compose.yml --profile tools up -d +``` + +Or from this example directory: + +```shell +docker compose -f ../docker-compose.yml --profile tools up -d +``` + +To start only PostgreSQL, omit `--profile tools`. + +PostgreSQL is available to applications running on the host at: + +```text +localhost:5432 +``` + +Adminer is available at: + +```text +http://localhost:8080 +``` + +Sign in to Adminer with: + +```text +System: PostgreSQL +Server: postgres +Username: postgres +Password: postgres +Database: postgres +``` + +> [!IMPORTANT] +> Use `postgres`, not `localhost`, in the **Server** field. Adminer connects to PostgreSQL through the Docker Compose +> network, where the database is discoverable by its service name. + +## Run + +From this example directory: + +```shell +go run . +``` + +Or from the repository root: + +```shell +go run ./examples/transactions +``` + +## Expected output + +```text +order ID: 1 +order status: new +promo code: PROMO2026 +promo applied: false +``` + +The order is committed because the unique-key error is confined to the savepoint. `InSavepoint` rolls the failed promo +insert back before the outer transaction decides that this specific error is non-fatal. + +## Inspect the result + +The embedded `setup.sql` recreates the `xpg_transactions_example` schema before each run and leaves the resulting data +available for inspection. + +In Adminer, open the `xpg_transactions_example` schema and inspect: + +```text +orders +promo_redemptions +``` + +## Stop services + +From the repository root: + +```shell +docker compose -f examples/docker-compose.yml --profile tools down --remove-orphans -v +``` + +Or from this example directory: + +```shell +docker compose -f ../docker-compose.yml --profile tools down --remove-orphans -v +``` \ No newline at end of file diff --git a/examples/transactions/go.mod b/examples/transactions/go.mod new file mode 100644 index 0000000..81a5c4c --- /dev/null +++ b/examples/transactions/go.mod @@ -0,0 +1,7 @@ +module transactions + +go 1.26 + +require ( + github.com/mkbeh/xpg v0.2.0 +) \ No newline at end of file diff --git a/examples/transactions/main.go b/examples/transactions/main.go new file mode 100644 index 0000000..5cc5c1a --- /dev/null +++ b/examples/transactions/main.go @@ -0,0 +1,193 @@ +package main + +import ( + "context" + _ "embed" + "errors" + "fmt" + "log" + "os" + + "github.com/jackc/pgx/v5" + "github.com/jackc/pgx/v5/pgconn" + "github.com/mkbeh/xpg" +) + +const ( + defaultDatabaseURL = "postgres://postgres:postgres@localhost:5432/postgres?sslmode=disable" + uniqueViolationCode = "23505" +) + +//go:embed setup.sql +var setupSQL string + +func main() { + if err := run(context.Background()); err != nil { + log.Fatal(err) + } +} + +func run(ctx context.Context) error { + pool, err := xpg.Open( + ctx, + databaseURL(), + xpg.WithName("transactions-example"), + ) + if err != nil { + return fmt.Errorf("create pool: %w", err) + } + defer pool.Close() + + if err := pool.Ping(ctx); err != nil { + return fmt.Errorf("ping PostgreSQL: %w", err) + } + + if err := prepareExample(ctx, pool); err != nil { + return fmt.Errorf("prepare example: %w", err) + } + + const ( + orderID = int64(1) + promoCode = "PROMO2026" + ) + + if err := processOrder(ctx, pool, orderID, promoCode); err != nil { + return fmt.Errorf("process order: %w", err) + } + + status, promoApplied, err := loadOrderResult( + ctx, + pool, + orderID, + promoCode, + ) + if err != nil { + return fmt.Errorf("load order result: %w", err) + } + + fmt.Printf("order ID: %d\n", orderID) + fmt.Printf("order status: %s\n", status) + fmt.Printf("promo code: %s\n", promoCode) + fmt.Printf("promo applied: %t\n", promoApplied) + + return nil +} + +func prepareExample( + ctx context.Context, + pool *xpg.Pool, +) error { + _, err := pool.Exec( + ctx, + setupSQL, + pgx.QueryExecModeSimpleProtocol, + ) + if err != nil { + return fmt.Errorf("execute setup SQL: %w", err) + } + + return nil +} + +func processOrder( + ctx context.Context, + pool *xpg.Pool, + orderID int64, + promoCode string, +) error { + return pool.InTx( + ctx, + pgx.TxOptions{}, + func(ctx context.Context, tx pgx.Tx) error { + // The order must be committed even when the optional promo fails. + if _, err := tx.Exec( + ctx, + `INSERT INTO xpg_transactions_example.orders (id, status) + VALUES ($1, $2)`, + orderID, + "new", + ); err != nil { + return fmt.Errorf("create order: %w", err) + } + + err := xpg.InSavepoint( + ctx, + tx, + func(ctx context.Context, savepoint pgx.Tx) error { + _, err := savepoint.Exec( + ctx, + ` + INSERT INTO xpg_transactions_example.promo_redemptions ( + code, + order_id + ) + VALUES ( + $1, + $2 + ) + `, + promoCode, + orderID, + ) + + return err + }, + ) + if err == nil { + return nil + } + + if pgErr, ok := errors.AsType[*pgconn.PgError](err); ok && + pgErr.Code == uniqueViolationCode { + // InSavepoint has already rolled back the failed promo insert. + return nil + } + + return fmt.Errorf("apply promo: %w", err) + }, + ) +} + +func loadOrderResult( + ctx context.Context, + pool *xpg.Pool, + orderID int64, + promoCode string, +) (string, bool, error) { + var ( + status string + promoApplied bool + ) + + err := pool.QueryRow( + ctx, + `SELECT + o.status, + EXISTS ( + SELECT 1 + FROM xpg_transactions_example.promo_redemptions AS p + WHERE p.order_id = o.id + AND p.code = $2 + ) + FROM xpg_transactions_example.orders AS o + WHERE o.id = $1`, + orderID, + promoCode, + ).Scan( + &status, + &promoApplied, + ) + if err != nil { + return "", false, err + } + + return status, promoApplied, nil +} + +func databaseURL() string { + if value := os.Getenv("XPG_DATABASE_URL"); value != "" { + return value + } + + return defaultDatabaseURL +} diff --git a/examples/transactions/setup.sql b/examples/transactions/setup.sql new file mode 100644 index 0000000..a2638ec --- /dev/null +++ b/examples/transactions/setup.sql @@ -0,0 +1,24 @@ +BEGIN; + +DROP SCHEMA IF EXISTS xpg_transactions_example CASCADE; + +CREATE SCHEMA xpg_transactions_example; + +CREATE TABLE xpg_transactions_example.orders +( + id bigint PRIMARY KEY, + status text NOT NULL +); + +CREATE TABLE xpg_transactions_example.promo_redemptions +( + code text PRIMARY KEY, + order_id bigint NOT NULL +); + +INSERT INTO xpg_transactions_example.promo_redemptions (code, + order_id) +VALUES ('PROMO2026', + 100); + +COMMIT; \ No newline at end of file diff --git a/tx.go b/tx.go new file mode 100644 index 0000000..178629d --- /dev/null +++ b/tx.go @@ -0,0 +1,77 @@ +package xpg + +import ( + "context" + "errors" + "fmt" + + "github.com/jackc/pgx/v5" +) + +// InTx executes fn in a transaction configured by txOptions. +// +// The transaction is committed when fn returns nil and rolled back when fn +// returns an error. If fn panics, rollback is attempted before the panic is +// propagated. The callback must not call Commit or Rollback; InTx owns +// transaction finalization. +// +// The callback receives the same context and an explicit pgx.Tx. Context +// cancellation does not automatically finalize the transaction while fn is +// running; fn should observe ctx and return promptly. +func (p *Pool) InTx( + ctx context.Context, + txOptions pgx.TxOptions, + fn func(context.Context, pgx.Tx) error, +) error { + if fn == nil { + return errors.New("xpg: transaction function is nil") + } + + err := pgx.BeginTxFunc( + ctx, + p.pool, + txOptions, + func(tx pgx.Tx) error { + return fn(ctx, tx) + }, + ) + if err != nil { + return fmt.Errorf("xpg: transaction: %w", err) + } + + return nil +} + +// InSavepoint executes fn in a pseudo-nested transaction implemented with a +// PostgreSQL savepoint. +// +// The savepoint is released when fn returns nil and rolled back when fn returns +// an error. If fn panics, rollback is attempted before the panic is propagated. +// The callback must not call Commit or Rollback; InSavepoint owns savepoint +// finalization. +func InSavepoint( + ctx context.Context, + tx pgx.Tx, + fn func(context.Context, pgx.Tx) error, +) error { + if tx == nil { + return errors.New("xpg: transaction is nil") + } + + if fn == nil { + return errors.New("xpg: savepoint function is nil") + } + + err := pgx.BeginFunc( + ctx, + tx, + func(savepoint pgx.Tx) error { + return fn(ctx, savepoint) + }, + ) + if err != nil { + return fmt.Errorf("xpg: savepoint: %w", err) + } + + return nil +} diff --git a/tx_test.go b/tx_test.go new file mode 100644 index 0000000..0565735 --- /dev/null +++ b/tx_test.go @@ -0,0 +1,86 @@ +package xpg + +import ( + "context" + "testing" + + "github.com/jackc/pgx/v5" +) + +type testTx struct { + pgx.Tx +} + +func TestPoolInTxNilFunc(t *testing.T) { + t.Parallel() + + err := (&Pool{}).InTx( + context.Background(), + pgx.TxOptions{}, + nil, + ) + + assertErrorMessage( + t, + err, + "xpg: transaction function is nil", + ) +} + +func TestInSavepointNilTx(t *testing.T) { + t.Parallel() + + called := false + + err := InSavepoint( + context.Background(), + nil, + func(context.Context, pgx.Tx) error { + called = true + + return nil + }, + ) + + assertErrorMessage( + t, + err, + "xpg: transaction is nil", + ) + + if called { + t.Fatal("savepoint function was called with a nil transaction") + } +} + +func TestInSavepointNilFunc(t *testing.T) { + t.Parallel() + + err := InSavepoint( + context.Background(), + &testTx{}, + nil, + ) + + assertErrorMessage( + t, + err, + "xpg: savepoint function is nil", + ) +} + +func assertErrorMessage( + t *testing.T, + err error, + want string, +) { + t.Helper() + + if err == nil { + t.Fatalf("expected error %q, got nil", want) + } + + if err.Error() != want { + t.Fatalf("unexpected error: got %q, want %q", err, want) + } +} From f86122cf1f8eb2393aa290bdf1fcb4fd2b0d7e0d Mon Sep 17 00:00:00 2001 From: mkbeh Date: Mon, 3 Aug 2026 17:40:40 +0300 Subject: [PATCH 03/41] docs: add basic pool usage example --- examples/basic/README.md | 91 +++++++++++++++++ examples/basic/go.mod | 8 ++ examples/basic/main.go | 211 +++++++++++++++++++++++++++++++++++++++ examples/basic/setup.sql | 14 +++ 4 files changed, 324 insertions(+) create mode 100644 examples/basic/README.md create mode 100644 examples/basic/go.mod create mode 100644 examples/basic/main.go create mode 100644 examples/basic/setup.sql diff --git a/examples/basic/README.md b/examples/basic/README.md new file mode 100644 index 0000000..a1f2a21 --- /dev/null +++ b/examples/basic/README.md @@ -0,0 +1,91 @@ +# Basic pool usage + +This example opens an `xpg.Pool`, checks PostgreSQL connectivity, writes two rows, and reads them back with the common +query methods exposed by the pool. + +**This example demonstrates:** + +* Creating and closing an `xpg.Pool` +* Explicitly checking connectivity with `Pool.Ping` +* Executing a statement with `Pool.Exec` +* Reading one row with `Pool.QueryRow` +* Iterating over rows returned by `Pool.Query` +* Using `Pool.Name` as logical pool metadata + +## Configuration + +The example uses the following connection string by default: + +```text +postgres://postgres:postgres@localhost:5432/postgres?sslmode=disable +``` + +Set `XPG_DATABASE_URL` to use another PostgreSQL instance: + +```shell +export XPG_DATABASE_URL='postgres://user:password@localhost:5432/database?sslmode=disable' +``` + +## Local PostgreSQL setup + +From the repository root, start PostgreSQL and Adminer: + +```shell +docker compose -f examples/docker-compose.yml --profile tools up -d +``` + +Or from this directory: + +```shell +docker compose -f ../docker-compose.yml --profile tools up -d +``` + +Adminer is available at . Sign in with: + +```text +System: PostgreSQL +Server: postgres +Username: postgres +Password: postgres +Database: postgres +``` + +> [!IMPORTANT] +> Use `postgres`, not `localhost`, in the **Server** field. Adminer connects through the Docker Compose network, where +> PostgreSQL is available by its service name. + +## Run + +From this directory: + +```shell +go run . +``` + +Or from the repository root: + +```shell +go run ./examples/basic +``` + +## Expected output + +```text +pool: basic-example +inserted users: 2 +selected user: 1 Alice active=true +active users: +- 1 Alice +- 2 Bob +``` + +The embedded `setup.sql` file recreates the `xpg_basic_example` schema before each run. The resulting data remains in +PostgreSQL so it can be inspected in Adminer. + +## Stop services + +```shell +docker compose -f examples/docker-compose.yml down +``` + +Add `-v` to remove the PostgreSQL volume as well. diff --git a/examples/basic/go.mod b/examples/basic/go.mod new file mode 100644 index 0000000..c1f1325 --- /dev/null +++ b/examples/basic/go.mod @@ -0,0 +1,8 @@ +module basic + +go 1.26 + +require ( + github.com/jackc/pgx/v5 v5.10.0 + github.com/mkbeh/xpg v0.2.0 +) diff --git a/examples/basic/main.go b/examples/basic/main.go new file mode 100644 index 0000000..fb7005f --- /dev/null +++ b/examples/basic/main.go @@ -0,0 +1,211 @@ +package main + +import ( + "context" + _ "embed" + "fmt" + "log" + "os" + + "github.com/jackc/pgx/v5" + "github.com/mkbeh/xpg" +) + +const defaultDatabaseURL = "postgres://postgres:postgres@localhost:5432/postgres?sslmode=disable" + +//go:embed setup.sql +var setupSQL string + +type user struct { + ID int64 + Name string + Email string + Active bool +} + +func main() { + if err := run(context.Background()); err != nil { + log.Fatal(err) + } +} + +func run(ctx context.Context) error { + pool, err := xpg.Open( + ctx, + databaseURL(), + xpg.WithName("basic-example"), + ) + if err != nil { + return fmt.Errorf("create pool: %w", err) + } + defer pool.Close() + + if err := pool.Ping(ctx); err != nil { + return fmt.Errorf("ping PostgreSQL: %w", err) + } + + if err := prepareExample(ctx, pool); err != nil { + return fmt.Errorf("prepare example: %w", err) + } + + inserted, err := insertUsers(ctx, pool) + if err != nil { + return fmt.Errorf("insert users: %w", err) + } + + selected, err := loadUser(ctx, pool, 1) + if err != nil { + return fmt.Errorf("load user: %w", err) + } + + users, err := listUsers(ctx, pool) + if err != nil { + return fmt.Errorf("list users: %w", err) + } + + fmt.Printf("pool: %s\n", pool.Name()) + fmt.Printf("inserted users: %d\n", inserted) + fmt.Printf( + "selected user: %d %s <%s> active=%t\n", + selected.ID, + selected.Name, + selected.Email, + selected.Active, + ) + fmt.Println("active users:") + + for _, current := range users { + fmt.Printf("- %d %s <%s>\n", current.ID, current.Name, current.Email) + } + + return nil +} + +func prepareExample( + ctx context.Context, + pool *xpg.Pool, +) error { + _, err := pool.Exec( + ctx, + setupSQL, + pgx.QueryExecModeSimpleProtocol, + ) + if err != nil { + return fmt.Errorf("execute setup SQL: %w", err) + } + + return nil +} + +func insertUsers( + ctx context.Context, + pool *xpg.Pool, +) (int64, error) { + tag, err := pool.Exec( + ctx, + `INSERT INTO xpg_basic_example.users ( + id, + name, + email, + active + ) + VALUES + ($1, $2, $3, $4), + ($5, $6, $7, $8)`, + int64(1), + "Alice", + "alice@example.com", + true, + int64(2), + "Bob", + "bob@example.com", + true, + ) + if err != nil { + return 0, err + } + + return tag.RowsAffected(), nil +} + +func loadUser( + ctx context.Context, + pool *xpg.Pool, + userID int64, +) (user, error) { + var selected user + + err := pool.QueryRow( + ctx, + `SELECT + id, + name, + email, + active + FROM xpg_basic_example.users + WHERE id = $1`, + userID, + ).Scan( + &selected.ID, + &selected.Name, + &selected.Email, + &selected.Active, + ) + if err != nil { + return user{}, err + } + + return selected, nil +} + +func listUsers( + ctx context.Context, + pool *xpg.Pool, +) ([]user, error) { + rows, err := pool.Query( + ctx, + `SELECT + id, + name, + email, + active + FROM xpg_basic_example.users + WHERE active + ORDER BY id`, + ) + if err != nil { + return nil, err + } + defer rows.Close() + + users := make([]user, 0, 2) + + for rows.Next() { + var current user + + if err := rows.Scan( + ¤t.ID, + ¤t.Name, + ¤t.Email, + ¤t.Active, + ); err != nil { + return nil, err + } + + users = append(users, current) + } + + if err := rows.Err(); err != nil { + return nil, err + } + + return users, nil +} + +func databaseURL() string { + if value := os.Getenv("XPG_DATABASE_URL"); value != "" { + return value + } + + return defaultDatabaseURL +} diff --git a/examples/basic/setup.sql b/examples/basic/setup.sql new file mode 100644 index 0000000..5be37b5 --- /dev/null +++ b/examples/basic/setup.sql @@ -0,0 +1,14 @@ +BEGIN; + +DROP SCHEMA IF EXISTS xpg_basic_example CASCADE; + +CREATE SCHEMA xpg_basic_example; + +CREATE TABLE xpg_basic_example.users ( + id bigint PRIMARY KEY, + name text NOT NULL, + email text NOT NULL UNIQUE, + active boolean NOT NULL +); + +COMMIT; From 4f0b0fce73a7d3f047b52f7cb5cbc9012ffb1455 Mon Sep 17 00:00:00 2001 From: mkbeh Date: Mon, 3 Aug 2026 20:05:44 +0300 Subject: [PATCH 04/41] feat: add error classifiers and advisory locks --- advisory.go | 49 ++++++++++++++ advisory_test.go | 120 ++++++++++++++++++++++++++++++++++ errors.go | 118 ++++++++++++++++++++++++++++++++++ errors_test.go | 164 +++++++++++++++++++++++++++++++++++++++++++++++ 4 files changed, 451 insertions(+) create mode 100644 advisory.go create mode 100644 advisory_test.go create mode 100644 errors_test.go diff --git a/advisory.go b/advisory.go new file mode 100644 index 0000000..a06bef4 --- /dev/null +++ b/advisory.go @@ -0,0 +1,49 @@ +package xpg + +import ( + "context" + "errors" + "fmt" + + "github.com/jackc/pgx/v5" +) + +const ( + advisoryXactLockSQL = "SELECT pg_advisory_xact_lock($1)" + tryAdvisoryXactLockSQL = "SELECT pg_try_advisory_xact_lock($1)" +) + +// AdvisoryXactLock acquires an exclusive transaction-level advisory lock. +// +// The call waits until the lock is available or ctx is canceled. PostgreSQL +// releases the lock automatically when tx is committed or rolled back. +func AdvisoryXactLock(ctx context.Context, tx pgx.Tx, key int64) error { + if tx == nil { + return errors.New("xpg: transaction is nil") + } + + if _, err := tx.Exec(ctx, advisoryXactLockSQL, key); err != nil { + return fmt.Errorf("xpg: acquire transaction advisory lock: %w", err) + } + + return nil +} + +// TryAdvisoryXactLock attempts to acquire an exclusive transaction-level +// advisory lock without waiting. +// +// PostgreSQL releases an acquired lock automatically when tx is committed or +// rolled back. +func TryAdvisoryXactLock(ctx context.Context, tx pgx.Tx, key int64) (bool, error) { + if tx == nil { + return false, errors.New("xpg: transaction is nil") + } + + var acquired bool + + if err := tx.QueryRow(ctx, tryAdvisoryXactLockSQL, key).Scan(&acquired); err != nil { + return false, fmt.Errorf("xpg: try transaction advisory lock: %w", err) + } + + return acquired, nil +} diff --git a/advisory_test.go b/advisory_test.go new file mode 100644 index 0000000..df8eb68 --- /dev/null +++ b/advisory_test.go @@ -0,0 +1,120 @@ +package xpg + +import ( + "context" + "errors" + "testing" + + "github.com/jackc/pgx/v5" + "github.com/jackc/pgx/v5/pgconn" +) + +type advisoryTestTx struct { + pgx.Tx + + execErr error + row pgx.Row +} + +func (tx *advisoryTestTx) Exec(context.Context, string, ...any) (pgconn.CommandTag, error) { + return pgconn.CommandTag{}, tx.execErr +} + +func (tx *advisoryTestTx) QueryRow(context.Context, string, ...any) pgx.Row { + return tx.row +} + +type advisoryTestRow struct { + acquired bool + err error +} + +func (row advisoryTestRow) Scan(dest ...any) error { + if row.err != nil { + return row.err + } + + acquired, ok := dest[0].(*bool) + if !ok { + return errors.New("unexpected destination type") + } + + *acquired = row.acquired + + return nil +} + +func TestAdvisoryXactLockNilTx(t *testing.T) { + t.Parallel() + + err := AdvisoryXactLock( + context.Background(), + nil, + 1, + ) + + assertErrorMessage(t, err, "xpg: transaction is nil") +} + +func TestAdvisoryXactLockPreservesError(t *testing.T) { + t.Parallel() + + expectedErr := errors.New("lock failed") + + err := AdvisoryXactLock( + context.Background(), + &advisoryTestTx{execErr: expectedErr}, + 1, + ) + if !errors.Is(err, expectedErr) { + t.Fatalf("original error was not preserved: %v", err) + } +} + +func TestTryAdvisoryXactLockNilTx(t *testing.T) { + t.Parallel() + + _, err := TryAdvisoryXactLock( + context.Background(), + nil, + 1, + ) + + assertErrorMessage(t, err, "xpg: transaction is nil") +} + +func TestTryAdvisoryXactLock(t *testing.T) { + t.Parallel() + + acquired, err := TryAdvisoryXactLock( + context.Background(), + &advisoryTestTx{ + row: advisoryTestRow{acquired: true}, + }, + 1, + ) + if err != nil { + t.Fatalf("TryAdvisoryXactLock returned an error: %v", err) + } + + if !acquired { + t.Fatal("TryAdvisoryXactLock returned false") + } +} + +func TestTryAdvisoryXactLockPreservesError(t *testing.T) { + t.Parallel() + + expectedErr := errors.New("try lock failed") + + _, err := TryAdvisoryXactLock( + context.Background(), + &advisoryTestTx{ + row: advisoryTestRow{err: expectedErr}, + }, + 1, + ) + if !errors.Is(err, expectedErr) { + t.Fatalf("original error was not preserved: %v", err) + } +} diff --git a/errors.go b/errors.go index d5c081c..ef374d7 100644 --- a/errors.go +++ b/errors.go @@ -1 +1,119 @@ package xpg + +import ( + "errors" + "io" + "net" + + "github.com/jackc/pgx/v5" + "github.com/jackc/pgx/v5/pgconn" +) + +const ( + sqlStateUniqueViolation = "23505" + sqlStateForeignKeyViolation = "23503" + sqlStateNotNullViolation = "23502" + sqlStateCheckViolation = "23514" + sqlStateSerializationFailure = "40001" + sqlStateDeadlockDetected = "40P01" + sqlStateLockNotAvailable = "55P03" + sqlStateQueryCanceled = "57014" +) + +// SQLState returns the PostgreSQL SQLSTATE code carried by err. +// It returns an empty string when the error tree does not contain a +// pgconn.PgError. +func SQLState(err error) string { + pgErr, ok := errors.AsType[*pgconn.PgError](err) + if !ok || pgErr == nil { + return "" + } + + return pgErr.Code +} + +// IsNoRows reports whether err indicates that a query returned no rows. +func IsNoRows(err error) bool { + return errors.Is(err, pgx.ErrNoRows) +} + +// IsUniqueViolation reports whether err is a PostgreSQL unique_violation. +func IsUniqueViolation(err error) bool { + return SQLState(err) == sqlStateUniqueViolation +} + +// IsForeignKeyViolation reports whether err is a PostgreSQL +// foreign_key_violation. +func IsForeignKeyViolation(err error) bool { + return SQLState(err) == sqlStateForeignKeyViolation +} + +// IsNotNullViolation reports whether err is a PostgreSQL not_null_violation. +func IsNotNullViolation(err error) bool { + return SQLState(err) == sqlStateNotNullViolation +} + +// IsCheckViolation reports whether err is a PostgreSQL check_violation. +func IsCheckViolation(err error) bool { + return SQLState(err) == sqlStateCheckViolation +} + +// IsSerializationFailure reports whether err is a PostgreSQL +// serialization_failure. +func IsSerializationFailure(err error) bool { + return SQLState(err) == sqlStateSerializationFailure +} + +// IsDeadlock reports whether err is a PostgreSQL deadlock_detected error. +func IsDeadlock(err error) bool { + return SQLState(err) == sqlStateDeadlockDetected +} + +// IsLockNotAvailable reports whether err is a PostgreSQL lock_not_available +// error. +func IsLockNotAvailable(err error) bool { + return SQLState(err) == sqlStateLockNotAvailable +} + +// IsQueryCanceled reports whether PostgreSQL canceled the query. +// +// Client-side context cancellation remains available through errors.Is with +// context.Canceled or context.DeadlineExceeded. +func IsQueryCanceled(err error) bool { + return SQLState(err) == sqlStateQueryCanceled +} + +// IsConnectionError reports whether err represents a PostgreSQL connection +// failure known to pgx or the Go networking stack. +func IsConnectionError(err error) bool { + if err == nil { + return false + } + + if errors.Is(err, pgconn.ErrConnClosed) || + errors.Is(err, io.EOF) || + errors.Is(err, io.ErrUnexpectedEOF) { + return true + } + + if _, ok := errors.AsType[*pgconn.ConnectError](err); ok { + return true + } + + if _, ok := errors.AsType[*net.OpError](err); ok { + return true + } + + state := SQLState(err) + + return len(state) >= 2 && state[:2] == "08" +} + +// IsRetryableTransaction reports whether PostgreSQL aborted the transaction +// because of a serialization failure or a deadlock. +// +// The entire transaction callback must still be safe to replay. Connection +// failures are deliberately not classified as transaction-retryable. +func IsRetryableTransaction(err error) bool { + return IsSerializationFailure(err) || IsDeadlock(err) +} diff --git a/errors_test.go b/errors_test.go new file mode 100644 index 0000000..6f74f34 --- /dev/null +++ b/errors_test.go @@ -0,0 +1,164 @@ +package xpg + +import ( + "errors" + "fmt" + "net" + "testing" + + "github.com/jackc/pgx/v5" + "github.com/jackc/pgx/v5/pgconn" +) + +func TestSQLState(t *testing.T) { + t.Parallel() + + err := fmt.Errorf( + "wrapped: %w", + &pgconn.PgError{Code: sqlStateUniqueViolation}, + ) + + if state := SQLState(err); state != sqlStateUniqueViolation { + t.Fatalf( + "unexpected SQLSTATE: got %q, want %q", + state, + sqlStateUniqueViolation, + ) + } +} + +func TestIsNoRows(t *testing.T) { + t.Parallel() + + if !IsNoRows(fmt.Errorf("wrapped: %w", pgx.ErrNoRows)) { + t.Fatal("IsNoRows returned false") + } +} + +func TestErrorClassifiers(t *testing.T) { + t.Parallel() + + tests := []struct { + name string + state string + classifier func(error) bool + }{ + { + name: "unique violation", + state: sqlStateUniqueViolation, + classifier: IsUniqueViolation, + }, + { + name: "foreign key violation", + state: sqlStateForeignKeyViolation, + classifier: IsForeignKeyViolation, + }, + { + name: "not null violation", + state: sqlStateNotNullViolation, + classifier: IsNotNullViolation, + }, + { + name: "check violation", + state: sqlStateCheckViolation, + classifier: IsCheckViolation, + }, + { + name: "serialization failure", + state: sqlStateSerializationFailure, + classifier: IsSerializationFailure, + }, + { + name: "deadlock", + state: sqlStateDeadlockDetected, + classifier: IsDeadlock, + }, + { + name: "lock not available", + state: sqlStateLockNotAvailable, + classifier: IsLockNotAvailable, + }, + { + name: "query canceled", + state: sqlStateQueryCanceled, + classifier: IsQueryCanceled, + }, + } + + for _, test := range tests { + t.Run(test.name, func(t *testing.T) { + t.Parallel() + + err := fmt.Errorf( + "wrapped: %w", + &pgconn.PgError{Code: test.state}, + ) + + if !test.classifier(err) { + t.Fatalf( + "classifier returned false for SQLSTATE %q", + test.state, + ) + } + }) + } +} + +func TestIsConnectionError(t *testing.T) { + t.Parallel() + + tests := []struct { + name string + err error + }{ + { + name: "SQLSTATE connection exception", + err: &pgconn.PgError{Code: "08006"}, + }, + { + name: "closed connection", + err: fmt.Errorf("wrapped: %w", pgconn.ErrConnClosed), + }, + { + name: "network operation", + err: &net.OpError{ + Op: "read", + Net: "tcp", + Err: errors.New("connection reset"), + }, + }, + } + + for _, test := range tests { + t.Run(test.name, func(t *testing.T) { + t.Parallel() + + if !IsConnectionError(test.err) { + t.Fatal("IsConnectionError returned false") + } + }) + } +} + +func TestIsRetryableTransaction(t *testing.T) { + t.Parallel() + + for _, state := range []string{ + sqlStateSerializationFailure, + sqlStateDeadlockDetected, + } { + err := &pgconn.PgError{Code: state} + if !IsRetryableTransaction(err) { + t.Fatalf( + "IsRetryableTransaction returned false for SQLSTATE %q", + state, + ) + } + } + + if IsRetryableTransaction( + &pgconn.PgError{Code: sqlStateUniqueViolation}, + ) { + t.Fatal("IsRetryableTransaction returned true for unique violation") + } +} From 4fbddaa1e63310c6d1db309d9c59a468bcb2a83e Mon Sep 17 00:00:00 2001 From: mkbeh Date: Mon, 3 Aug 2026 20:05:56 +0300 Subject: [PATCH 05/41] docs: add advisory lock example --- examples/README.md | 1 + examples/advisory/README.md | 125 ++++++++++++++++ examples/advisory/go.mod | 8 ++ examples/advisory/main.go | 280 ++++++++++++++++++++++++++++++++++++ examples/advisory/setup.sql | 14 ++ 5 files changed, 428 insertions(+) create mode 100644 examples/advisory/README.md create mode 100644 examples/advisory/go.mod create mode 100644 examples/advisory/main.go create mode 100644 examples/advisory/setup.sql diff --git a/examples/README.md b/examples/README.md index 6fe2f93..e44c38d 100644 --- a/examples/README.md +++ b/examples/README.md @@ -6,6 +6,7 @@ This directory contains runnable examples demonstrating the main features and us |:-------------------------------|:------------------------------------------------------------------------------------------| | [`basic`](basic) | Pool lifecycle and common query methods | | [`transactions`](transactions) | Committing an outer transaction after an optional operation is rolled back to a savepoint | +| [`advisory`](advisory) | Coordinating concurrent transactions with PostgreSQL advisory locks | ## Running the examples diff --git a/examples/advisory/README.md b/examples/advisory/README.md new file mode 100644 index 0000000..dfea9b2 --- /dev/null +++ b/examples/advisory/README.md @@ -0,0 +1,125 @@ +# Transaction advisory locks + +This example coordinates concurrent workers with PostgreSQL transaction-level advisory locks. The first worker acquires +an advisory lock and holds it until its transaction commits. A second worker uses the non-blocking try variant and +cannot +enter the protected section while the lock is held. After the first transaction commits, a third worker acquires the +same +lock successfully. + +**This example demonstrates:** + +* Acquiring a transaction-level lock with `xpg.AdvisoryXactLock` +* Trying to acquire a lock without waiting with `xpg.TryAdvisoryXactLock` +* Holding a lock for the lifetime of an explicit `pgx.Tx` +* Releasing a transaction-level lock automatically on commit or rollback +* Coordinating concurrent database work without `pg_sleep` + +## Configuration + +The example uses the following connection string by default: + +```text +postgres://postgres:postgres@localhost:5432/postgres?sslmode=disable +``` + +Set `XPG_DATABASE_URL` to use another PostgreSQL instance: + +```shell +export XPG_DATABASE_URL='postgres://user:password@localhost:5432/database?sslmode=disable' +``` + +> [!NOTE] +> The example runs two transactions concurrently, so the pool must allow at least two connections. The default pgxpool +> configuration satisfies this requirement. + +## Local PostgreSQL setup + +From the repository root, start PostgreSQL and Adminer: + +```shell +docker compose -f examples/docker-compose.yml --profile tools up -d +``` + +Or from this directory: + +```shell +docker compose -f ../docker-compose.yml --profile tools up -d +``` + +Adminer is available at: + +```text +http://localhost:8080 +``` + +Sign in to Adminer with: + +```text +System: PostgreSQL +Server: postgres +Username: postgres +Password: postgres +Database: postgres +``` + +> [!IMPORTANT] +> Use `postgres`, not `localhost`, in the **Server** field. Adminer connects through the Docker Compose network, where +> PostgreSQL is available by its service name. + +## Run + +From this directory: + +```shell +go run . +``` + +Or from the repository root: + +```shell +go run ./examples/advisory +``` + +## Flow + +This example is easier to follow as a sequence: + +```text +1. Reset the example state +2. Worker A acquires the advisory lock +3. Worker B tries the same lock without waiting +4. Worker A commits and releases the lock +5. Worker C acquires the released lock +6. Read the committed job runs +``` + +## Expected output + +```text +worker-a acquired the lock +worker-b acquired the lock: false +worker-a committed and released the lock +worker-c acquired the lock: true +recorded job runs: +- worker-a (lock key: 2026) +- worker-c (lock key: 2026) +``` + +> [!IMPORTANT] +> Advisory lock keys are application-defined `int64` values. Use a stable key mapping and keep the protected transaction +> short because it holds both the lock and a pool connection until commit or rollback. + +## Stop services + +From the repository root: + +```shell +docker compose -f examples/docker-compose.yml --profile tools down --remove-orphans -v +``` + +Or from this directory: + +```shell +docker compose -f ../docker-compose.yml --profile tools down --remove-orphans -v +``` \ No newline at end of file diff --git a/examples/advisory/go.mod b/examples/advisory/go.mod new file mode 100644 index 0000000..1907213 --- /dev/null +++ b/examples/advisory/go.mod @@ -0,0 +1,8 @@ +module advisory + +go 1.26 + +require ( + github.com/jackc/pgx/v5 v5.10.0 + github.com/mkbeh/xpg v0.2.0 +) \ No newline at end of file diff --git a/examples/advisory/main.go b/examples/advisory/main.go new file mode 100644 index 0000000..8df67db --- /dev/null +++ b/examples/advisory/main.go @@ -0,0 +1,280 @@ +package main + +import ( + "context" + _ "embed" + "errors" + "fmt" + "log" + "os" + + "github.com/jackc/pgx/v5" + "github.com/mkbeh/xpg" +) + +const ( + defaultDatabaseURL = "postgres://postgres:postgres@localhost:5432/postgres?sslmode=disable" + jobLockKey = int64(2026) +) + +//go:embed setup.sql +var setupSQL string + +type jobRun struct { + Worker string + LockKey int64 +} + +func main() { + if err := run(context.Background()); err != nil { + log.Fatal(err) + } +} + +func run(ctx context.Context) error { + pool, err := xpg.Open( + ctx, + databaseURL(), + xpg.WithName("advisory-example"), + ) + if err != nil { + return fmt.Errorf("create pool: %w", err) + } + defer pool.Close() + + if err := pool.Ping(ctx); err != nil { + return fmt.Errorf("ping PostgreSQL: %w", err) + } + + if err := prepareExample(ctx, pool); err != nil { + return fmt.Errorf("prepare example: %w", err) + } + + lockAcquired := make(chan struct{}) + releaseLock := make(chan struct{}) + holderDone := make(chan error, 1) + + go func() { + holderDone <- holdJobLock( + ctx, + pool, + "worker-a", + jobLockKey, + lockAcquired, + releaseLock, + ) + }() + + select { + case <-lockAcquired: + fmt.Println("worker-a acquired the lock") + case err := <-holderDone: + if err == nil { + return errors.New("worker-a exited before acquiring the lock") + } + + return fmt.Errorf("worker-a: %w", err) + case <-ctx.Done(): + return ctx.Err() + } + + workerBAcquired, workerBErr := tryRunJob( + ctx, + pool, + "worker-b", + jobLockKey, + ) + fmt.Printf("worker-b acquired the lock: %t\n", workerBAcquired) + + close(releaseLock) + holderErr := <-holderDone + + if workerBErr != nil { + return fmt.Errorf("worker-b: %w", workerBErr) + } + + if holderErr != nil { + return fmt.Errorf("worker-a: %w", holderErr) + } + + if workerBAcquired { + return errors.New("worker-b acquired a lock that should still be held") + } + + fmt.Println("worker-a committed and released the lock") + + workerCAcquired, err := tryRunJob( + ctx, + pool, + "worker-c", + jobLockKey, + ) + if err != nil { + return fmt.Errorf("worker-c: %w", err) + } + + if !workerCAcquired { + return errors.New("worker-c did not acquire the released lock") + } + + fmt.Printf("worker-c acquired the lock: %t\n", workerCAcquired) + + runs, err := loadJobRuns(ctx, pool) + if err != nil { + return fmt.Errorf("load job runs: %w", err) + } + + fmt.Println("recorded job runs:") + for _, current := range runs { + fmt.Printf("- %s (lock key: %d)\n", current.Worker, current.LockKey) + } + + return nil +} + +func prepareExample( + ctx context.Context, + pool *xpg.Pool, +) error { + _, err := pool.Exec( + ctx, + setupSQL, + pgx.QueryExecModeSimpleProtocol, + ) + if err != nil { + return fmt.Errorf("execute setup SQL: %w", err) + } + + return nil +} + +func holdJobLock( + ctx context.Context, + pool *xpg.Pool, + worker string, + lockKey int64, + lockAcquired chan<- struct{}, + releaseLock <-chan struct{}, +) error { + return pool.InTx( + ctx, + pgx.TxOptions{}, + func(ctx context.Context, tx pgx.Tx) error { + if err := xpg.AdvisoryXactLock(ctx, tx, lockKey); err != nil { + return err + } + + if err := recordJobRun(ctx, tx, worker, lockKey); err != nil { + return fmt.Errorf("record job run: %w", err) + } + + close(lockAcquired) + + select { + case <-releaseLock: + return nil + case <-ctx.Done(): + return ctx.Err() + } + }, + ) +} + +func tryRunJob( + ctx context.Context, + pool *xpg.Pool, + worker string, + lockKey int64, +) (bool, error) { + var acquired bool + + err := pool.InTx( + ctx, + pgx.TxOptions{}, + func(ctx context.Context, tx pgx.Tx) error { + var err error + + acquired, err = xpg.TryAdvisoryXactLock(ctx, tx, lockKey) + if err != nil { + return err + } + + if !acquired { + return nil + } + + if err := recordJobRun(ctx, tx, worker, lockKey); err != nil { + return fmt.Errorf("record job run: %w", err) + } + + return nil + }, + ) + if err != nil { + return false, err + } + + return acquired, nil +} + +func recordJobRun( + ctx context.Context, + tx pgx.Tx, + worker string, + lockKey int64, +) error { + _, err := tx.Exec( + ctx, + `INSERT INTO xpg_advisory_example.job_runs (worker, lock_key) + VALUES ($1, $2)`, + worker, + lockKey, + ) + + return err +} + +func loadJobRuns( + ctx context.Context, + pool *xpg.Pool, +) ([]jobRun, error) { + rows, err := pool.Query( + ctx, + `SELECT worker, lock_key + FROM xpg_advisory_example.job_runs + ORDER BY id`, + ) + if err != nil { + return nil, err + } + defer rows.Close() + + runs := make([]jobRun, 0, 2) + + for rows.Next() { + var current jobRun + + if err := rows.Scan( + ¤t.Worker, + ¤t.LockKey, + ); err != nil { + return nil, err + } + + runs = append(runs, current) + } + + if err := rows.Err(); err != nil { + return nil, err + } + + return runs, nil +} + +func databaseURL() string { + if value := os.Getenv("XPG_DATABASE_URL"); value != "" { + return value + } + + return defaultDatabaseURL +} diff --git a/examples/advisory/setup.sql b/examples/advisory/setup.sql new file mode 100644 index 0000000..bab9a39 --- /dev/null +++ b/examples/advisory/setup.sql @@ -0,0 +1,14 @@ +BEGIN; + +DROP SCHEMA IF EXISTS xpg_advisory_example CASCADE; + +CREATE SCHEMA xpg_advisory_example; + +CREATE TABLE xpg_advisory_example.job_runs ( + id bigint GENERATED ALWAYS AS IDENTITY PRIMARY KEY, + worker text NOT NULL, + lock_key bigint NOT NULL, + created_at timestamptz NOT NULL DEFAULT clock_timestamp() +); + +COMMIT; From a8d47e0804e8cffc4cd98679853aac7c37922a49 Mon Sep 17 00:00:00 2001 From: mkbeh Date: Wed, 5 Aug 2026 13:34:12 +0300 Subject: [PATCH 06/41] feat: add pluggable pool metrics --- metrics.go | 15 +++++++++++ options.go | 45 ++++++++++++++++++++++++++----- pool.go | 54 ++++++++++++++++++++++++++++++++----- stats.go | 79 ++++++++++++++++++++++++++++++++++++++++++++++++++++++ 4 files changed, 180 insertions(+), 13 deletions(-) create mode 100644 metrics.go create mode 100644 stats.go diff --git a/metrics.go b/metrics.go new file mode 100644 index 0000000..c100f03 --- /dev/null +++ b/metrics.go @@ -0,0 +1,15 @@ +package xpg + +// PoolMetrics registers metrics for a Pool. +// +// Implementations are expected to be immutable and safe to reuse for multiple +// pools. Register is called after the underlying pgxpool.Pool has been created. +type PoolMetrics interface { + Register(pool *Pool) (PoolMetricsRegistration, error) +} + +// PoolMetricsRegistration owns a metrics registration associated with one +// Pool. Close is called once before the underlying pgxpool.Pool is closed. +type PoolMetricsRegistration interface { + Close() +} diff --git a/options.go b/options.go index 561f696..164522b 100644 --- a/options.go +++ b/options.go @@ -3,6 +3,9 @@ package xpg import ( "errors" "fmt" + "maps" + "net" + "strconv" "strings" ) @@ -20,8 +23,25 @@ func (option optionFunc) apply(settings *settings) error { } type settings struct { - name string - labels map[string]string + name string + labels map[string]string + metrics PoolMetrics +} + +func (s settings) poolName(host string, port uint16, database string) string { + if s.name != "" { + return s.name + } + + address := net.JoinHostPort( + host, + strconv.Itoa(int(port)), + ) + if database == "" { + return address + } + + return address + "/" + database } func defaultSettings() *settings { @@ -98,16 +118,29 @@ func WithLabel(key, value string) Option { }) } +// WithMetrics attaches one metrics implementation to the pool. +// +// Metrics are registered during New and unregistered automatically when the +// Pool is closed. +func WithMetrics(metrics PoolMetrics) Option { + return optionFunc(func(settings *settings) error { + if metrics == nil { + return errors.New("pool metrics is nil") + } + + settings.metrics = metrics + + return nil + }) +} + func cloneLabels(labels map[string]string) map[string]string { if len(labels) == 0 { return nil } cloned := make(map[string]string, len(labels)) - - for key, value := range labels { - cloned[key] = value - } + maps.Copy(cloned, labels) return cloned } diff --git a/pool.go b/pool.go index 3ff64a1..3ed4223 100644 --- a/pool.go +++ b/pool.go @@ -13,7 +13,8 @@ import ( // Pool is a concurrency-safe PostgreSQL connection pool backed by pgxpool. type Pool struct { - pool *pgxpool.Pool + pool *pgxpool.Pool + metrics PoolMetricsRegistration name string labels map[string]string @@ -50,16 +51,30 @@ func New(ctx context.Context, config *pgxpool.Config, options ...Option) (*Pool, return nil, err } - pgxPool, err := pgxpool.NewWithConfig(ctx, config.Copy()) + poolConfig := config.Copy() + connConfig := poolConfig.ConnConfig + + pgxPool, err := pgxpool.NewWithConfig(ctx, poolConfig) if err != nil { return nil, fmt.Errorf("xpg: create pool: %w", err) } - return &Pool{ - pool: pgxPool, - name: settings.name, + pool := &Pool{ + pool: pgxPool, + name: settings.poolName( + connConfig.Host, + connConfig.Port, + connConfig.Database, + ), labels: cloneLabels(settings.labels), - }, nil + } + + if err := pool.registerMetrics(settings.metrics); err != nil { + pool.Close() + return nil, fmt.Errorf("xpg: register pool metrics: %w", err) + } + + return pool, nil } // Name returns the logical pool name configured with WithName. @@ -67,6 +82,10 @@ func (p *Pool) Name() string { return p.name } +func (p *Pool) Labels() map[string]string { + return cloneLabels(p.labels) +} + // Raw returns the underlying pgxpool.Pool. func (p *Pool) Raw() *pgxpool.Pool { return p.pool @@ -79,7 +98,13 @@ func (p *Pool) Ping(ctx context.Context) error { // Close closes the underlying pool and waits for acquired connections to be // returned. Close is safe to call multiple times. func (p *Pool) Close() { - p.closeOnce.Do(p.pool.Close) + p.closeOnce.Do(func() { + if p.metrics != nil { + p.metrics.Close() + } + + p.pool.Close() + }) } func (p *Pool) Exec(ctx context.Context, sql string, arguments ...any) (pgconn.CommandTag, error) { @@ -101,3 +126,18 @@ func (p *Pool) SendBatch(ctx context.Context, batch *pgx.Batch) pgx.BatchResults func (p *Pool) CopyFrom(ctx context.Context, tableName pgx.Identifier, columnNames []string, rowSrc pgx.CopyFromSource) (int64, error) { return p.pool.CopyFrom(ctx, tableName, columnNames, rowSrc) } + +func (p *Pool) registerMetrics(metrics PoolMetrics) error { + if metrics == nil { + return nil + } + + registration, err := metrics.Register(p) + if err != nil { + return err + } + + p.metrics = registration + + return nil +} diff --git a/stats.go b/stats.go new file mode 100644 index 0000000..353ff2b --- /dev/null +++ b/stats.go @@ -0,0 +1,79 @@ +package xpg + +import "time" + +// PoolStats is a detached point-in-time snapshot of connection pool statistics. +type PoolStats struct { + // Current state. + + // AcquiredConns is the number of connections currently checked out from the + // pool. + AcquiredConns int32 + + // ConstructingConns is the number of connections currently being created. + ConstructingConns int32 + + // IdleConns is the number of currently idle connections. + IdleConns int32 + + // MaxConns is the maximum number of connections allowed by the pool. + MaxConns int32 + + // TotalConns is the number of acquired, idle, and constructing connections. + TotalConns int32 + + // Acquire lifecycle. + + // AcquireCount is the cumulative number of successful connection acquires. + AcquireCount int64 + + // AcquireDuration is the cumulative duration of successful connection + // acquires. + AcquireDuration time.Duration + + // CanceledAcquireCount is the cumulative number of connection acquires + // canceled by context cancellation. + CanceledAcquireCount int64 + + // EmptyAcquireCount is the cumulative number of successful acquires that + // waited because the pool was empty. + EmptyAcquireCount int64 + + // EmptyAcquireWaitTime is the cumulative time spent waiting on successful + // acquires while the pool was empty. + EmptyAcquireWaitTime time.Duration + + // Connection lifecycle. + + // NewConnsCount is the cumulative number of connections opened by the pool. + NewConnsCount int64 + + // MaxIdleDestroyCount is the cumulative number of connections closed after + // exceeding MaxConnIdleTime. + MaxIdleDestroyCount int64 + + // MaxLifetimeDestroyCount is the cumulative number of connections closed + // after exceeding MaxConnLifetime. + MaxLifetimeDestroyCount int64 +} + +// Stats returns a detached snapshot of the current pool statistics. +func (p *Pool) Stats() PoolStats { + stats := p.pool.Stat() + + return PoolStats{ + AcquiredConns: stats.AcquiredConns(), + ConstructingConns: stats.ConstructingConns(), + IdleConns: stats.IdleConns(), + MaxConns: stats.MaxConns(), + TotalConns: stats.TotalConns(), + AcquireCount: stats.AcquireCount(), + AcquireDuration: stats.AcquireDuration(), + CanceledAcquireCount: stats.CanceledAcquireCount(), + EmptyAcquireCount: stats.EmptyAcquireCount(), + EmptyAcquireWaitTime: stats.EmptyAcquireWaitTime(), + NewConnsCount: stats.NewConnsCount(), + MaxIdleDestroyCount: stats.MaxIdleDestroyCount(), + MaxLifetimeDestroyCount: stats.MaxLifetimeDestroyCount(), + } +} From a4c2849988f3ede47b2cb460f469b48a2a5cc83b Mon Sep 17 00:00:00 2001 From: mkbeh Date: Wed, 5 Aug 2026 13:34:21 +0300 Subject: [PATCH 07/41] feat: add OpenTelemetry pool metrics --- metrics/otel/go.mod | 9 + metrics/otel/metrics.go | 428 +++++++++++++++++++++++++++++++++++ metrics/otel/options.go | 54 +++++ metrics/otel/options_test.go | 57 +++++ 4 files changed, 548 insertions(+) create mode 100644 metrics/otel/go.mod create mode 100644 metrics/otel/metrics.go create mode 100644 metrics/otel/options.go create mode 100644 metrics/otel/options_test.go diff --git a/metrics/otel/go.mod b/metrics/otel/go.mod new file mode 100644 index 0000000..966131f --- /dev/null +++ b/metrics/otel/go.mod @@ -0,0 +1,9 @@ +module github.com/mkbeh/xpg/metrics/otel + +go 1.26 + +require ( + github.com/mkbeh/xpg v0.2.0 + go.opentelemetry.io/otel v1.45.0 + go.opentelemetry.io/otel/metric v1.45.0 +) diff --git a/metrics/otel/metrics.go b/metrics/otel/metrics.go new file mode 100644 index 0000000..706599b --- /dev/null +++ b/metrics/otel/metrics.go @@ -0,0 +1,428 @@ +package xpgotel + +import ( + "context" + "errors" + "fmt" + "slices" + "sync" + + "github.com/mkbeh/xpg" + "go.opentelemetry.io/otel" + "go.opentelemetry.io/otel/attribute" + "go.opentelemetry.io/otel/metric" +) + +const instrumentationName = "github.com/mkbeh/xpg/otel" + +const ( + connectionCountMetricName = "db.client.connection.count" + connectionMaxMetricName = "db.client.connection.max" + connectionConstructingMetricName = "xpg.pool.connection.constructing" + connectionAcquireCountMetricName = "xpg.pool.connection.acquire.count" + connectionAcquireTimeMetricName = "xpg.pool.connection.acquire.time" + connectionAcquireCanceledCountMetricName = "xpg.pool.connection.acquire.canceled.count" + connectionAcquireEmptyCountMetricName = "xpg.pool.connection.acquire.empty.count" + connectionAcquireEmptyWaitTimeMetricName = "xpg.pool.connection.acquire.empty.wait_time" + connectionCreateCountMetricName = "xpg.pool.connection.create.count" + connectionDestroyCountMetricName = "xpg.pool.connection.destroy.count" +) + +const ( + dbSystemNameAttribute = "db.system.name" + poolNameAttribute = "db.client.connection.pool.name" + connectionStateAttribute = "db.client.connection.state" + destroyReasonAttribute = "xpg.pool.connection.destroy.reason" +) + +const ( + dbSystemPostgreSQL = "postgresql" + + connectionStateIdle = "idle" + connectionStateUsed = "used" + + destroyReasonIdleTimeout = "idle_timeout" + destroyReasonLifetime = "lifetime" +) + +type poolMetricInstruments struct { + connectionCount metric.Int64ObservableUpDownCounter + connectionMax metric.Int64ObservableUpDownCounter + constructingConnections metric.Int64ObservableGauge + acquireCount metric.Int64ObservableCounter + acquireTime metric.Float64ObservableCounter + canceledAcquireCount metric.Int64ObservableCounter + emptyAcquireCount metric.Int64ObservableCounter + emptyAcquireWaitTime metric.Float64ObservableCounter + createdConnections metric.Int64ObservableCounter + destroyedConnections metric.Int64ObservableCounter +} + +type poolMetricAttributes struct { + base metric.ObserveOption + idle metric.ObserveOption + used metric.ObserveOption + destroyedIdle metric.ObserveOption + destroyedLifetime metric.ObserveOption +} + +// Metrics exports xpg pool statistics through OpenTelemetry. +type Metrics struct { + meterProvider metric.MeterProvider +} + +type poolMetrics struct { + registration metric.Registration + closeOnce sync.Once +} + +// Register registers metrics for one xpg Pool. +func (m *Metrics) Register( + pool *xpg.Pool, +) (xpg.PoolMetricsRegistration, error) { + if m == nil { + return nil, errors.New("xpg/otel: metrics is nil") + } + + provider := m.meterProvider + if provider == nil { + provider = otel.GetMeterProvider() + } + + return registerPoolMetrics(pool, provider) +} + +func registerPoolMetrics( + pool *xpg.Pool, + provider metric.MeterProvider, +) (xpg.PoolMetricsRegistration, error) { + meter := provider.Meter(instrumentationName) + + instruments, err := newPoolMetricInstruments(meter) + if err != nil { + return nil, err + } + + attributes := newPoolMetricAttributes(pool.Name(), pool.Labels()) + + registration, err := meter.RegisterCallback( + func(_ context.Context, observer metric.Observer) error { + instruments.observe(observer, pool.Stats(), attributes) + return nil + }, + instruments.observables()..., + ) + if err != nil { + return nil, fmt.Errorf("xpg/otel: register pool metrics callback: %w", err) + } + + return &poolMetrics{ + registration: registration, + }, nil +} + +func (m *poolMetrics) Close() { + if m == nil || m.registration == nil { + return + } + + m.closeOnce.Do(func() { + if err := m.registration.Unregister(); err != nil { + otel.Handle( + fmt.Errorf("xpg/otel: unregister pool metrics: %w", err), + ) + } + }) +} + +func (i poolMetricInstruments) observe( + observer metric.Observer, + stats xpg.PoolStats, + attributes poolMetricAttributes, +) { + observer.ObserveInt64( + i.connectionCount, + int64(stats.IdleConns), + attributes.idle, + ) + observer.ObserveInt64( + i.connectionCount, + int64(stats.AcquiredConns), + attributes.used, + ) + observer.ObserveInt64( + i.connectionMax, + int64(stats.MaxConns), + attributes.base, + ) + observer.ObserveInt64( + i.constructingConnections, + int64(stats.ConstructingConns), + attributes.base, + ) + observer.ObserveInt64( + i.acquireCount, + stats.AcquireCount, + attributes.base, + ) + observer.ObserveFloat64( + i.acquireTime, + stats.AcquireDuration.Seconds(), + attributes.base, + ) + observer.ObserveInt64( + i.canceledAcquireCount, + stats.CanceledAcquireCount, + attributes.base, + ) + observer.ObserveInt64( + i.emptyAcquireCount, + stats.EmptyAcquireCount, + attributes.base, + ) + observer.ObserveFloat64( + i.emptyAcquireWaitTime, + stats.EmptyAcquireWaitTime.Seconds(), + attributes.base, + ) + observer.ObserveInt64( + i.createdConnections, + stats.NewConnsCount, + attributes.base, + ) + observer.ObserveInt64( + i.destroyedConnections, + stats.MaxIdleDestroyCount, + attributes.destroyedIdle, + ) + observer.ObserveInt64( + i.destroyedConnections, + stats.MaxLifetimeDestroyCount, + attributes.destroyedLifetime, + ) +} + +func (i poolMetricInstruments) observables() []metric.Observable { + return []metric.Observable{ + i.connectionCount, + i.connectionMax, + i.constructingConnections, + i.acquireCount, + i.acquireTime, + i.canceledAcquireCount, + i.emptyAcquireCount, + i.emptyAcquireWaitTime, + i.createdConnections, + i.destroyedConnections, + } +} + +func newPoolMetricInstruments(meter metric.Meter) (poolMetricInstruments, error) { + var instruments poolMetricInstruments + + var err error + + instruments.connectionCount, err = meter.Int64ObservableUpDownCounter( + connectionCountMetricName, + metric.WithDescription( + "The number of connections currently used or idle in the pool.", + ), + metric.WithUnit("{connection}"), + ) + if err != nil { + return poolMetricInstruments{}, fmt.Errorf( + "xpg/otel: create %s: %w", + connectionCountMetricName, + err, + ) + } + + instruments.connectionMax, err = meter.Int64ObservableUpDownCounter( + connectionMaxMetricName, + metric.WithDescription( + "The maximum number of open connections allowed by the pool.", + ), + metric.WithUnit("{connection}"), + ) + if err != nil { + return poolMetricInstruments{}, fmt.Errorf( + "xpg/otel: create %s: %w", + connectionMaxMetricName, + err, + ) + } + + instruments.constructingConnections, err = meter.Int64ObservableGauge( + connectionConstructingMetricName, + metric.WithDescription( + "The number of connections currently being created by the pool.", + ), + metric.WithUnit("{connection}"), + ) + if err != nil { + return poolMetricInstruments{}, fmt.Errorf( + "xpg/otel: create %s: %w", + connectionConstructingMetricName, + err, + ) + } + + instruments.acquireCount, err = meter.Int64ObservableCounter( + connectionAcquireCountMetricName, + metric.WithDescription( + "The cumulative number of successful connection acquires.", + ), + metric.WithUnit("{request}"), + ) + if err != nil { + return poolMetricInstruments{}, fmt.Errorf( + "xpg/otel: create %s: %w", + connectionAcquireCountMetricName, + err, + ) + } + + instruments.acquireTime, err = meter.Float64ObservableCounter( + connectionAcquireTimeMetricName, + metric.WithDescription( + "The cumulative time spent acquiring connections.", + ), + metric.WithUnit("s"), + ) + if err != nil { + return poolMetricInstruments{}, fmt.Errorf( + "xpg/otel: create %s: %w", + connectionAcquireTimeMetricName, + err, + ) + } + + instruments.canceledAcquireCount, err = meter.Int64ObservableCounter( + connectionAcquireCanceledCountMetricName, + metric.WithDescription( + "The cumulative number of connection acquires canceled by context.", + ), + metric.WithUnit("{request}"), + ) + if err != nil { + return poolMetricInstruments{}, fmt.Errorf( + "xpg/otel: create %s: %w", + connectionAcquireCanceledCountMetricName, + err, + ) + } + + instruments.emptyAcquireCount, err = meter.Int64ObservableCounter( + connectionAcquireEmptyCountMetricName, + metric.WithDescription( + "The cumulative number of successful acquires that waited because the pool was empty.", + ), + metric.WithUnit("{request}"), + ) + if err != nil { + return poolMetricInstruments{}, fmt.Errorf( + "xpg/otel: create %s: %w", + connectionAcquireEmptyCountMetricName, + err, + ) + } + + instruments.emptyAcquireWaitTime, err = meter.Float64ObservableCounter( + connectionAcquireEmptyWaitTimeMetricName, + metric.WithDescription( + "The cumulative time spent waiting for a connection while the pool was empty.", + ), + metric.WithUnit("s"), + ) + if err != nil { + return poolMetricInstruments{}, fmt.Errorf( + "xpg/otel: create %s: %w", + connectionAcquireEmptyWaitTimeMetricName, + err, + ) + } + + instruments.createdConnections, err = meter.Int64ObservableCounter( + connectionCreateCountMetricName, + metric.WithDescription( + "The cumulative number of connections opened by the pool.", + ), + metric.WithUnit("{connection}"), + ) + if err != nil { + return poolMetricInstruments{}, fmt.Errorf( + "xpg/otel: create %s: %w", + connectionCreateCountMetricName, + err, + ) + } + + instruments.destroyedConnections, err = meter.Int64ObservableCounter( + connectionDestroyCountMetricName, + metric.WithDescription( + "The cumulative number of connections destroyed by pool lifecycle limits.", + ), + metric.WithUnit("{connection}"), + ) + if err != nil { + return poolMetricInstruments{}, fmt.Errorf( + "xpg/otel: create %s: %w", + connectionDestroyCountMetricName, + err, + ) + } + + return instruments, nil +} + +func newPoolMetricAttributes(name string, labels map[string]string) poolMetricAttributes { + var base []attribute.KeyValue + + for key, value := range labels { + base = append(base, attribute.String(key, value)) + } + + // System attributes are appended last, so xpg-controlled values win when + // user labels contain duplicate keys. + base = append( + base, + attribute.String(dbSystemNameAttribute, dbSystemPostgreSQL), + attribute.String(poolNameAttribute, name), + ) + + option := func(extra ...attribute.KeyValue) metric.ObserveOption { + return metric.WithAttributeSet( + attribute.NewSet( + slices.Concat(base, extra)..., + ), + ) + } + + return poolMetricAttributes{ + base: option(), + idle: option( + attribute.String( + connectionStateAttribute, + connectionStateIdle, + ), + ), + used: option( + attribute.String( + connectionStateAttribute, + connectionStateUsed, + ), + ), + destroyedIdle: option( + attribute.String( + destroyReasonAttribute, + destroyReasonIdleTimeout, + ), + ), + destroyedLifetime: option( + attribute.String( + destroyReasonAttribute, + destroyReasonLifetime, + ), + ), + } +} diff --git a/metrics/otel/options.go b/metrics/otel/options.go new file mode 100644 index 0000000..dac9ab0 --- /dev/null +++ b/metrics/otel/options.go @@ -0,0 +1,54 @@ +package xpgotel + +import ( + "go.opentelemetry.io/otel/metric" +) + +// MetricsOption configures OpenTelemetry pool metrics. +// +// The interface is sealed so options can only be created by this package. +type MetricsOption interface { + apply(*metricsSettings) +} + +type metricsOptionFunc func(*metricsSettings) + +func (option metricsOptionFunc) apply(settings *metricsSettings) { + option(settings) +} + +type metricsSettings struct { + meterProvider metric.MeterProvider +} + +// NewMetrics creates an OpenTelemetry pool metrics implementation. +// +// By default, metrics use the global OpenTelemetry MeterProvider. The returned +// value is immutable and may be reused for multiple pools. +func NewMetrics(options ...MetricsOption) *Metrics { + settings := metricsSettings{} + + for _, option := range options { + if option == nil { + continue + } + + option.apply(&settings) + } + + return &Metrics{ + meterProvider: settings.meterProvider, + } +} + +// WithMeterProvider configures the MeterProvider used for pool metrics. +// +// The caller owns the provider and must shut it down after all instrumented +// pools have been closed. +func WithMeterProvider(provider metric.MeterProvider) MetricsOption { + return metricsOptionFunc(func(settings *metricsSettings) { + if provider != nil { + settings.meterProvider = provider + } + }) +} diff --git a/metrics/otel/options_test.go b/metrics/otel/options_test.go new file mode 100644 index 0000000..da0d53e --- /dev/null +++ b/metrics/otel/options_test.go @@ -0,0 +1,57 @@ +package xpgotel + +import ( + "context" + "strings" + "testing" + + "github.com/jackc/pgx/v5/pgxpool" + "github.com/mkbeh/xpg" + "go.opentelemetry.io/otel/metric/noop" +) + +func TestWithMetrics(t *testing.T) { + t.Parallel() + + config, err := pgxpool.ParseConfig("") + if err != nil { + t.Fatalf("parse pool config: %v", err) + } + + pool, err := xpg.New( + context.Background(), + config, + xpg.WithName("test"), + WithMetrics( + WithMeterProvider(noop.NewMeterProvider()), + ), + ) + if err != nil { + t.Fatalf("create pool: %v", err) + } + + pool.Close() + pool.Close() +} + +func TestWithMetricsNilProvider(t *testing.T) { + t.Parallel() + + config, err := pgxpool.ParseConfig("") + if err != nil { + t.Fatalf("parse pool config: %v", err) + } + + pool, err := xpg.New( + context.Background(), + config, + WithMetrics(WithMeterProvider(nil)), + ) + if pool != nil { + pool.Close() + t.Fatal("New returned a pool with a nil MeterProvider") + } + if err == nil || !strings.Contains(err.Error(), "meter provider is nil") { + t.Fatalf("unexpected error: %v", err) + } +} From 5d812de111a2e3e8c02a46747728ed9a811d31da Mon Sep 17 00:00:00 2001 From: mkbeh Date: Wed, 5 Aug 2026 13:34:31 +0300 Subject: [PATCH 08/41] docs: add OpenTelemetry metrics example --- examples/README.md | 1 + examples/basic/README.md | 19 ++-- examples/docker-compose.yml | 2 +- examples/otel/README.md | 140 ++++++++++++++++++++++++++ examples/otel/go.mod | 37 +++++++ examples/otel/main.go | 190 ++++++++++++++++++++++++++++++++++++ examples/otel/metrics.go | 70 +++++++++++++ 7 files changed, 450 insertions(+), 9 deletions(-) create mode 100644 examples/otel/README.md create mode 100644 examples/otel/go.mod create mode 100644 examples/otel/main.go create mode 100644 examples/otel/metrics.go diff --git a/examples/README.md b/examples/README.md index e44c38d..b6bca32 100644 --- a/examples/README.md +++ b/examples/README.md @@ -7,6 +7,7 @@ This directory contains runnable examples demonstrating the main features and us | [`basic`](basic) | Pool lifecycle and common query methods | | [`transactions`](transactions) | Committing an outer transaction after an optional operation is rolled back to a savepoint | | [`advisory`](advisory) | Coordinating concurrent transactions with PostgreSQL advisory locks | +| [`otel`](otel) | Exporting pool metrics through OpenTelemetry and Prometheus | ## Running the examples diff --git a/examples/basic/README.md b/examples/basic/README.md index a1f2a21..adef86e 100644 --- a/examples/basic/README.md +++ b/examples/basic/README.md @@ -26,21 +26,28 @@ Set `XPG_DATABASE_URL` to use another PostgreSQL instance: export XPG_DATABASE_URL='postgres://user:password@localhost:5432/database?sslmode=disable' ``` -## Local PostgreSQL setup +## Local setup -From the repository root, start PostgreSQL and Adminer: +Start PostgreSQL and Adminer from the repository root: ```shell docker compose -f examples/docker-compose.yml --profile tools up -d ``` -Or from this directory: +Or from this example directory: ```shell docker compose -f ../docker-compose.yml --profile tools up -d ``` -Adminer is available at . Sign in with: +Services are available at: + +```text +PostgreSQL: localhost:5432 +Adminer: http://localhost:8080 +``` + +Sign in to Adminer with: ```text System: PostgreSQL @@ -50,10 +57,6 @@ Password: postgres Database: postgres ``` -> [!IMPORTANT] -> Use `postgres`, not `localhost`, in the **Server** field. Adminer connects through the Docker Compose network, where -> PostgreSQL is available by its service name. - ## Run From this directory: diff --git a/examples/docker-compose.yml b/examples/docker-compose.yml index 4cb6461..24523d9 100644 --- a/examples/docker-compose.yml +++ b/examples/docker-compose.yml @@ -17,7 +17,7 @@ services: start_period: 5s adminer: - image: adminer:5.4.2-standalone + image: adminer:standalone profiles: - tools environment: diff --git a/examples/otel/README.md b/examples/otel/README.md new file mode 100644 index 0000000..4e94ae5 --- /dev/null +++ b/examples/otel/README.md @@ -0,0 +1,140 @@ +# OpenTelemetry Metrics Example + +This example shows how to export `xpg` connection pool metrics through the OpenTelemetry Prometheus exporter. + +**This example demonstrates:** + +* Exporting `xpg` pool metrics with OpenTelemetry and Prometheus +* Registering pool metrics through the global `MeterProvider` +* Generating pool contention through the `/load` endpoint +* Shutting down the pool and metrics provider in the correct order + +## Configuration + +The example uses the following connection string by default: + +```text +postgres://postgres:postgres@localhost:5432/postgres?sslmode=disable +``` + +Set `XPG_DATABASE_URL` to use another PostgreSQL instance: + +```shell +export XPG_DATABASE_URL='postgres://user:password@localhost:5432/database?sslmode=disable' +``` + +## Local setup + +Start PostgreSQL and Adminer from the repository root: + +```shell +docker compose -f examples/docker-compose.yml --profile tools up -d +``` + +Or from this example directory: + +```shell +docker compose -f ../docker-compose.yml --profile tools up -d +``` + +Services are available at: + +```text +PostgreSQL: localhost:5432 +Adminer: http://localhost:8080 +``` + +Sign in to Adminer with: + +```text +System: PostgreSQL +Server: postgres +Username: postgres +Password: postgres +Database: postgres +``` + +## Run + +From this directory: + +```shell +go run . +``` + +Or from the repository root: + +```shell +go run ./examples/otel +``` + +The HTTP server starts on: + +```text +http://localhost:9464 +``` + +## View metrics + +Open the Prometheus endpoint: + +```shell +curl 'http://localhost:9464/metrics' +``` + +Show only database connection pool and `xpg` metrics: + +```shell +curl -s 'http://localhost:9464/metrics' \ + | grep -E '^(db_client_connection|xpg_pool_connection_)' +``` + +The example exports these metric families: + +```text +db_client_connection_count +db_client_connection_max +xpg_pool_connection_constructing +xpg_pool_connection_acquire_count_total +xpg_pool_connection_acquire_time_seconds_total +xpg_pool_connection_acquire_canceled_count_total +xpg_pool_connection_acquire_empty_count_total +xpg_pool_connection_acquire_empty_wait_time_seconds_total +xpg_pool_connection_create_count_total +xpg_pool_connection_destroy_count_total +``` + +The Prometheus exporter converts OpenTelemetry dotted instrument names to Prometheus-compatible names and adds unit and +counter suffixes where required. + +## Generate load + +Run the debug workload: + +```shell +curl -X POST 'http://localhost:9464/load' +``` + +While it is running, inspect the pool metrics from another terminal: + +```shell +curl -s 'http://localhost:9464/metrics' \ + | grep -E 'db_client_connection_count|xpg_pool_connection_acquire_' +``` + +The workload runs six concurrent queries against a pool limited to two connections, making connection usage and wait +metrics visible. + +## Stop services + +From the repository root: + +```shell +docker compose -f examples/docker-compose.yml --profile tools down --remove-orphans -v +``` + +Or from this example directory: + +```shell +docker compose -f ../docker-compose.yml --profile tools down --remove-orphans -v +``` \ No newline at end of file diff --git a/examples/otel/go.mod b/examples/otel/go.mod new file mode 100644 index 0000000..a0aabe9 --- /dev/null +++ b/examples/otel/go.mod @@ -0,0 +1,37 @@ +module observability + +go 1.26 + +require ( + github.com/jackc/pgx/v5 v5.10.0 + github.com/mkbeh/xpg v0.2.0 + github.com/prometheus/client_golang v1.24.1 + go.opentelemetry.io/otel v1.44.0 + go.opentelemetry.io/otel/exporters/prometheus v0.67.0 + go.opentelemetry.io/otel/sdk v1.45.0 + go.opentelemetry.io/otel/sdk/metric v1.45.0 +) + +require ( + github.com/beorn7/perks v1.0.1 // indirect + github.com/cespare/xxhash/v2 v2.3.0 // indirect + github.com/go-logr/logr v1.4.3 // indirect + github.com/go-logr/stdr v1.2.2 // indirect + github.com/google/uuid v1.6.0 // indirect + github.com/jackc/pgpassfile v1.0.0 // indirect + github.com/jackc/pgservicefile v0.0.0-20240606120523-5a60cdf6a761 // indirect + github.com/jackc/puddle/v2 v2.2.2 // indirect + github.com/munnerz/goautoneg v0.0.0-20191010083416-a7dc8b61c822 // indirect + github.com/prometheus/client_model v0.6.2 // indirect + github.com/prometheus/common v0.69.0 // indirect + github.com/prometheus/otlptranslator v1.0.0 // indirect + github.com/prometheus/procfs v0.20.1 // indirect + go.opentelemetry.io/auto/sdk v1.2.1 // indirect + go.opentelemetry.io/otel/metric v1.44.0 // indirect + go.opentelemetry.io/otel/metric/x v0.66.0 // indirect + go.opentelemetry.io/otel/trace v1.44.0 // indirect + golang.org/x/sync v0.22.0 // indirect + golang.org/x/sys v0.46.0 // indirect + golang.org/x/text v0.40.0 // indirect + google.golang.org/protobuf v1.36.11 // indirect +) diff --git a/examples/otel/main.go b/examples/otel/main.go new file mode 100644 index 0000000..13f0f6f --- /dev/null +++ b/examples/otel/main.go @@ -0,0 +1,190 @@ +package main + +import ( + "context" + "errors" + "fmt" + "log" + "net/http" + "os" + "os/signal" + "syscall" + "time" + + "github.com/jackc/pgx/v5/pgxpool" + "github.com/mkbeh/xpg" + xpgotel "github.com/mkbeh/xpg/metrics/otel" +) + +const ( + defaultDatabaseURL = "postgres://postgres:postgres@localhost:5432/postgres?sslmode=disable" + defaultHTTPAddress = "localhost:9464" +) + +func main() { + ctx, stop := signal.NotifyContext(context.Background(), os.Interrupt, syscall.SIGTERM) + defer stop() + + if err := run(ctx); err != nil { + log.Fatal(err) + } +} + +func run(ctx context.Context) (runErr error) { + metrics, err := newMetricsRuntime(ctx) + if err != nil { + return fmt.Errorf("initialize metrics: %w", err) + } + defer func() { + shutdownCtx, cancel := context.WithTimeout(context.Background(), 5*time.Second) + defer cancel() + + runErr = errors.Join( + runErr, + metrics.Shutdown(shutdownCtx), + ) + }() + + config, err := pgxpool.ParseConfig(databaseURL()) + if err != nil { + return fmt.Errorf("parse PostgreSQL config: %w", err) + } + + // A small pool makes contention visible when POST /load runs. + config.MaxConns = 2 + + pool, err := xpg.New( + ctx, + config, + xpg.WithName("otel-example"), + xpg.WithLabel("xpg.pool.role", "primary"), + xpg.WithMetrics( + xpgotel.NewMetrics(), + ), + ) + if err != nil { + return fmt.Errorf("create pool: %w", err) + } + defer pool.Close() + + if err := pool.Ping(ctx); err != nil { + return fmt.Errorf("ping PostgreSQL: %w", err) + } + + mux := http.NewServeMux() + + mux.Handle("GET /metrics", metrics.Handler()) + mux.HandleFunc("POST /load", loadHandler(pool)) + + server := &http.Server{ + Addr: httpAddress(), + Handler: mux, + ReadHeaderTimeout: 5 * time.Second, + } + + log.Printf("OpenTelemetry example listening on http://%s", server.Addr) + + if err := serveHTTP(ctx, server); err != nil { + return fmt.Errorf("serve HTTP: %w", err) + } + + return nil +} + +func serveHTTP(ctx context.Context, server *http.Server) error { + serverErr := make(chan error, 1) + + go func() { + serverErr <- server.ListenAndServe() + }() + + select { + case err := <-serverErr: + if errors.Is(err, http.ErrServerClosed) { + return nil + } + + return err + + case <-ctx.Done(): + } + + shutdownCtx, cancel := context.WithTimeout(context.Background(), 5*time.Second) + defer cancel() + + if err := server.Shutdown(shutdownCtx); err != nil { + return fmt.Errorf("shutdown: %w", err) + } + + if err := <-serverErr; !errors.Is(err, http.ErrServerClosed) { + return err + } + + return nil +} + +func loadHandler(pool *xpg.Pool) http.HandlerFunc { + return func(w http.ResponseWriter, request *http.Request) { + startedAt := time.Now() + + if err := runWorkload(request.Context(), pool); err != nil { + http.Error( + w, + fmt.Sprintf("run workload: %v", err), + http.StatusInternalServerError, + ) + + return + } + + _, _ = fmt.Fprintf( + w, + "workload completed in %s\n", + time.Since(startedAt), + ) + } +} + +func runWorkload(ctx context.Context, pool *xpg.Pool) error { + const workerCount = 6 + + start := make(chan struct{}) + results := make(chan error, workerCount) + + for range workerCount { + go func() { + <-start + _, err := pool.Exec(ctx, "SELECT pg_sleep(1)") + results <- err + }() + } + + close(start) + + var workloadErr error + + for range workerCount { + workloadErr = errors.Join( + workloadErr, + <-results, + ) + } + + return workloadErr +} + +func databaseURL() string { + if value := os.Getenv("XPG_DATABASE_URL"); value != "" { + return value + } + + return defaultDatabaseURL +} + +func httpAddress() string { + if value := os.Getenv("HTTP_ADDR"); value != "" { + return value + } + + return defaultHTTPAddress +} diff --git a/examples/otel/metrics.go b/examples/otel/metrics.go new file mode 100644 index 0000000..29f9cbc --- /dev/null +++ b/examples/otel/metrics.go @@ -0,0 +1,70 @@ +package main + +import ( + "context" + "fmt" + "net/http" + + promclient "github.com/prometheus/client_golang/prometheus" + "github.com/prometheus/client_golang/prometheus/promhttp" + "go.opentelemetry.io/otel" + otelprom "go.opentelemetry.io/otel/exporters/prometheus" + sdkmetric "go.opentelemetry.io/otel/sdk/metric" + "go.opentelemetry.io/otel/sdk/resource" + semconv "go.opentelemetry.io/otel/semconv/v1.37.0" +) + +type metricsRuntime struct { + handler http.Handler + meterProvider *sdkmetric.MeterProvider +} + +func newMetricsRuntime(ctx context.Context) (*metricsRuntime, error) { + registry := promclient.NewRegistry() + + exporter, err := otelprom.New( + otelprom.WithRegisterer(registry), + otelprom.WithoutScopeInfo(), + ) + if err != nil { + return nil, fmt.Errorf("create Prometheus exporter: %w", err) + } + + res, err := resource.New( + ctx, + resource.WithFromEnv(), + resource.WithTelemetrySDK(), + resource.WithAttributes( + semconv.ServiceName( + "xpg-observability-example", + ), + semconv.ServiceVersion("dev"), + ), + ) + if err != nil { + return nil, fmt.Errorf("create OpenTelemetry resource: %w", err) + } + + meterProvider := sdkmetric.NewMeterProvider( + sdkmetric.WithResource(res), + sdkmetric.WithReader(exporter), + ) + + otel.SetMeterProvider(meterProvider) + + return &metricsRuntime{ + handler: promhttp.HandlerFor( + registry, + promhttp.HandlerOpts{}, + ), + meterProvider: meterProvider, + }, nil +} + +func (m *metricsRuntime) Handler() http.Handler { + return m.handler +} + +func (m *metricsRuntime) Shutdown(ctx context.Context) error { + return m.meterProvider.Shutdown(ctx) +} From a2c0ee013c5c768cff370bcc7351a17c3befb2d1 Mon Sep 17 00:00:00 2001 From: mkbeh Date: Fri, 7 Aug 2026 18:03:40 +0300 Subject: [PATCH 09/41] feat(cluster): add primary and replica routing --- cluster/cluster.go | 130 ++++++++++++++++++++++++++++++++++++++++++++ cluster/doc.go | 6 ++ cluster/errors.go | 7 +++ cluster/resolver.go | 112 ++++++++++++++++++++++++++++++++++++++ cluster/selector.go | 102 ++++++++++++++++++++++++++++++++++ cluster/tx.go | 54 ++++++++++++++++++ 6 files changed, 411 insertions(+) create mode 100644 cluster/cluster.go create mode 100644 cluster/doc.go create mode 100644 cluster/errors.go create mode 100644 cluster/resolver.go create mode 100644 cluster/selector.go create mode 100644 cluster/tx.go diff --git a/cluster/cluster.go b/cluster/cluster.go new file mode 100644 index 0000000..153d00c --- /dev/null +++ b/cluster/cluster.go @@ -0,0 +1,130 @@ +package cluster + +import ( + "errors" + "fmt" + "slices" + "sync" + + "github.com/mkbeh/xpg" +) + +// Config configures a Cluster from independently created pools. +// +// New takes ownership of Primary and Replicas only after it returns +// successfully. Cluster.Close closes the owned pools. +type Config struct { + Primary *xpg.Pool + Replicas []*xpg.Pool + Selector ReplicaSelector +} + +// Cluster routes operations between one primary pool and optional replica +// pools. +// +// A deployment with one PostgreSQL endpoint is represented by a Cluster with +// one Primary and no Replicas. +// +// Cluster does not inspect SQL, retry failed queries, promote replicas, or +// discover PostgreSQL nodes. Those responsibilities remain with the caller or +// the surrounding high-availability infrastructure. +type Cluster struct { + primary *xpg.Pool + replicas []*xpg.Pool + + metadata replicaMetadata + selector ReplicaSelector + + closeOnce sync.Once +} + +// New creates a Cluster from independently configured pools. +// +// Primary is required. Replicas may be omitted. When Selector is nil, +// replicas are selected using round-robin. +func New(config Config) (*Cluster, error) { + if config.Primary != nil && invalidPool(config.Primary) { + return nil, errors.New("xpg/cluster: primary pool is invalid") + } + + if config.Primary == nil && len(config.Replicas) == 0 { + return nil, errors.New("xpg/cluster: at least one pool is required") + } + + replicas := slices.Clone(config.Replicas) + metadata := make(replicaMetadata, len(replicas)) + + for index, replica := range replicas { + if invalidPool(replica) { + return nil, fmt.Errorf("xpg/cluster: replica %d is nil", index) + } + + metadata[index] = ReplicaInfo{ + name: replica.Name(), + labels: cloneLabels(replica.Labels()), + } + } + + selector := config.Selector + if selector == nil { + selector = RoundRobinSelector() + } + + return &Cluster{ + primary: config.Primary, + replicas: replicas, + metadata: metadata, + selector: selector, + }, nil +} + +func invalidPool(pool *xpg.Pool) bool { + return pool == nil || pool.Raw() == nil +} + +// Primary returns the primary pool owned by the cluster. +// +// The returned pool is borrowed and must not be closed separately. +func (c *Cluster) Primary() *xpg.Pool { + if c == nil { + return nil + } + + return c.primary +} + +// ReplicaCount returns the number of replicas registered in the cluster. +func (c *Cluster) ReplicaCount() int { + if c == nil { + return 0 + } + + return len(c.replicas) +} + +// ReplicaAt returns the replica at index in registration order. +// +// The returned pool is borrowed and must not be closed separately. ReplicaAt +// panics when c is nil or index is outside the replica set, matching ordinary +// slice indexing semantics. +func (c *Cluster) ReplicaAt(index int) *xpg.Pool { + return c.replicas[index] +} + +// Close closes all replica pools in reverse registration order and then closes +// the primary pool. +// +// Close is safe to call multiple times. +func (c *Cluster) Close() { + if c == nil { + return + } + + c.closeOnce.Do(func() { + for index := len(c.replicas) - 1; index >= 0; index-- { + c.replicas[index].Close() + } + + c.primary.Close() + }) +} diff --git a/cluster/doc.go b/cluster/doc.go new file mode 100644 index 0000000..609a36b --- /dev/null +++ b/cluster/doc.go @@ -0,0 +1,6 @@ +// Package cluster provides explicit routing between one PostgreSQL primary +// pool and zero or more replica pools. +// +// The package does not inspect SQL, retry failed queries, promote replicas, +// or discover PostgreSQL nodes. +package cluster diff --git a/cluster/errors.go b/cluster/errors.go new file mode 100644 index 0000000..b205250 --- /dev/null +++ b/cluster/errors.go @@ -0,0 +1,7 @@ +package cluster + +import "errors" + +// ErrNoReplica indicates that a read policy required a replica but no replica +// was available for selection. +var ErrNoReplica = errors.New("xpg/cluster: no replica available") diff --git a/cluster/resolver.go b/cluster/resolver.go new file mode 100644 index 0000000..e3514a7 --- /dev/null +++ b/cluster/resolver.go @@ -0,0 +1,112 @@ +package cluster + +import ( + "context" + "errors" + "fmt" + + "github.com/mkbeh/xpg" +) + +// ReadPolicy defines where a Cluster resolves a read operation. +type ReadPolicy uint8 + +const ( + // ReadPrimary always resolves reads to the primary pool. + ReadPrimary ReadPolicy = iota + + // ReadReplicaPreferred prefers a replica and falls back to the primary when + // no replica can be selected. + ReadReplicaPreferred + + // ReadReplicaRequired requires a replica and returns ErrNoReplica when none + // can be selected. + ReadReplicaRequired +) + +const ( + readPolicyPrimary = "primary" + readPolicyReplicaPreferred = "replica_preferred" + readPolicyReplicaRequired = "replica_required" +) + +// ParsePolicy parses a ReadPolicy from its string representation. +func ParsePolicy(value string) (ReadPolicy, error) { + switch value { + case readPolicyPrimary: + return ReadPrimary, nil + case readPolicyReplicaPreferred: + return ReadReplicaPreferred, nil + case readPolicyReplicaRequired: + return ReadReplicaRequired, nil + default: + return 0, fmt.Errorf( + "xpg/cluster: unknown read policy %q", + value, + ) + } +} + +// String returns the string representation of the read policy. +func (policy ReadPolicy) String() string { + switch policy { + case ReadPrimary: + return readPolicyPrimary + case ReadReplicaPreferred: + return readPolicyReplicaPreferred + case ReadReplicaRequired: + return readPolicyReplicaRequired + default: + return "unknown" + } +} + +// ReadPool returns a pool for a read operation according to policy. +// +// ReadReplicaPreferred falls back to the primary only when replica selection +// returns ErrNoReplica. Other selector errors are returned to the caller. +func (c *Cluster) ReadPool(ctx context.Context, policy ReadPolicy) (*xpg.Pool, error) { + if c == nil || c.primary == nil { + return nil, errors.New("xpg/cluster: cluster is nil") + } + + switch policy { + case ReadPrimary: + return c.primary, nil + + case ReadReplicaPreferred: + replica, err := c.selectReplica(ctx) + if errors.Is(err, ErrNoReplica) { + return c.primary, nil + } + + return replica, err + + case ReadReplicaRequired: + return c.selectReplica(ctx) + + default: + return nil, fmt.Errorf("xpg/cluster: unsupported read policy %d", policy) + } +} + +func (c *Cluster) selectReplica(ctx context.Context) (*xpg.Pool, error) { + if len(c.replicas) == 0 { + return nil, ErrNoReplica + } + + index, err := c.selector.Select(ctx, c.metadata) + if err != nil { + return nil, fmt.Errorf("xpg/cluster: select replica: %w", err) + } + + if index < 0 || index >= len(c.replicas) { + return nil, fmt.Errorf( + "xpg/cluster: replica selector returned index %d for %d replicas", + index, + len(c.replicas), + ) + } + + return c.replicas[index], nil +} diff --git a/cluster/selector.go b/cluster/selector.go new file mode 100644 index 0000000..5f0fb85 --- /dev/null +++ b/cluster/selector.go @@ -0,0 +1,102 @@ +package cluster + +import ( + "context" + "errors" + "maps" + "sync/atomic" +) + +// ReplicaInfo contains immutable metadata captured from one replica pool when +// the Cluster is created. +type ReplicaInfo struct { + name string + labels map[string]string +} + +// Name returns the stable logical pool name. +func (info ReplicaInfo) Name() string { + return info.name +} + +// Label returns one replica label without allocating a copy of all labels. +func (info ReplicaInfo) Label(key string) (string, bool) { + value, ok := info.labels[key] + return value, ok +} + +// Labels returns a defensive copy of replica labels. +func (info ReplicaInfo) Labels() map[string]string { + return cloneLabels(info.labels) +} + +// ReplicaSet provides read-only access to replica metadata. +type ReplicaSet interface { + Len() int + At(index int) ReplicaInfo +} + +type replicaMetadata []ReplicaInfo + +func (replicas replicaMetadata) Len() int { + return len(replicas) +} + +func (replicas replicaMetadata) At(index int) ReplicaInfo { + return replicas[index] +} + +// ReplicaSelector selects one replica index from the supplied metadata. +// Implementations used by concurrent callers must be concurrency-safe. +type ReplicaSelector interface { + Select(ctx context.Context, replicas ReplicaSet) (index int, err error) +} + +// ReplicaSelectorFunc adapts a function to ReplicaSelector. +type ReplicaSelectorFunc func(context.Context, ReplicaSet) (int, error) + +// Select calls selector. +func (selector ReplicaSelectorFunc) Select(ctx context.Context, replicas ReplicaSet) (int, error) { + if selector == nil { + return -1, errors.New("xpg/cluster: replica selector function is nil") + } + + return selector(ctx, replicas) +} + +type roundRobinSelector struct { + next atomic.Uint64 +} + +// RoundRobinSelector returns a concurrency-safe selector that distributes +// selections across replicas in registration order. +func RoundRobinSelector() ReplicaSelector { + return &roundRobinSelector{} +} + +func (selector *roundRobinSelector) Select(_ context.Context, replicas ReplicaSet) (int, error) { + length := replicas.Len() + + if length == 0 { + return -1, ErrNoReplica + } + + if length == 1 { + return 0, nil + } + + next := selector.next.Add(1) - 1 + + return int(next % uint64(length)), nil +} + +func cloneLabels(labels map[string]string) map[string]string { + if len(labels) == 0 { + return nil + } + + cloned := make(map[string]string, len(labels)) + maps.Copy(cloned, labels) + + return cloned +} diff --git a/cluster/tx.go b/cluster/tx.go new file mode 100644 index 0000000..9d5a430 --- /dev/null +++ b/cluster/tx.go @@ -0,0 +1,54 @@ +package cluster + +import ( + "context" + "errors" + + "github.com/jackc/pgx/v5" +) + +// ReadTxOptions configures a read-only transaction. +// +// AccessMode, BeginQuery, and CommitQuery are intentionally controlled by the +// cluster. +type ReadTxOptions struct { + IsoLevel pgx.TxIsoLevel + DeferrableMode pgx.TxDeferrableMode +} + +// InPrimaryTx executes fn in a transaction on the primary pool. +func (c *Cluster) InPrimaryTx( + ctx context.Context, + options pgx.TxOptions, + fn func(context.Context, pgx.Tx) error, +) error { + if c == nil || c.primary == nil { + return errors.New("xpg/cluster: cluster is not initialized") + } + + return c.primary.InTx(ctx, options, fn) +} + +// InReadTx selects a pool according to policy and executes fn in a read-only +// transaction on that pool. +func (c *Cluster) InReadTx( + ctx context.Context, + policy ReadPolicy, + options ReadTxOptions, + fn func(context.Context, pgx.Tx) error, +) error { + pool, err := c.ReadPool(ctx, policy) + if err != nil { + return err + } + + return pool.InTx( + ctx, + pgx.TxOptions{ + IsoLevel: options.IsoLevel, + AccessMode: pgx.ReadOnly, + DeferrableMode: options.DeferrableMode, + }, + fn, + ) +} From 4d2bdab9f9a9d78fa8a14874cee1caedae6a05b5 Mon Sep 17 00:00:00 2001 From: mkbeh Date: Fri, 7 Aug 2026 18:04:11 +0300 Subject: [PATCH 10/41] docs(examples): add cluster routing example --- examples/cluster/README.md | 138 +++++++++++++ examples/cluster/docker-compose.yml | 75 +++++++ examples/cluster/go.mod | 8 + examples/cluster/main.go | 294 +++++++++++++++++++++++++++ examples/cluster/sql/primary.sql | 20 ++ examples/cluster/sql/replica-one.sql | 18 ++ examples/cluster/sql/replica-two.sql | 18 ++ 7 files changed, 571 insertions(+) create mode 100644 examples/cluster/README.md create mode 100644 examples/cluster/docker-compose.yml create mode 100644 examples/cluster/go.mod create mode 100644 examples/cluster/main.go create mode 100644 examples/cluster/sql/primary.sql create mode 100644 examples/cluster/sql/replica-one.sql create mode 100644 examples/cluster/sql/replica-two.sql diff --git a/examples/cluster/README.md b/examples/cluster/README.md new file mode 100644 index 0000000..24c9684 --- /dev/null +++ b/examples/cluster/README.md @@ -0,0 +1,138 @@ +# Cluster routing + +This example routes reads and transactions across one primary pool and two replica pools with `cluster.Cluster`. + +```text + ┌─ primary +application ─ cluster + ├─ replica-one + └─ replica-two +``` + +**This example demonstrates:** + +* Creating a cluster from primary and replica pools +* Routing reads to the primary or replicas +* Distributing replica reads with round-robin +* Running primary and read-only replica transactions + +> [!NOTE] +> The local containers are independent PostgreSQL instances used to demonstrate routing. They do not configure streaming +replication. + +## Configuration + +The example uses the following connection strings by default: + +```text +XPG_PRIMARY_DATABASE_URL=postgres://postgres:postgres@localhost:55432/postgres?sslmode=disable&target_session_attrs=read-write +XPG_REPLICA_ONE_DATABASE_URL=postgres://postgres:postgres@localhost:55433/postgres?sslmode=disable&target_session_attrs=read-only +XPG_REPLICA_TWO_DATABASE_URL=postgres://postgres:postgres@localhost:55434/postgres?sslmode=disable&target_session_attrs=read-only +``` + +Set the corresponding environment variables to use different PostgreSQL endpoints: + +```shell +export XPG_PRIMARY_DATABASE_URL='postgres://user:password@primary.example.com:5432/database?sslmode=disable&target_session_attrs=read-write' +export XPG_REPLICA_ONE_DATABASE_URL='postgres://user:password@replica-one.example.com:5432/database?sslmode=disable&target_session_attrs=read-only' +export XPG_REPLICA_TWO_DATABASE_URL='postgres://user:password@replica-two.example.com:5432/database?sslmode=disable&target_session_attrs=read-only' +``` + +## Local setup + +Start the primary, both replica endpoints, and Adminer from the repository root: + +```shell +docker compose -f examples/cluster/docker-compose.yml --profile tools up -d +``` + +Or from this example directory: + +```shell +docker compose --profile tools up -d +``` + +Services are available at: + +```text +Primary: localhost:55432 +Replica 1: localhost:55433 +Replica 2: localhost:55434 +Adminer: http://localhost:58080 +``` + +Sign in to Adminer with: + +```text +System: PostgreSQL +Server: postgres-primary +Username: postgres +Password: postgres +Database: postgres +``` + +Use `postgres-replica-one` or `postgres-replica-two` in the **Server** field to inspect the replica endpoints. + +## Run + +From this directory: + +```shell +go run . +``` + +Or from the repository root: + +```shell +go run ./examples/basic +``` + +## Run + +From this directory: + +```shell +go run . +``` + +Or from the repository root: + +```shell +go run ./examples/cluster +``` + +## Expected output + +```text +primary: +- pool=cluster.primary node=primary role=primary +replica reads: +- pool=cluster.replica-one node=replica-one role=replica +- pool=cluster.replica-two node=replica-two role=replica +transactions: +- primary node=primary read_only=off +- replica node=replica-one read_only=on +``` + +The example performs one complete round-robin pass across the configured replicas before starting the read-only +transaction. + +Pools remain owned by the caller until `cluster.New` succeeds. After successful cluster creation, `cluster.Cluster` owns +the pools and closes them when `Cluster.Close` is called. + +## Stop services + +From the repository root: + +```shell +docker compose \ + -f examples/cluster/docker-compose.yml \ + --profile tools \ + down --remove-orphans -v +``` + +Or from this example directory: + +```shell +docker compose --profile tools down --remove-orphans -v +``` diff --git a/examples/cluster/docker-compose.yml b/examples/cluster/docker-compose.yml new file mode 100644 index 0000000..00dc2bc --- /dev/null +++ b/examples/cluster/docker-compose.yml @@ -0,0 +1,75 @@ +services: + postgres-primary: + image: postgres:18-alpine + environment: + POSTGRES_DB: postgres + POSTGRES_USER: postgres + POSTGRES_PASSWORD: postgres + ports: + - "55432:5432" + volumes: + - primary-data:/var/lib/postgresql + - ./sql/primary.sql:/docker-entrypoint-initdb.d/10-node.sql:ro + healthcheck: + test: "pg_isready -U $${POSTGRES_USER} -d $${POSTGRES_DB}" + interval: 2s + timeout: 5s + retries: 10 + start_period: 5s + + postgres-replica-one: + image: postgres:18-alpine + environment: + POSTGRES_DB: postgres + POSTGRES_USER: postgres + POSTGRES_PASSWORD: postgres + ports: + - "55433:5432" + volumes: + - replica-one-data:/var/lib/postgresql + - ./sql/replica-one.sql:/docker-entrypoint-initdb.d/10-node.sql:ro + healthcheck: + test: "pg_isready -U $${POSTGRES_USER} -d $${POSTGRES_DB}" + interval: 2s + timeout: 5s + retries: 10 + start_period: 5s + + postgres-replica-two: + image: postgres:18-alpine + environment: + POSTGRES_DB: postgres + POSTGRES_USER: postgres + POSTGRES_PASSWORD: postgres + ports: + - "55434:5432" + volumes: + - replica-two-data:/var/lib/postgresql + - ./sql/replica-two.sql:/docker-entrypoint-initdb.d/10-node.sql:ro + healthcheck: + test: "pg_isready -U $${POSTGRES_USER} -d $${POSTGRES_DB}" + interval: 2s + timeout: 5s + retries: 10 + start_period: 5s + + adminer: + image: adminer:standalone + profiles: + - tools + environment: + ADMINER_DEFAULT_SERVER: postgres-primary + ports: + - "8080:8080" + depends_on: + postgres-primary: + condition: service_healthy + postgres-replica-one: + condition: service_healthy + postgres-replica-two: + condition: service_healthy + +volumes: + primary-data: + replica-one-data: + replica-two-data: diff --git a/examples/cluster/go.mod b/examples/cluster/go.mod new file mode 100644 index 0000000..da620f5 --- /dev/null +++ b/examples/cluster/go.mod @@ -0,0 +1,8 @@ +module cluster + +go 1.26 + +require ( + github.com/jackc/pgx/v5 v5.10.0 + github.com/mkbeh/xpg v0.2.0 +) diff --git a/examples/cluster/main.go b/examples/cluster/main.go new file mode 100644 index 0000000..8165240 --- /dev/null +++ b/examples/cluster/main.go @@ -0,0 +1,294 @@ +package main + +import ( + "context" + "fmt" + "log" + "os" + + "github.com/jackc/pgx/v5" + "github.com/mkbeh/xpg" + "github.com/mkbeh/xpg/cluster" +) + +const ( + defaultPrimaryDatabaseURL = "postgres://postgres:postgres@localhost:55432/postgres?sslmode=disable&target_session_attrs=read-write" + defaultReplicaOneDatabaseURL = "postgres://postgres:postgres@localhost:55433/postgres?sslmode=disable&target_session_attrs=read-only" + defaultReplicaTwoDatabaseURL = "postgres://postgres:postgres@localhost:55434/postgres?sslmode=disable&target_session_attrs=read-only" +) + +type nodeInfo struct { + Name string + Role string +} + +type queryRower interface { + QueryRow(context.Context, string, ...any) pgx.Row +} + +func main() { + if err := run(context.Background()); err != nil { + log.Fatal(err) + } +} + +func run(ctx context.Context) error { + dbCluster, err := openCluster(ctx) + if err != nil { + return err + } + defer dbCluster.Close() + + if err := showRouting(ctx, dbCluster); err != nil { + return err + } + + if err := showTransactions(ctx, dbCluster); err != nil { + return err + } + + return nil +} + +func openCluster(ctx context.Context) (*cluster.Cluster, error) { + nodes := []struct { + databaseURL string + name string + role string + }{ + { + databaseURL: environment( + "XPG_PRIMARY_DATABASE_URL", + defaultPrimaryDatabaseURL, + ), + name: "cluster.primary", + role: "primary", + }, + { + databaseURL: environment( + "XPG_REPLICA_ONE_DATABASE_URL", + defaultReplicaOneDatabaseURL, + ), + name: "cluster.replica-one", + role: "replica", + }, + { + databaseURL: environment( + "XPG_REPLICA_TWO_DATABASE_URL", + defaultReplicaTwoDatabaseURL, + ), + name: "cluster.replica-two", + role: "replica", + }, + } + + pools := make([]*xpg.Pool, 0, len(nodes)) + + for _, node := range nodes { + pool, err := openPool( + ctx, + node.databaseURL, + node.name, + node.role, + ) + if err != nil { + closePools(pools) + + return nil, fmt.Errorf("open %s pool: %w", node.name, err) + } + + pools = append(pools, pool) + } + + // Build a replicated cluster and explicitly use round-robin selection for + // replica reads. + dbCluster, err := cluster.New( + cluster.Config{ + Primary: pools[0], + Replicas: pools[1:], + Selector: cluster.RoundRobinSelector(), + }, + ) + if err != nil { + return nil, fmt.Errorf("create cluster: %w", err) + } + + return dbCluster, nil +} + +func openPool( + ctx context.Context, + databaseURL string, + name string, + role string, +) (*xpg.Pool, error) { + pool, err := xpg.Open( + ctx, + databaseURL, + xpg.WithName(name), + xpg.WithLabel("role", role), + ) + if err != nil { + return nil, err + } + + if err := pool.Ping(ctx); err != nil { + pool.Close() + + return nil, fmt.Errorf("ping %s: %w", name, err) + } + + return pool, nil +} + +func closePools(pools []*xpg.Pool) { + for index := len(pools) - 1; index >= 0; index-- { + pools[index].Close() + } +} + +func showRouting(ctx context.Context, dbCluster *cluster.Cluster) error { + primary := dbCluster.Primary() + + node, err := loadNode(ctx, primary) + if err != nil { + return fmt.Errorf("read primary node: %w", err) + } + + fmt.Println("primary:") + fmt.Printf( + "- pool=%s node=%s role=%s\n", + primary.Name(), + node.Name, + node.Role, + ) + + fmt.Println("replica reads:") + + // Read once per registered replica to demonstrate one complete round-robin + // cycle without hard-coding the cluster size. + for range dbCluster.ReplicaCount() { + pool, err := dbCluster.ReadPool( + ctx, + cluster.ReadReplicaRequired, + ) + if err != nil { + return fmt.Errorf("resolve replica read: %w", err) + } + + node, err := loadNode(ctx, pool) + if err != nil { + return fmt.Errorf("read replica node: %w", err) + } + + fmt.Printf( + "- pool=%s node=%s role=%s\n", + pool.Name(), + node.Name, + node.Role, + ) + } + + return nil +} + +func showTransactions(ctx context.Context, dbCluster *cluster.Cluster) error { + var ( + primaryNode nodeInfo + primaryReadOnly string + ) + + // Primary transactions use regular pgx transaction options and may perform + // both reads and writes. + err := dbCluster.InPrimaryTx( + ctx, + pgx.TxOptions{}, + func(ctx context.Context, tx pgx.Tx) error { + var err error + + primaryNode, err = loadNode(ctx, tx) + if err != nil { + return err + } + + return tx.QueryRow( + ctx, + "SHOW transaction_read_only", + ).Scan(&primaryReadOnly) + }, + ) + if err != nil { + return fmt.Errorf("run primary transaction: %w", err) + } + + var ( + replicaNode nodeInfo + replicaReadOnly string + ) + + // Read transactions resolve their pool through ReadPolicy and always start + // PostgreSQL transactions in READ ONLY mode. + err = dbCluster.InReadTx( + ctx, + cluster.ReadReplicaRequired, + cluster.ReadTxOptions{ + IsoLevel: pgx.RepeatableRead, + }, + func(ctx context.Context, tx pgx.Tx) error { + var err error + + replicaNode, err = loadNode(ctx, tx) + if err != nil { + return err + } + + return tx.QueryRow( + ctx, + "SHOW transaction_read_only", + ).Scan(&replicaReadOnly) + }, + ) + if err != nil { + return fmt.Errorf("run replica transaction: %w", err) + } + + fmt.Println("transactions:") + fmt.Printf( + "- primary node=%s read_only=%s\n", + primaryNode.Name, + primaryReadOnly, + ) + fmt.Printf( + "- replica node=%s read_only=%s\n", + replicaNode.Name, + replicaReadOnly, + ) + + return nil +} + +func loadNode(ctx context.Context, db queryRower) (nodeInfo, error) { + var node nodeInfo + + err := db.QueryRow( + ctx, + `SELECT node_name, node_role + FROM xpg_cluster_example.node_info`, + ).Scan( + &node.Name, + &node.Role, + ) + if err != nil { + return nodeInfo{}, err + } + + return node, nil +} + +func environment(key, fallback string) string { + if value := os.Getenv(key); value != "" { + return value + } + + return fallback +} diff --git a/examples/cluster/sql/primary.sql b/examples/cluster/sql/primary.sql new file mode 100644 index 0000000..d53fe6b --- /dev/null +++ b/examples/cluster/sql/primary.sql @@ -0,0 +1,20 @@ +CREATE SCHEMA xpg_cluster_example; + +CREATE TABLE xpg_cluster_example.node_info ( + node_name text PRIMARY KEY, + node_role text NOT NULL +); + +INSERT INTO xpg_cluster_example.node_info ( + node_name, + node_role +) +VALUES ( + 'primary', + 'primary' +); + +CREATE TABLE xpg_cluster_example.primary_writes ( + id bigint GENERATED ALWAYS AS IDENTITY PRIMARY KEY, + created_at timestamptz NOT NULL DEFAULT now() +); diff --git a/examples/cluster/sql/replica-one.sql b/examples/cluster/sql/replica-one.sql new file mode 100644 index 0000000..0973062 --- /dev/null +++ b/examples/cluster/sql/replica-one.sql @@ -0,0 +1,18 @@ +CREATE SCHEMA xpg_cluster_example; + +CREATE TABLE xpg_cluster_example.node_info ( + node_name text PRIMARY KEY, + node_role text NOT NULL +); + +INSERT INTO xpg_cluster_example.node_info ( + node_name, + node_role +) +VALUES ( + 'replica-one', + 'replica' +); + +ALTER DATABASE postgres +SET default_transaction_read_only = on; diff --git a/examples/cluster/sql/replica-two.sql b/examples/cluster/sql/replica-two.sql new file mode 100644 index 0000000..968eb46 --- /dev/null +++ b/examples/cluster/sql/replica-two.sql @@ -0,0 +1,18 @@ +CREATE SCHEMA xpg_cluster_example; + +CREATE TABLE xpg_cluster_example.node_info ( + node_name text PRIMARY KEY, + node_role text NOT NULL +); + +INSERT INTO xpg_cluster_example.node_info ( + node_name, + node_role +) +VALUES ( + 'replica-two', + 'replica' +); + +ALTER DATABASE postgres +SET default_transaction_read_only = on; From 2325fce5eb71a2e4a9136c062634296f2b4489e5 Mon Sep 17 00:00:00 2001 From: mkbeh Date: Tue, 11 Aug 2026 16:42:42 +0300 Subject: [PATCH 11/41] feat(shard): add topology and shard routing --- cluster/cluster.go | 83 ++++++++++++++-- cluster/doc.go | 4 +- cluster/errors.go | 12 ++- cluster/resolver.go | 30 ++++-- cluster/tx.go | 8 +- shard/doc.go | 15 +++ shard/errors.go | 53 +++++++++++ shard/foreach.go | 119 +++++++++++++++++++++++ shard/group.go | 87 +++++++++++++++++ shard/resolver.go | 8 ++ shard/resolver/custom.go | 60 ++++++++++++ shard/resolver/doc.go | 8 ++ shard/resolver/encoder.go | 91 ++++++++++++++++++ shard/resolver/hash.go | 177 +++++++++++++++++++++++++++++++++++ shard/resolver/range.go | 151 ++++++++++++++++++++++++++++++ shard/resolver/time_range.go | 158 +++++++++++++++++++++++++++++++ shard/resolver/validation.go | 27 ++++++ shard/shard.go | 97 +++++++++++++++++++ shard/topology.go | 132 ++++++++++++++++++++++++++ 19 files changed, 1296 insertions(+), 24 deletions(-) create mode 100644 shard/doc.go create mode 100644 shard/errors.go create mode 100644 shard/foreach.go create mode 100644 shard/group.go create mode 100644 shard/resolver.go create mode 100644 shard/resolver/custom.go create mode 100644 shard/resolver/doc.go create mode 100644 shard/resolver/encoder.go create mode 100644 shard/resolver/hash.go create mode 100644 shard/resolver/range.go create mode 100644 shard/resolver/time_range.go create mode 100644 shard/resolver/validation.go create mode 100644 shard/shard.go create mode 100644 shard/topology.go diff --git a/cluster/cluster.go b/cluster/cluster.go index 153d00c..f7b9911 100644 --- a/cluster/cluster.go +++ b/cluster/cluster.go @@ -9,26 +9,37 @@ import ( "github.com/mkbeh/xpg" ) +// ID identifies one logical PostgreSQL cluster. +type ID string + // Config configures a Cluster from independently created pools. // -// New takes ownership of Primary and Replicas only after it returns -// successfully. Cluster.Close closes the owned pools. +// ID and Labels are optional cluster metadata. New takes ownership of Primary +// and Replicas only after it returns successfully. Cluster.Close closes the +// owned pools. type Config struct { + ID ID + Labels map[string]string + Primary *xpg.Pool Replicas []*xpg.Pool Selector ReplicaSelector } -// Cluster routes operations between one primary pool and optional replica -// pools. +// Cluster routes operations across an optional primary pool and zero or more +// replica pools. // // A deployment with one PostgreSQL endpoint is represented by a Cluster with -// one Primary and no Replicas. +// one Primary and no Replicas. A read-only deployment may omit Primary and +// configure only Replicas. // // Cluster does not inspect SQL, retry failed queries, promote replicas, or // discover PostgreSQL nodes. Those responsibilities remain with the caller or // the surrounding high-availability infrastructure. type Cluster struct { + id ID + labels map[string]string + primary *xpg.Pool replicas []*xpg.Pool @@ -40,8 +51,8 @@ type Cluster struct { // New creates a Cluster from independently configured pools. // -// Primary is required. Replicas may be omitted. When Selector is nil, -// replicas are selected using round-robin. +// At least one pool is required. When Selector is nil, replicas are selected +// using round-robin. func New(config Config) (*Cluster, error) { if config.Primary != nil && invalidPool(config.Primary) { return nil, errors.New("xpg/cluster: primary pool is invalid") @@ -51,6 +62,10 @@ func New(config Config) (*Cluster, error) { return nil, errors.New("xpg/cluster: at least one pool is required") } + if err := validateLabels(config.Labels); err != nil { + return nil, fmt.Errorf("xpg/cluster: %w", err) + } + replicas := slices.Clone(config.Replicas) metadata := make(replicaMetadata, len(replicas)) @@ -71,6 +86,8 @@ func New(config Config) (*Cluster, error) { } return &Cluster{ + id: config.ID, + labels: cloneLabels(config.Labels), primary: config.Primary, replicas: replicas, metadata: metadata, @@ -82,9 +99,53 @@ func invalidPool(pool *xpg.Pool) bool { return pool == nil || pool.Raw() == nil } +func validateLabels(labels map[string]string) error { + for key, value := range labels { + if key == "" { + return errors.New("label key must not be empty") + } + + if value == "" { + return fmt.Errorf("label %q value must not be empty", key) + } + } + + return nil +} + +// ID returns the stable logical cluster ID. +func (c *Cluster) ID() ID { + if c == nil { + return "" + } + + return c.id +} + +// Label returns one cluster label without allocating a copy of all labels. +func (c *Cluster) Label(key string) (string, bool) { + if c == nil { + return "", false + } + + value, ok := c.labels[key] + + return value, ok +} + +// Labels returns a defensive copy of cluster labels. +func (c *Cluster) Labels() map[string]string { + if c == nil { + return nil + } + + return cloneLabels(c.labels) +} + // Primary returns the primary pool owned by the cluster. // -// The returned pool is borrowed and must not be closed separately. +// Primary returns nil when no primary is configured. The returned pool is +// borrowed and must not be closed separately. func (c *Cluster) Primary() *xpg.Pool { if c == nil { return nil @@ -112,7 +173,7 @@ func (c *Cluster) ReplicaAt(index int) *xpg.Pool { } // Close closes all replica pools in reverse registration order and then closes -// the primary pool. +// the primary pool when one is configured. // // Close is safe to call multiple times. func (c *Cluster) Close() { @@ -125,6 +186,8 @@ func (c *Cluster) Close() { c.replicas[index].Close() } - c.primary.Close() + if c.primary != nil { + c.primary.Close() + } }) } diff --git a/cluster/doc.go b/cluster/doc.go index 609a36b..1286c8b 100644 --- a/cluster/doc.go +++ b/cluster/doc.go @@ -1,5 +1,5 @@ -// Package cluster provides explicit routing between one PostgreSQL primary -// pool and zero or more replica pools. +// Package cluster provides explicit routing between a PostgreSQL primary pool, +// when configured, and zero or more replica pools. // // The package does not inspect SQL, retry failed queries, promote replicas, // or discover PostgreSQL nodes. diff --git a/cluster/errors.go b/cluster/errors.go index b205250..8267d4a 100644 --- a/cluster/errors.go +++ b/cluster/errors.go @@ -2,6 +2,12 @@ package cluster import "errors" -// ErrNoReplica indicates that a read policy required a replica but no replica -// was available for selection. -var ErrNoReplica = errors.New("xpg/cluster: no replica available") +var ( + // ErrNoPrimary is returned when an operation requires a primary pool but + // the cluster has no primary configured. + ErrNoPrimary = errors.New("xpg/cluster: no primary available") + + // ErrNoReplica is returned when an operation requires a replica but no + // replica can be selected. + ErrNoReplica = errors.New("xpg/cluster: no replica available") +) diff --git a/cluster/resolver.go b/cluster/resolver.go index e3514a7..024b192 100644 --- a/cluster/resolver.go +++ b/cluster/resolver.go @@ -63,30 +63,46 @@ func (policy ReadPolicy) String() string { // ReadPool returns a pool for a read operation according to policy. // -// ReadReplicaPreferred falls back to the primary only when replica selection -// returns ErrNoReplica. Other selector errors are returned to the caller. +// ReadReplicaPreferred falls back to the primary when no replica can be +// selected. If the cluster has no primary, it returns ErrNoPrimary. +// Other selector errors are returned to the caller. func (c *Cluster) ReadPool(ctx context.Context, policy ReadPolicy) (*xpg.Pool, error) { - if c == nil || c.primary == nil { + if c == nil { return nil, errors.New("xpg/cluster: cluster is nil") } switch policy { case ReadPrimary: + if c.primary == nil { + return nil, ErrNoPrimary + } + return c.primary, nil case ReadReplicaPreferred: replica, err := c.selectReplica(ctx) - if errors.Is(err, ErrNoReplica) { - return c.primary, nil + if err == nil { + return replica, nil } - return replica, err + if !errors.Is(err, ErrNoReplica) { + return nil, err + } + + if c.primary == nil { + return nil, ErrNoPrimary + } + + return c.primary, nil case ReadReplicaRequired: return c.selectReplica(ctx) default: - return nil, fmt.Errorf("xpg/cluster: unsupported read policy %d", policy) + return nil, fmt.Errorf( + "xpg/cluster: unsupported read policy %d", + policy, + ) } } diff --git a/cluster/tx.go b/cluster/tx.go index 9d5a430..6a925fa 100644 --- a/cluster/tx.go +++ b/cluster/tx.go @@ -22,8 +22,12 @@ func (c *Cluster) InPrimaryTx( options pgx.TxOptions, fn func(context.Context, pgx.Tx) error, ) error { - if c == nil || c.primary == nil { - return errors.New("xpg/cluster: cluster is not initialized") + if c == nil { + return errors.New("xpg/cluster: cluster is nil") + } + + if c.primary == nil { + return ErrNoPrimary } return c.primary.InTx(ctx, options, fn) diff --git a/shard/doc.go b/shard/doc.go new file mode 100644 index 0000000..55fad02 --- /dev/null +++ b/shard/doc.go @@ -0,0 +1,15 @@ +// Package shard provides explicit application-level routing across PostgreSQL +// clusters. +// +// A Topology owns an ordered set of logical shards. Every Shard is backed by a +// cluster.Cluster, and typed resolvers map application keys directly to shards. +// Built-in resolvers support rendezvous hashing, ordered ranges, and time +// ranges; applications may also provide custom routing logic. Resolvers borrow +// their topology and do not own its clusters. +// +// The package also provides shard grouping, colocation checks, bounded fan-out, +// and connection-budget diagnostics. +// +// The package does not inspect SQL, hide shard keys in contexts, move data, +// replicate reference tables, or provide distributed transactions. +package shard diff --git a/shard/errors.go b/shard/errors.go new file mode 100644 index 0000000..fcf0411 --- /dev/null +++ b/shard/errors.go @@ -0,0 +1,53 @@ +package shard + +import ( + "errors" + "fmt" +) + +var ( + // ErrNoShard indicates that a resolver could not map a key to any shard. + ErrNoShard = errors.New("xpg/shard: no shard resolved") + + // ErrUnknownShard indicates that routing configuration or custom routing + // logic referenced a shard that does not exist in the topology. + ErrUnknownShard = errors.New("xpg/shard: unknown shard") + + // ErrShardMismatch indicates that keys expected to be colocated resolved to + // different shards. + ErrShardMismatch = errors.New("xpg/shard: keys resolve to different shards") +) + +// UnknownShardError identifies a shard that does not exist in a topology. +type UnknownShardError struct { + ShardID ID +} + +func (e *UnknownShardError) Error() string { + return fmt.Sprintf("xpg/shard: unknown shard %q", e.ShardID) +} + +func (e *UnknownShardError) Unwrap() error { + return ErrUnknownShard +} + +// MismatchError describes the first key that resolved to a different shard +// than the first key. +type MismatchError struct { + Expected ID + Actual ID + Index int +} + +func (e *MismatchError) Error() string { + return fmt.Sprintf( + "xpg/shard: key %d resolved to shard %q instead of %q", + e.Index, + e.Actual, + e.Expected, + ) +} + +func (e *MismatchError) Unwrap() error { + return ErrShardMismatch +} diff --git a/shard/foreach.go b/shard/foreach.go new file mode 100644 index 0000000..41cc265 --- /dev/null +++ b/shard/foreach.go @@ -0,0 +1,119 @@ +package shard + +import ( + "context" + "errors" + "fmt" + "sync" +) + +// ForEachShardResult contains the result of one shard callback invocation. +type ForEachShardResult struct { + ShardID ID + Err error +} + +// ForEachShardResults contains results in topology registration order. +type ForEachShardResults []ForEachShardResult + +// Err returns all shard failures joined in registration order. +func (results ForEachShardResults) Err() error { + errs := make([]error, 0, len(results)) + + for _, result := range results { + if result.Err == nil { + continue + } + + errs = append( + errs, + fmt.Errorf( + "xpg/shard: shard %q callback: %w", + result.ShardID, + result.Err, + ), + ) + } + + return errors.Join(errs...) +} + +// ForEachShard invokes fn for each shard with at most concurrency callbacks +// running at once. Results are returned in topology registration order. +// +// Once context cancellation is observed, callbacks that have not started are +// skipped and their results contain ctx.Err(). Callbacks already running are +// responsible for observing ctx. +func (t *Topology) ForEachShard( + ctx context.Context, + concurrency int, + fn func(context.Context, Shard) error, +) (ForEachShardResults, error) { + if t == nil || len(t.shards) == 0 { + return nil, errors.New("xpg/shard: topology is nil or empty") + } + + if concurrency <= 0 { + return nil, errors.New("xpg/shard: concurrency must be positive") + } + + if fn == nil { + return nil, errors.New("xpg/shard: callback is nil") + } + + results := make(ForEachShardResults, len(t.shards)) + + for index, shard := range t.shards { + results[index].ShardID = shard.ID() + } + + workerCount := min(concurrency, len(t.shards)) + jobs := make(chan int) + + var workers sync.WaitGroup + + for range workerCount { + workers.Go(func() { + for index := range jobs { + if err := ctx.Err(); err != nil { + results[index].Err = err + continue + } + + results[index].Err = fn(ctx, t.shards[index]) + } + }) + } + + nextIndex := 0 + +schedule: + for nextIndex < len(t.shards) { + // Check cancellation before entering select so that a ready worker does + // not repeatedly win against an already canceled context. + if ctx.Err() != nil { + break + } + + select { + case jobs <- nextIndex: + nextIndex++ + + case <-ctx.Done(): + break schedule + } + } + + close(jobs) + workers.Wait() + + if err := ctx.Err(); err != nil { + // Every index before nextIndex was handed to exactly one worker. + // Everything from nextIndex onward was never scheduled. + for index := nextIndex; index < len(results); index++ { + results[index].Err = err + } + } + + return results, nil +} diff --git a/shard/group.go b/shard/group.go new file mode 100644 index 0000000..870ad17 --- /dev/null +++ b/shard/group.go @@ -0,0 +1,87 @@ +package shard + +import ( + "errors" + "fmt" +) + +// SameShard resolves the keys and verifies that they all belong to the same +// shard. It returns the resolved shard when all keys are colocated. +func SameShard[K any](resolver Resolver[K], keys ...K) (Shard, error) { + if resolver == nil { + return Shard{}, errors.New("xpg/shard: resolver is nil") + } + + if len(keys) == 0 { + return Shard{}, ErrNoShard + } + + expected, err := resolver.Resolve(keys[0]) + if err != nil { + return Shard{}, fmt.Errorf("xpg/shard: resolve key 0: %w", err) + } + + expectedID := expected.ID() + + for index := 1; index < len(keys); index++ { + actual, resolveErr := resolver.Resolve(keys[index]) + if resolveErr != nil { + return Shard{}, fmt.Errorf("xpg/shard: resolve key %d: %w", index, resolveErr) + } + + actualID := actual.ID() + if actualID != expectedID { + return Shard{}, &MismatchError{ + Expected: expectedID, + Actual: actualID, + Index: index, + } + } + } + + return expected, nil +} + +// Group contains input keys that resolve to one shard. Keys preserve their +// original relative order. +type Group[K any] struct { + Shard Shard + Keys []K +} + +// GroupByShard resolves every key once and returns groups in order of each +// shard's first appearance in the input. +func GroupByShard[K any](resolver Resolver[K], keys []K) ([]Group[K], error) { + if resolver == nil { + return nil, errors.New("xpg/shard: resolver is nil") + } + + groups := make([]Group[K], 0) + indexByID := make(map[ID]int) + + for keyIndex, key := range keys { + resolved, err := resolver.Resolve(key) + if err != nil { + return nil, fmt.Errorf("xpg/shard: resolve key %d: %w", keyIndex, err) + } + + id := resolved.ID() + + groupIndex, exists := indexByID[id] + if !exists { + groupIndex = len(groups) + indexByID[id] = groupIndex + + groups = append( + groups, + Group[K]{ + Shard: resolved, + }, + ) + } + + groups[groupIndex].Keys = append(groups[groupIndex].Keys, key) + } + + return groups, nil +} diff --git a/shard/resolver.go b/shard/resolver.go new file mode 100644 index 0000000..67f287a --- /dev/null +++ b/shard/resolver.go @@ -0,0 +1,8 @@ +package shard + +// Resolver maps a typed application key to one shard. +// +// Implementations shared by concurrent callers must be concurrency-safe. +type Resolver[K any] interface { + Resolve(key K) (Shard, error) +} diff --git a/shard/resolver/custom.go b/shard/resolver/custom.go new file mode 100644 index 0000000..7c07c15 --- /dev/null +++ b/shard/resolver/custom.go @@ -0,0 +1,60 @@ +package resolver + +import ( + "errors" + + "github.com/mkbeh/xpg/shard" +) + +// ResolveFunc maps an application key to a shard ID within topology. +// +// Implementations shared by concurrent callers must be deterministic and +// concurrency-safe. They should not perform hidden I/O or modify topology. +type ResolveFunc[K any] func(key K, topology *shard.Topology) (shard.ID, error) + +// CustomResolver adapts ResolveFunc to shard.Resolver. +type CustomResolver[K any] struct { + topology *shard.Topology + resolve ResolveFunc[K] +} + +// NewCustom binds custom routing logic to one immutable topology. +func NewCustom[K any](topology *shard.Topology, resolve ResolveFunc[K]) (*CustomResolver[K], error) { + if err := requireTopology(topology); err != nil { + return nil, err + } + + if resolve == nil { + return nil, errors.New("xpg/shard/resolver: custom resolve function is nil") + } + + return &CustomResolver[K]{ + topology: topology, + resolve: resolve, + }, nil +} + +// Resolve maps key to a shard and rejects IDs absent from the bound topology. +func (resolver *CustomResolver[K]) Resolve(key K) (shard.Shard, error) { + if resolver == nil || + resolver.topology == nil || + resolver.resolve == nil { + return shard.Shard{}, errors.New( + "xpg/shard/resolver: custom resolver is not initialized", + ) + } + + id, err := resolver.resolve(key, resolver.topology) + if err != nil { + return shard.Shard{}, err + } + + resolved, ok := resolver.topology.Shard(id) + if !ok { + return shard.Shard{}, &shard.UnknownShardError{ + ShardID: id, + } + } + + return resolved, nil +} diff --git a/shard/resolver/doc.go b/shard/resolver/doc.go new file mode 100644 index 0000000..047e8ca --- /dev/null +++ b/shard/resolver/doc.go @@ -0,0 +1,8 @@ +// Package resolver provides routing strategies for shard.Topology. +// +// Resolvers are bound to an immutable topology and return shard.Shard values. +// The package provides rendezvous hashing, ordered numeric or string ranges, +// time ranges, and an adapter for custom routing functions. +// +// Resolvers borrow their topology and must not outlive it. +package resolver diff --git a/shard/resolver/encoder.go b/shard/resolver/encoder.go new file mode 100644 index 0000000..a0a1c3e --- /dev/null +++ b/shard/resolver/encoder.go @@ -0,0 +1,91 @@ +package resolver + +import ( + "bytes" + "encoding/binary" + "errors" +) + +// KeyEncoder converts a typed key into stable canonical bytes. +// +// Implementations used for persistent shard placement must remain +// deterministic across processes and releases. Changing an encoder changes +// hash placement and may require data migration. +type KeyEncoder[K any] interface { + Encode(K) ([]byte, error) +} + +// KeyEncoderFunc adapts a function to KeyEncoder. +type KeyEncoderFunc[K any] func(K) ([]byte, error) + +// Encode encodes key using the adapted function. +func (encoder KeyEncoderFunc[K]) Encode(key K) ([]byte, error) { + if encoder == nil { + return nil, errors.New("xpg/shard/resolver: key encoder function is nil") + } + + return encoder(key) +} + +// StringKeyEncoder encodes the exact bytes stored in a string without +// normalization. +func StringKeyEncoder() KeyEncoder[string] { + return KeyEncoderFunc[string]( + func(key string) ([]byte, error) { + return []byte(key), nil + }, + ) +} + +// BytesKeyEncoder encodes exact bytes and returns a defensive copy. +func BytesKeyEncoder() KeyEncoder[[]byte] { + return KeyEncoderFunc[[]byte]( + func(key []byte) ([]byte, error) { + return bytes.Clone(key), nil + }, + ) +} + +// Int64KeyEncoder encodes a signed integer as big-endian two's-complement. +func Int64KeyEncoder() KeyEncoder[int64] { + return KeyEncoderFunc[int64]( + func(key int64) ([]byte, error) { + encoded := make([]byte, 8) + + binary.BigEndian.PutUint64(encoded, uint64(key)) + + return encoded, nil + }, + ) +} + +// Uint64KeyEncoder encodes an unsigned integer as big-endian bytes. +func Uint64KeyEncoder() KeyEncoder[uint64] { + return KeyEncoderFunc[uint64]( + func(key uint64) ([]byte, error) { + encoded := make([]byte, 8) + + binary.BigEndian.PutUint64(encoded, key) + + return encoded, nil + }, + ) +} + +// Bytes16KeyEncoder encodes a 16-byte key exactly. +func Bytes16KeyEncoder() KeyEncoder[[16]byte] { + return KeyEncoderFunc[[16]byte]( + func(key [16]byte) ([]byte, error) { + return bytes.Clone(key[:]), nil + }, + ) +} + +// Bytes32KeyEncoder encodes a 32-byte key exactly. +func Bytes32KeyEncoder() KeyEncoder[[32]byte] { + return KeyEncoderFunc[[32]byte]( + func(key [32]byte) ([]byte, error) { + return bytes.Clone(key[:]), nil + }, + ) +} diff --git a/shard/resolver/hash.go b/shard/resolver/hash.go new file mode 100644 index 0000000..6bb3e70 --- /dev/null +++ b/shard/resolver/hash.go @@ -0,0 +1,177 @@ +package resolver + +import ( + "bytes" + "crypto/sha256" + "encoding/binary" + "errors" + "fmt" + "math" + "strings" + + "github.com/mkbeh/xpg/shard" +) + +const ( + // rendezvousDomain identifies the persistent hash-placement algorithm. + // Changing it changes shard placement and therefore requires data migration. + rendezvousDomain = "xpg.shard.rendezvous.v1" + + // rendezvousLengthSize is the size of each uint32 length prefix in bytes. + rendezvousLengthSize = 4 +) + +// HashResolver implements rendezvous/HRW routing with SHA-256 and stable named +// shard IDs. +type HashResolver[K any] struct { + shards []shard.Shard + prefix []byte + encoder KeyEncoder[K] + maxIDLength int +} + +// NewHash creates the version-1 rendezvous resolver bound to topology. +// +// Namespace is part of the persistent placement contract. Changing it changes +// shard placement and may require data migration. +func NewHash[K any]( + topology *shard.Topology, + namespace string, + encoder KeyEncoder[K], +) (*HashResolver[K], error) { + if err := requireTopology(topology); err != nil { + return nil, err + } + + if encoder == nil { + return nil, errors.New("xpg/shard/resolver: key encoder is nil") + } + + trimmedNamespace := strings.TrimSpace(namespace) + + if trimmedNamespace == "" { + return nil, errors.New("xpg/shard/resolver: hash namespace must not be blank") + } + + if trimmedNamespace != namespace { + return nil, errors.New("xpg/shard/resolver: hash namespace must not contain surrounding whitespace") + } + + if len(namespace) > math.MaxUint32 { + return nil, errors.New("xpg/shard/resolver: hash namespace is too large") + } + + shards := topology.Shards() + maxIDLength := 0 + + for _, candidate := range shards { + id := candidate.ID() + + if len(id) > math.MaxUint32 { + return nil, errors.New("xpg/shard/resolver: shard ID is too large") + } + + maxIDLength = max(maxIDLength, len(id)) + } + + // Prefix is invariant for the lifetime of the resolver: + // + // domain || namespace_length || namespace + prefix := make([]byte, len(rendezvousDomain)+rendezvousLengthSize+len(namespace)) + + offset := copy(prefix, rendezvousDomain) + + binary.BigEndian.PutUint32( + prefix[offset:offset+rendezvousLengthSize], + uint32(len(namespace)), + ) + offset += rendezvousLengthSize + + copy(prefix[offset:], namespace) + + return &HashResolver[K]{ + shards: shards, + prefix: prefix, + encoder: encoder, + maxIDLength: maxIDLength, + }, nil +} + +// Resolve selects the shard with the lexicographically greatest SHA-256 score. +// +// Resolve performs only in-memory routing. It does not acquire a connection or +// execute a PostgreSQL query. +func (resolver *HashResolver[K]) Resolve(key K) (shard.Shard, error) { + if resolver == nil || + len(resolver.shards) == 0 || + resolver.encoder == nil { + return shard.Shard{}, errors.New("xpg/shard/resolver: hash resolver is not initialized") + } + + encoded, err := resolver.encoder.Encode(key) + if err != nil { + return shard.Shard{}, fmt.Errorf("xpg/shard/resolver: encode hash key: %w", err) + } + + if len(encoded) > math.MaxUint32 { + return shard.Shard{}, errors.New("xpg/shard/resolver: encoded key is too large") + } + + // Build the candidate-independent prefix once. The shard ID suffix is + // overwritten for each candidate below. + keyLengthOffset := len(resolver.prefix) + keyOffset := keyLengthOffset + rendezvousLengthSize + idLengthOffset := keyOffset + len(encoded) + idOffset := idLengthOffset + rendezvousLengthSize + + scoreInput := make([]byte, idOffset+resolver.maxIDLength) + + copy(scoreInput, resolver.prefix) + + binary.BigEndian.PutUint32( + scoreInput[keyLengthOffset:keyOffset], + uint32(len(encoded)), + ) + + copy(scoreInput[keyOffset:idLengthOffset], encoded) + + var ( + selected shard.Shard + best [sha256.Size]byte + bestID shard.ID + hasBest bool + ) + + for _, candidate := range resolver.shards { + candidateID := candidate.ID() + inputEnd := idOffset + len(candidateID) + + binary.BigEndian.PutUint32( + scoreInput[idLengthOffset:idOffset], + uint32(len(candidateID)), + ) + + copy(scoreInput[idOffset:inputEnd], candidateID) + + score := sha256.Sum256(scoreInput[:inputEnd]) + + comparison := bytes.Compare(score[:], best[:]) + + // Shard ID is the deterministic tie-breaker, so placement does not + // depend on topology registration order when scores are equal. + if !hasBest || + comparison > 0 || + comparison == 0 && candidateID < bestID { + selected = candidate + best = score + bestID = candidateID + hasBest = true + } + } + + if !hasBest { + return shard.Shard{}, shard.ErrNoShard + } + + return selected, nil +} diff --git a/shard/resolver/range.go b/shard/resolver/range.go new file mode 100644 index 0000000..70f0aa2 --- /dev/null +++ b/shard/resolver/range.go @@ -0,0 +1,151 @@ +package resolver + +import ( + "cmp" + "errors" + "fmt" + "slices" + "sort" + + "github.com/mkbeh/xpg/shard" +) + +// Range maps the bounded half-open interval [Start, End) to one shard. +// +// Start must be less than End. Ranges may be supplied in any order. +// Gaps are allowed and resolve to shard.ErrNoShard. +type Range[K cmp.Ordered] struct { + Start K + End K + ShardID shard.ID +} + +// RangeResolver resolves ordered keys through bounded, non-overlapping ranges. +type RangeResolver[K cmp.Ordered] struct { + ranges []rangeEntry[K] +} + +type rangeEntry[K cmp.Ordered] struct { + start K + end K + shard shard.Shard + + sourceIndex int +} + +// NewRange creates a resolver from bounded, non-overlapping half-open ranges. +// +// NewRange copies the supplied ranges into an internal representation, sorts +// them by Start, and validates that they do not overlap. The caller's slice is +// not modified. +func NewRange[K cmp.Ordered](topology *shard.Topology, ranges []Range[K]) (*RangeResolver[K], error) { + if err := requireTopology(topology); err != nil { + return nil, err + } + + if len(ranges) == 0 { + return nil, errors.New("xpg/shard/resolver: range resolver requires at least one range") + } + + entries := make([]rangeEntry[K], len(ranges)) + + for index, valueRange := range ranges { + if err := requireShardID(valueRange.ShardID); err != nil { + return nil, fmt.Errorf("xpg/shard/resolver: range %d: %w", index, err) + } + + // This rejects empty, reversed, and NaN-bounded ranges. + valid := valueRange.Start < valueRange.End + if !valid { + return nil, fmt.Errorf("xpg/shard/resolver: range %d must satisfy start < end", index) + } + + resolved, ok := topology.Shard(valueRange.ShardID) + if !ok { + return nil, fmt.Errorf( + "xpg/shard/resolver: range %d: %w", + index, + &shard.UnknownShardError{ + ShardID: valueRange.ShardID, + }, + ) + } + + entries[index] = rangeEntry[K]{ + start: valueRange.Start, + end: valueRange.End, + shard: resolved, + sourceIndex: index, + } + } + + // sourceIndex provides deterministic ordering for equal starts and keeps + // overlap errors tied to the caller's original slice. + slices.SortFunc( + entries, + func(left, right rangeEntry[K]) int { + if order := cmp.Compare(left.start, right.start); order != 0 { + return order + } + + return cmp.Compare(left.sourceIndex, right.sourceIndex) + }, + ) + + // Once ranges are sorted by Start, checking adjacent entries is sufficient + // to detect every overlap. + for index := 1; index < len(entries); index++ { + previous := entries[index-1] + current := entries[index] + + // Adjacent half-open ranges are valid: + // + // [0, 100) and [100, 200) + if previous.end <= current.start { + continue + } + + return nil, fmt.Errorf( + "xpg/shard/resolver: ranges %d and %d overlap", + previous.sourceIndex, + current.sourceIndex, + ) + } + + return &RangeResolver[K]{ + ranges: entries, + }, nil +} + +// Resolve returns the shard whose configured range contains key. +// +// Resolve performs only an in-memory lookup. It does not acquire a connection +// or execute a PostgreSQL query. +func (resolver *RangeResolver[K]) Resolve(key K) (shard.Shard, error) { + if resolver == nil || len(resolver.ranges) == 0 { + return shard.Shard{}, shard.ErrNoShard + } + + // Non-overlap validation guarantees strictly increasing upper boundaries, + // making this search predicate monotonic. + index := sort.Search( + len(resolver.ranges), + func(index int) bool { + return key < resolver.ranges[index].end + }, + ) + + if index == len(resolver.ranges) { + return shard.Shard{}, shard.ErrNoShard + } + + entry := resolver.ranges[index] + + // The first range ending after key may still start after it when there is a + // gap between configured ranges. + if key < entry.start { + return shard.Shard{}, shard.ErrNoShard + } + + return entry.shard, nil +} diff --git a/shard/resolver/time_range.go b/shard/resolver/time_range.go new file mode 100644 index 0000000..79a54a2 --- /dev/null +++ b/shard/resolver/time_range.go @@ -0,0 +1,158 @@ +package resolver + +import ( + "cmp" + "errors" + "fmt" + "slices" + "sort" + "time" + + "github.com/mkbeh/xpg/shard" +) + +// TimeRange maps the bounded half-open interval [Start, End) to one shard. +// +// Start must be before End. Ranges may be supplied in any order. +// Gaps are allowed and resolve to shard.ErrNoShard. +type TimeRange struct { + Start time.Time + End time.Time + ShardID shard.ID +} + +// TimeRangeResolver resolves time instants through bounded, non-overlapping +// ranges. +type TimeRangeResolver struct { + ranges []timeRangeEntry +} + +type timeRangeEntry struct { + start time.Time + end time.Time + shard shard.Shard + + sourceIndex int +} + +// NewTimeRange creates a resolver from bounded, non-overlapping half-open time +// ranges. +// +// Range boundaries are normalized to UTC. The caller's slice and time values +// are not modified. +func NewTimeRange(topology *shard.Topology, ranges []TimeRange) (*TimeRangeResolver, error) { + if err := requireTopology(topology); err != nil { + return nil, err + } + + if len(ranges) == 0 { + return nil, errors.New("xpg/shard/resolver: time range resolver requires at least one range") + } + + entries := make([]timeRangeEntry, len(ranges)) + + for index, valueRange := range ranges { + if err := requireShardID(valueRange.ShardID); err != nil { + return nil, fmt.Errorf("xpg/shard/resolver: time range %d: %w", index, err) + } + + start := normalizeTime(valueRange.Start) + end := normalizeTime(valueRange.End) + + if !start.Before(end) { + return nil, fmt.Errorf("xpg/shard/resolver: time range %d must satisfy start < end", index) + } + + resolved, ok := topology.Shard(valueRange.ShardID) + if !ok { + return nil, fmt.Errorf( + "xpg/shard/resolver: time range %d: %w", + index, + &shard.UnknownShardError{ + ShardID: valueRange.ShardID, + }, + ) + } + + entries[index] = timeRangeEntry{ + start: start, + end: end, + shard: resolved, + sourceIndex: index, + } + } + + // sourceIndex keeps overlap diagnostics tied to the caller's original + // slice and provides deterministic ordering for equal starts. + slices.SortFunc( + entries, + func(left, right timeRangeEntry) int { + if order := left.start.Compare(right.start); order != 0 { + return order + } + + return cmp.Compare(left.sourceIndex, right.sourceIndex) + }, + ) + + // Once ranges are sorted by Start, checking adjacent entries is sufficient + // to detect every overlap. + for index := 1; index < len(entries); index++ { + previous := entries[index-1] + current := entries[index] + + // Adjacent half-open ranges are valid: + // + // [00:00, 01:00) and [01:00, 02:00) + if previous.end.After(current.start) { + return nil, fmt.Errorf( + "xpg/shard/resolver: time ranges %d and %d overlap", + previous.sourceIndex, + current.sourceIndex, + ) + } + } + + return &TimeRangeResolver{ + ranges: entries, + }, nil +} + +// Resolve returns the shard whose configured time range contains key. +// +// Resolve performs only an in-memory lookup. It does not acquire a connection +// or execute a PostgreSQL query. +func (resolver *TimeRangeResolver) Resolve(key time.Time) (shard.Shard, error) { + if resolver == nil || len(resolver.ranges) == 0 { + return shard.Shard{}, shard.ErrNoShard + } + + key = normalizeTime(key) + + // Non-overlap validation guarantees strictly increasing upper boundaries, + // making this search predicate monotonic. + index := sort.Search( + len(resolver.ranges), + func(index int) bool { + return key.Before(resolver.ranges[index].end) + }, + ) + + if index == len(resolver.ranges) { + return shard.Shard{}, shard.ErrNoShard + } + + entry := resolver.ranges[index] + + // The first range ending after key may still start after it when there is a + // gap between configured ranges. + if key.Before(entry.start) { + return shard.Shard{}, shard.ErrNoShard + } + + return entry.shard, nil +} + +func normalizeTime(value time.Time) time.Time { + return value.UTC() +} diff --git a/shard/resolver/validation.go b/shard/resolver/validation.go new file mode 100644 index 0000000..dce5f31 --- /dev/null +++ b/shard/resolver/validation.go @@ -0,0 +1,27 @@ +package resolver + +import ( + "errors" + + "github.com/mkbeh/xpg/shard" +) + +func requireTopology(topology *shard.Topology) error { + if topology == nil || topology.Len() == 0 { + return errors.New( + "xpg/shard/resolver: topology is nil or empty", + ) + } + + return nil +} + +func requireShardID(id shard.ID) error { + if id == "" { + return errors.New( + "shard ID must not be empty", + ) + } + + return nil +} diff --git a/shard/shard.go b/shard/shard.go new file mode 100644 index 0000000..3394e5b --- /dev/null +++ b/shard/shard.go @@ -0,0 +1,97 @@ +package shard + +import ( + "context" + + "github.com/jackc/pgx/v5" + "github.com/mkbeh/xpg" + "github.com/mkbeh/xpg/cluster" +) + +// ID identifies one logical shard. +type ID = cluster.ID + +// Shard is a borrowed handle to one cluster registered in a Topology. +// +// Shard exposes shard-local operations without exposing cluster lifecycle or +// replica-set management. +type Shard struct { + cluster *cluster.Cluster +} + +// ID returns the stable logical shard ID. +func (s Shard) ID() ID { + if s.cluster == nil { + return "" + } + + return s.cluster.ID() +} + +// Label returns one shard label without allocating a copy of all labels. +func (s Shard) Label(key string) (string, bool) { + if s.cluster == nil { + return "", false + } + + return s.cluster.Label(key) +} + +// Labels returns a defensive copy of shard labels. +func (s Shard) Labels() map[string]string { + if s.cluster == nil { + return nil + } + + return s.cluster.Labels() +} + +// Primary returns the shard primary pool, or nil when the shard cluster has no +// primary configured. The returned pool is borrowed and remains owned by the +// shard cluster. +func (s Shard) Primary() *xpg.Pool { + if s.cluster == nil { + return nil + } + + return s.cluster.Primary() +} + +// ReadPool returns a borrowed pool for a read operation according to policy. +func (s Shard) ReadPool( + ctx context.Context, + policy cluster.ReadPolicy, +) (*xpg.Pool, error) { + if s.cluster == nil { + return nil, ErrNoShard + } + + return s.cluster.ReadPool(ctx, policy) +} + +// InPrimaryTx executes fn in a transaction on the shard primary. +func (s Shard) InPrimaryTx( + ctx context.Context, + options pgx.TxOptions, + fn func(context.Context, pgx.Tx) error, +) error { + if s.cluster == nil { + return ErrNoShard + } + + return s.cluster.InPrimaryTx(ctx, options, fn) +} + +// InReadTx executes fn in a read-only transaction resolved within this shard. +func (s Shard) InReadTx( + ctx context.Context, + policy cluster.ReadPolicy, + options cluster.ReadTxOptions, + fn func(context.Context, pgx.Tx) error, +) error { + if s.cluster == nil { + return ErrNoShard + } + + return s.cluster.InReadTx(ctx, policy, options, fn) +} diff --git a/shard/topology.go b/shard/topology.go new file mode 100644 index 0000000..d79ec20 --- /dev/null +++ b/shard/topology.go @@ -0,0 +1,132 @@ +package shard + +import ( + "errors" + "fmt" + "slices" + "sync" + + "github.com/mkbeh/xpg/cluster" +) + +// Config registers one Cluster as a logical shard in a Topology. +// +// The shard ID and labels are provided by Cluster. Config is intentionally +// retained as the extension point for future topology-specific options. +type Config struct { + Cluster *cluster.Cluster +} + +// Topology is an immutable ordered set of logical shards and their PostgreSQL +// clusters. +// +// NewTopology takes ownership of all configured clusters only after it returns +// successfully. Close closes every owned cluster exactly once. +type Topology struct { + shards []Shard + indexByID map[ID]int + + closeOnce sync.Once +} + +// NewTopology validates and creates an immutable topology. Shards retain their +// registration order. Every cluster must have a non-empty and unique ID. +func NewTopology(configs []Config) (*Topology, error) { + if len(configs) == 0 { + return nil, errors.New("xpg/shard: topology must contain at least one shard") + } + + shards := make([]Shard, 0, len(configs)) + indexByID := make(map[ID]int, len(configs)) + + for index, config := range configs { + resolved, err := newShard(config) + if err != nil { + return nil, fmt.Errorf("xpg/shard: shard %d: %w", index, err) + } + + id := resolved.ID() + + if _, exists := indexByID[id]; exists { + return nil, fmt.Errorf("xpg/shard: duplicate shard ID %q", id) + } + + indexByID[id] = len(shards) + shards = append(shards, resolved) + } + + return &Topology{ + shards: shards, + indexByID: indexByID, + }, nil +} + +// Len returns the number of registered shards. +func (t *Topology) Len() int { + if t == nil { + return 0 + } + + return len(t.shards) +} + +// At returns the shard at index in registration order. At panics when index is +// outside the topology, matching ordinary slice indexing semantics. +func (t *Topology) At(index int) Shard { + return t.shards[index] +} + +// Shards returns a defensive copy of shards in registration order. +func (t *Topology) Shards() []Shard { + if t == nil { + return nil + } + + return slices.Clone(t.shards) +} + +// Shard returns one shard by stable ID. +func (t *Topology) Shard(id ID) (Shard, bool) { + return t.lookup(id) +} + +func (t *Topology) lookup(id ID) (Shard, bool) { + if t == nil { + return Shard{}, false + } + + index, ok := t.indexByID[id] + if !ok { + return Shard{}, false + } + + return t.shards[index], true +} + +// Close closes owned clusters in reverse registration order. Close is safe to +// call multiple times. +func (t *Topology) Close() { + if t == nil { + return + } + + t.closeOnce.Do(func() { + for index := len(t.shards) - 1; index >= 0; index-- { + t.shards[index].cluster.Close() + } + }) +} + +func newShard(config Config) (Shard, error) { + if config.Cluster == nil { + return Shard{}, errors.New("cluster is nil") + } + + if config.Cluster.ID() == "" { + return Shard{}, errors.New("cluster ID must not be empty") + } + + return Shard{ + cluster: config.Cluster, + }, nil +} From fddb1b54b1d78d691833dd99b542e8b8516a43a5 Mon Sep 17 00:00:00 2001 From: mkbeh Date: Tue, 11 Aug 2026 16:44:41 +0300 Subject: [PATCH 12/41] docs(examples): add sharding example --- examples/README.md | 2 + examples/shard/README.md | 133 +++++++++++++++++++ examples/shard/docker-compose.yml | 68 ++++++++++ examples/shard/go.mod | 8 ++ examples/shard/main.go | 66 ++++++++++ examples/shard/operations.go | 160 +++++++++++++++++++++++ examples/shard/routing.go | 133 +++++++++++++++++++ examples/shard/setup.go | 171 +++++++++++++++++++++++++ examples/shard/sql/shard-a-primary.sql | 2 + examples/shard/sql/shard-b-primary.sql | 2 + examples/shard/sql/shard-b-replica.sql | 5 + examples/shard/sql/shard.sql | 21 +++ 12 files changed, 771 insertions(+) create mode 100644 examples/shard/README.md create mode 100644 examples/shard/docker-compose.yml create mode 100644 examples/shard/go.mod create mode 100644 examples/shard/main.go create mode 100644 examples/shard/operations.go create mode 100644 examples/shard/routing.go create mode 100644 examples/shard/setup.go create mode 100644 examples/shard/sql/shard-a-primary.sql create mode 100644 examples/shard/sql/shard-b-primary.sql create mode 100644 examples/shard/sql/shard-b-replica.sql create mode 100644 examples/shard/sql/shard.sql diff --git a/examples/README.md b/examples/README.md index b6bca32..32c21d1 100644 --- a/examples/README.md +++ b/examples/README.md @@ -8,6 +8,8 @@ This directory contains runnable examples demonstrating the main features and us | [`transactions`](transactions) | Committing an outer transaction after an optional operation is rolled back to a savepoint | | [`advisory`](advisory) | Coordinating concurrent transactions with PostgreSQL advisory locks | | [`otel`](otel) | Exporting pool metrics through OpenTelemetry and Prometheus | +| [`cluster`](cluster) | Routing reads and transactions across primary and replica pools | +| [`sharding`](shard) | Typed routing across immutable standalone and cluster shard targets | ## Running the examples diff --git a/examples/shard/README.md b/examples/shard/README.md new file mode 100644 index 0000000..cd06850 --- /dev/null +++ b/examples/shard/README.md @@ -0,0 +1,133 @@ +# Sharding + +This example routes typed application keys across two logical PostgreSQL shards. Resolvers are bound to an immutable +`shard.Topology` and return a `shard.Shard`, which delegates database operations to its `cluster.Cluster`. + +```text + Topology []Shard + │ +application key ─── Resolver + │ + Shard + │ + Cluster + │ + primary / replicas +``` + +**This example demonstrates:** + +* Building a topology from primary-only and primary/replica clusters +* Routing `uint64` keys through bounded numeric ranges +* Running shard-local primary and read-only replica transactions +* Checking key colocation and grouping keys by shard +* Reading reference-table copies across shards with bounded fan-out + +> [!NOTE] +> The `shard-b` replica is an independent read-only PostgreSQL instance used to +> demonstrate routing. The local setup does not configure streaming replication. + +## Configuration + +The example uses the following connection strings by default: + +```text +XPG_SHARD_A_PRIMARY_DATABASE_URL=postgres://postgres:postgres@localhost:56431/postgres?sslmode=disable&target_session_attrs=read-write +XPG_SHARD_B_PRIMARY_DATABASE_URL=postgres://postgres:postgres@localhost:56432/postgres?sslmode=disable&target_session_attrs=read-write +XPG_SHARD_B_REPLICA_DATABASE_URL=postgres://postgres:postgres@localhost:56433/postgres?sslmode=disable&target_session_attrs=read-only +``` + +Set the corresponding environment variables to use different PostgreSQL endpoints: + +```shell +export XPG_SHARD_A_PRIMARY_DATABASE_URL='postgres://user:password@shard-a.example.com:5432/database?sslmode=disable&target_session_attrs=read-write' +export XPG_SHARD_B_PRIMARY_DATABASE_URL='postgres://user:password@shard-b-primary.example.com:5432/database?sslmode=disable&target_session_attrs=read-write' +export XPG_SHARD_B_REPLICA_DATABASE_URL='postgres://user:password@shard-b-replica.example.com:5432/database?sslmode=disable&target_session_attrs=read-only' +``` + +## Local setup + +Start all PostgreSQL endpoints and Adminer from the repository root: + +```shell +docker compose -f examples/shard/docker-compose.yml --profile tools up -d +``` + +Or from this example directory: + +```shell +docker compose --profile tools up -d +``` + +Services are available at: + +```text +Shard A primary: localhost:56431 +Shard B primary: localhost:56432 +Shard B replica: localhost:56433 +Adminer: http://localhost:58082 +``` + +Sign in to Adminer with: + +```text +System: PostgreSQL +Server: postgres-shard-a-primary +Username: postgres +Password: postgres +Database: postgres +``` + +Use `postgres-shard-b-primary` or `postgres-shard-b-replica` in the **Server** field to inspect another shard endpoint. + +## Run + +From this directory: + +```shell +go run . +``` + +Or from the repository root: + +```shell +go run ./examples/shard +``` + +## Expected output + +```text +range routing and primary transactions: +- user_id=42 shard=shard-a primary_pool=shard.shard-a.primary +- user_id=142 shard=shard-b primary_pool=shard.shard-b.primary + +grouping and colocation: +- shard=shard-b user_ids=[142 143] +- shard=shard-a user_ids=[42 43] +- colocated shard=shard-a +- cross-shard SameShard returns ErrShardMismatch=true + +replica routing and read-only transaction: +- user_id=142 shard=shard-b read_pool=shard.shard-b.replica read_node=shard-b-replica tx_node=shard-b-replica role=replica read_only=on + +reference table copies: +- shard=shard-a countries=2 +- shard=shard-b countries=2 +``` + +## Stop services + +From the repository root: + +```shell +docker compose \ + -f examples/shard/docker-compose.yml \ + --profile tools \ + down --remove-orphans -v +``` + +Or from this example directory: + +```shell +docker compose --profile tools down --remove-orphans -v +``` diff --git a/examples/shard/docker-compose.yml b/examples/shard/docker-compose.yml new file mode 100644 index 0000000..d379c3e --- /dev/null +++ b/examples/shard/docker-compose.yml @@ -0,0 +1,68 @@ +services: + postgres-shard-a-primary: + image: postgres:18-alpine + environment: + POSTGRES_DB: postgres + POSTGRES_USER: postgres + POSTGRES_PASSWORD: postgres + ports: + - "56431:5432" + volumes: + - shard-a-primary-data:/var/lib/postgresql + - ./sql/shard.sql:/docker-entrypoint-initdb.d/10-shard.sql:ro + - ./sql/shard-a-primary.sql:/docker-entrypoint-initdb.d/20-node.sql:ro + healthcheck: &postgres-healthcheck + test: "pg_isready -U $${POSTGRES_USER} -d $${POSTGRES_DB}" + interval: 2s + timeout: 5s + retries: 10 + start_period: 5s + + postgres-shard-b-primary: + image: postgres:18-alpine + environment: + POSTGRES_DB: postgres + POSTGRES_USER: postgres + POSTGRES_PASSWORD: postgres + ports: + - "56432:5432" + volumes: + - shard-b-primary-data:/var/lib/postgresql + - ./sql/shard.sql:/docker-entrypoint-initdb.d/10-shard.sql:ro + - ./sql/shard-b-primary.sql:/docker-entrypoint-initdb.d/20-node.sql:ro + healthcheck: *postgres-healthcheck + + postgres-shard-b-replica: + image: postgres:18-alpine + environment: + POSTGRES_DB: postgres + POSTGRES_USER: postgres + POSTGRES_PASSWORD: postgres + ports: + - "56433:5432" + volumes: + - shard-b-replica-data:/var/lib/postgresql + - ./sql/shard.sql:/docker-entrypoint-initdb.d/10-shard.sql:ro + - ./sql/shard-b-replica.sql:/docker-entrypoint-initdb.d/20-node.sql:ro + healthcheck: *postgres-healthcheck + + adminer: + image: adminer:standalone + profiles: + - tools + environment: + ADMINER_DEFAULT_SERVER: postgres-shard-a-primary + ports: + - "58082:8080" + depends_on: + postgres-shard-a-primary: + condition: service_healthy + postgres-shard-b-primary: + condition: service_healthy + postgres-shard-b-replica: + condition: service_healthy + +volumes: + shard-a-primary-data: + shard-b-primary-data: + shard-b-replica-data: diff --git a/examples/shard/go.mod b/examples/shard/go.mod new file mode 100644 index 0000000..2215122 --- /dev/null +++ b/examples/shard/go.mod @@ -0,0 +1,8 @@ +module github.com/mkbeh/xpg/examples/shard + +go 1.26 + +require ( + github.com/jackc/pgx/v5 v5.10.0 + github.com/mkbeh/xpg v0.2.0 +) diff --git a/examples/shard/main.go b/examples/shard/main.go new file mode 100644 index 0000000..8b7b00a --- /dev/null +++ b/examples/shard/main.go @@ -0,0 +1,66 @@ +package main + +import ( + "context" + "fmt" + "log" + + "github.com/mkbeh/xpg/shard/resolver" +) + +const ( + userIDRangeStart uint64 = 0 + userIDBoundary uint64 = 100 + userIDRangeEnd uint64 = 200 +) + +func main() { + if err := run(context.Background()); err != nil { + log.Fatal(err) + } +} + +func run(ctx context.Context) error { + topology, err := openTopology(ctx) + if err != nil { + return err + } + defer topology.Close() + + userResolver, err := resolver.NewRange( + topology, + []resolver.Range[uint64]{ + { + Start: userIDRangeStart, + End: userIDBoundary, + ShardID: shardAID, + }, + { + Start: userIDBoundary, + End: userIDRangeEnd, + ShardID: shardBID, + }, + }, + ) + if err != nil { + return fmt.Errorf("create user resolver: %w", err) + } + + if err := showRangeRouting(ctx, userResolver); err != nil { + return err + } + + if err := showGrouping(userResolver); err != nil { + return err + } + + if err := showReplicaRead(ctx, userResolver, shardBUserID); err != nil { + return err + } + + if err := showReferenceTables(ctx, topology); err != nil { + return err + } + + return nil +} diff --git a/examples/shard/operations.go b/examples/shard/operations.go new file mode 100644 index 0000000..3e2db85 --- /dev/null +++ b/examples/shard/operations.go @@ -0,0 +1,160 @@ +package main + +import ( + "context" + "fmt" + "sync" + + "github.com/jackc/pgx/v5" + "github.com/mkbeh/xpg/cluster" + "github.com/mkbeh/xpg/shard" +) + +const referenceTableConcurrency = 2 + +type nodeInfo struct { + name string + role string +} + +type rowQuerier interface { + QueryRow(context.Context, string, ...any) pgx.Row +} + +func showReplicaRead( + ctx context.Context, + userResolver shard.Resolver[uint64], + key uint64, +) error { + resolved, err := userResolver.Resolve(key) + if err != nil { + return fmt.Errorf("resolve replica key %d: %w", key, err) + } + + pool, err := resolved.ReadPool( + ctx, + cluster.ReadReplicaRequired, + ) + if err != nil { + return fmt.Errorf("resolve shard replica: %w", err) + } + + node, err := loadNode(ctx, pool) + if err != nil { + return fmt.Errorf("read replica node: %w", err) + } + + var ( + transactionNode nodeInfo + readOnly string + ) + if err := resolved.InReadTx( + ctx, + cluster.ReadReplicaRequired, + cluster.ReadTxOptions{}, + func(ctx context.Context, tx pgx.Tx) error { + var err error + + transactionNode, err = loadNode(ctx, tx) + if err != nil { + return err + } + + return tx.QueryRow( + ctx, + "SHOW transaction_read_only", + ).Scan(&readOnly) + }, + ); err != nil { + return fmt.Errorf("run shard read transaction: %w", err) + } + + fmt.Println() + fmt.Println("replica routing and read-only transaction:") + fmt.Printf( + "- user_id=%d shard=%s read_pool=%s read_node=%s tx_node=%s role=%s read_only=%s\n", + key, + resolved.ID(), + pool.Name(), + node.name, + transactionNode.name, + node.role, + readOnly, + ) + + return nil +} + +func showReferenceTables( + ctx context.Context, + topology *shard.Topology, +) error { + counts := make(map[shard.ID]int, topology.Len()) + var countsMu sync.Mutex + + // Fan out with bounded concurrency while preserving registration-order + // results for deterministic reporting. + results, err := topology.ForEachShard( + ctx, + referenceTableConcurrency, + func(ctx context.Context, resolved shard.Shard) error { + primary := resolved.Primary() + if primary == nil { + return fmt.Errorf("shard %q has no primary", resolved.ID()) + } + + var count int + if err := primary.QueryRow( + ctx, + `SELECT count(*) + FROM xpg_shard_example.countries`, + ).Scan(&count); err != nil { + return err + } + + countsMu.Lock() + counts[resolved.ID()] = count + countsMu.Unlock() + + return nil + }, + ) + if err != nil { + return fmt.Errorf("schedule reference-table reads: %w", err) + } + + if err := results.Err(); err != nil { + return fmt.Errorf("read reference tables: %w", err) + } + + fmt.Println() + fmt.Println("reference table copies:") + for index := 0; index < topology.Len(); index++ { + resolved := topology.At(index) + fmt.Printf( + "- shard=%s countries=%d\n", + resolved.ID(), + counts[resolved.ID()], + ) + } + + return nil +} + +func loadNode( + ctx context.Context, + db rowQuerier, +) (nodeInfo, error) { + var node nodeInfo + + err := db.QueryRow( + ctx, + `SELECT node_name, node_role + FROM xpg_shard_example.node_info`, + ).Scan(&node.name, &node.role) + if err != nil { + return nodeInfo{}, err + } + + return node, nil +} diff --git a/examples/shard/routing.go b/examples/shard/routing.go new file mode 100644 index 0000000..6079a8d --- /dev/null +++ b/examples/shard/routing.go @@ -0,0 +1,133 @@ +package main + +import ( + "context" + "errors" + "fmt" + + "github.com/jackc/pgx/v5" + "github.com/mkbeh/xpg/shard" +) + +const ( + shardAUserID uint64 = 42 + shardASecondUserID uint64 = 43 + shardBUserID uint64 = 142 + shardBSecondUserID uint64 = 143 +) + +func showRangeRouting( + ctx context.Context, + userResolver shard.Resolver[uint64], +) error { + fmt.Println("range routing and primary transactions:") + + for _, userID := range []uint64{ + shardAUserID, + shardBUserID, + } { + resolved, err := userResolver.Resolve(userID) + if err != nil { + return fmt.Errorf("resolve user %d: %w", userID, err) + } + + primary := resolved.Primary() + if primary == nil { + return fmt.Errorf("shard %q has no primary", resolved.ID()) + } + + name := fmt.Sprintf("user-%d", userID) + + if err := resolved.InPrimaryTx( + ctx, + pgx.TxOptions{}, + func(ctx context.Context, tx pgx.Tx) error { + _, err := tx.Exec( + ctx, + `INSERT INTO xpg_shard_example.users (id, name) + VALUES ($1, $2) + ON CONFLICT (id) DO UPDATE + SET name = EXCLUDED.name`, + userID, + name, + ) + + return err + }, + ); err != nil { + return fmt.Errorf("write user %d: %w", userID, err) + } + + fmt.Printf( + "- user_id=%d shard=%s primary_pool=%s\n", + userID, + resolved.ID(), + primary.Name(), + ) + } + + return nil +} + +func showGrouping( + userResolver shard.Resolver[uint64], +) error { + // The first occurrence belongs to shard-b, so GroupByShard returns + // shard-b before shard-a. + keys := []uint64{ + shardBUserID, + shardAUserID, + shardBSecondUserID, + shardASecondUserID, + } + + groups, err := shard.GroupByShard( + userResolver, + keys, + ) + if err != nil { + return fmt.Errorf("group users: %w", err) + } + + colocated, err := shard.SameShard( + userResolver, + shardAUserID, + shardASecondUserID, + ) + if err != nil { + return fmt.Errorf("verify colocated users: %w", err) + } + + _, mismatchErr := shard.SameShard( + userResolver, + shardAUserID, + shardBUserID, + ) + if !errors.Is(mismatchErr, shard.ErrShardMismatch) { + return fmt.Errorf( + "verify cross-shard users: expected ErrShardMismatch, got %v", + mismatchErr, + ) + } + + fmt.Println() + fmt.Println("grouping and colocation:") + + for _, group := range groups { + fmt.Printf( + "- shard=%s user_ids=%v\n", + group.Shard.ID(), + group.Keys, + ) + } + + fmt.Printf( + "- colocated shard=%s\n", + colocated.ID(), + ) + fmt.Println( + "- cross-shard SameShard returns ErrShardMismatch=true", + ) + + return nil +} diff --git a/examples/shard/setup.go b/examples/shard/setup.go new file mode 100644 index 0000000..b72cf3f --- /dev/null +++ b/examples/shard/setup.go @@ -0,0 +1,171 @@ +package main + +import ( + "context" + "fmt" + "os" + + "github.com/mkbeh/xpg" + "github.com/mkbeh/xpg/cluster" + "github.com/mkbeh/xpg/shard" +) + +const ( + defaultShardAPrimaryDatabaseURL = "postgres://postgres:postgres@localhost:56431/postgres?sslmode=disable&target_session_attrs=read-write" + defaultShardBPrimaryDatabaseURL = "postgres://postgres:postgres@localhost:56432/postgres?sslmode=disable&target_session_attrs=read-write" + defaultShardBReplicaDatabaseURL = "postgres://postgres:postgres@localhost:56433/postgres?sslmode=disable&target_session_attrs=read-only" + + shardAID shard.ID = "shard-a" + shardBID shard.ID = "shard-b" +) + +type poolSpec struct { + databaseURL string + name string +} + +type shardSpec struct { + id shard.ID + primary poolSpec + replicas []poolSpec +} + +func openTopology(ctx context.Context) (*shard.Topology, error) { + specs := []shardSpec{ + { + id: shardAID, + primary: poolSpec{ + databaseURL: environment( + "XPG_SHARD_A_PRIMARY_DATABASE_URL", + defaultShardAPrimaryDatabaseURL, + ), + name: "shard.shard-a.primary", + }, + }, + { + id: shardBID, + primary: poolSpec{ + databaseURL: environment( + "XPG_SHARD_B_PRIMARY_DATABASE_URL", + defaultShardBPrimaryDatabaseURL, + ), + name: "shard.shard-b.primary", + }, + replicas: []poolSpec{ + { + databaseURL: environment( + "XPG_SHARD_B_REPLICA_DATABASE_URL", + defaultShardBReplicaDatabaseURL, + ), + name: "shard.shard-b.replica", + }, + }, + }, + } + + clusters := make([]*cluster.Cluster, 0, len(specs)) + configs := make([]shard.Config, 0, len(specs)) + + for _, spec := range specs { + dbCluster, err := openCluster(ctx, spec) + if err != nil { + closeClusters(clusters) + + return nil, fmt.Errorf("open %s cluster: %w", spec.id, err) + } + + clusters = append(clusters, dbCluster) + configs = append(configs, shard.Config{Cluster: dbCluster}) + } + + // Topology takes ownership of the clusters only after successful creation. + topology, err := shard.NewTopology(configs) + if err != nil { + closeClusters(clusters) + + return nil, fmt.Errorf("create topology: %w", err) + } + + return topology, nil +} + +func openCluster(ctx context.Context, spec shardSpec) (*cluster.Cluster, error) { + primary, err := openPool(ctx, spec.primary) + if err != nil { + return nil, fmt.Errorf("open primary %s: %w", spec.primary.name, err) + } + + pools := []*xpg.Pool{primary} + replicas := make([]*xpg.Pool, 0, len(spec.replicas)) + + for _, replicaSpec := range spec.replicas { + replica, err := openPool(ctx, replicaSpec) + if err != nil { + closePools(pools) + + return nil, fmt.Errorf("open replica %s: %w", replicaSpec.name, err) + } + + pools = append(pools, replica) + replicas = append(replicas, replica) + } + + config := cluster.Config{ + ID: spec.id, + Primary: primary, + Replicas: replicas, + } + + if len(replicas) > 0 { + // Keep replica selection explicit in the runnable example. + config.Selector = cluster.RoundRobinSelector() + } + + dbCluster, err := cluster.New(config) + if err != nil { + closePools(pools) + + return nil, fmt.Errorf("create cluster: %w", err) + } + + return dbCluster, nil +} + +func openPool(ctx context.Context, spec poolSpec) (*xpg.Pool, error) { + pool, err := xpg.Open( + ctx, + spec.databaseURL, + xpg.WithName(spec.name), + ) + if err != nil { + return nil, fmt.Errorf("open %s: %w", spec.name, err) + } + + if err := pool.Ping(ctx); err != nil { + pool.Close() + + return nil, fmt.Errorf("ping %s: %w", spec.name, err) + } + + return pool, nil +} + +func closePools(pools []*xpg.Pool) { + for index := len(pools) - 1; index >= 0; index-- { + pools[index].Close() + } +} + +func closeClusters(clusters []*cluster.Cluster) { + for index := len(clusters) - 1; index >= 0; index-- { + clusters[index].Close() + } +} + +func environment(name, fallback string) string { + if value := os.Getenv(name); value != "" { + return value + } + + return fallback +} diff --git a/examples/shard/sql/shard-a-primary.sql b/examples/shard/sql/shard-a-primary.sql new file mode 100644 index 0000000..7baca90 --- /dev/null +++ b/examples/shard/sql/shard-a-primary.sql @@ -0,0 +1,2 @@ +INSERT INTO xpg_shard_example.node_info (node_name, node_role) +VALUES ('shard-a-primary', 'primary'); diff --git a/examples/shard/sql/shard-b-primary.sql b/examples/shard/sql/shard-b-primary.sql new file mode 100644 index 0000000..252c162 --- /dev/null +++ b/examples/shard/sql/shard-b-primary.sql @@ -0,0 +1,2 @@ +INSERT INTO xpg_shard_example.node_info (node_name, node_role) +VALUES ('shard-b-primary', 'primary'); diff --git a/examples/shard/sql/shard-b-replica.sql b/examples/shard/sql/shard-b-replica.sql new file mode 100644 index 0000000..14a8bc2 --- /dev/null +++ b/examples/shard/sql/shard-b-replica.sql @@ -0,0 +1,5 @@ +INSERT INTO xpg_shard_example.node_info (node_name, node_role) +VALUES ('shard-b-replica', 'replica'); + +ALTER DATABASE postgres +SET default_transaction_read_only = on; diff --git a/examples/shard/sql/shard.sql b/examples/shard/sql/shard.sql new file mode 100644 index 0000000..ddc8b69 --- /dev/null +++ b/examples/shard/sql/shard.sql @@ -0,0 +1,21 @@ +CREATE SCHEMA xpg_shard_example; + +CREATE TABLE xpg_shard_example.node_info ( + node_name text PRIMARY KEY, + node_role text NOT NULL +); + +CREATE TABLE xpg_shard_example.users ( + id bigint PRIMARY KEY, + name text NOT NULL +); + +CREATE TABLE xpg_shard_example.countries ( + code text PRIMARY KEY, + name text NOT NULL +); + +INSERT INTO xpg_shard_example.countries (code, name) +VALUES + ('NL', 'Netherlands'), + ('DE', 'Germany'); From 2746d72c393e994a20dc0ea4db84238315ff3a8f Mon Sep 17 00:00:00 2001 From: mkbeh Date: Wed, 12 Aug 2026 21:53:36 +0300 Subject: [PATCH 13/41] feat(cache): add typed in-process cache --- cache/cache.go | 359 +++++++++++++++++++++++++++++++++++++++++++++++ cache/config.go | 98 +++++++++++++ cache/doc.go | 9 ++ cache/loader.go | 21 +++ cache/storage.go | 330 +++++++++++++++++++++++++++++++++++++++++++ go.mod | 8 +- go.sum | 4 +- 7 files changed, 824 insertions(+), 5 deletions(-) create mode 100644 cache/cache.go create mode 100644 cache/config.go create mode 100644 cache/doc.go create mode 100644 cache/loader.go create mode 100644 cache/storage.go diff --git a/cache/cache.go b/cache/cache.go new file mode 100644 index 0000000..64a49c5 --- /dev/null +++ b/cache/cache.go @@ -0,0 +1,359 @@ +package cache + +import ( + "cmp" + "context" + "errors" + "math/rand/v2" + "slices" + "sync" + "time" + + "golang.org/x/sync/singleflight" +) + +// Cache is a concurrency-safe bounded in-process cache. +// +// Cache must be created with New and must not be copied after first use. +type Cache[V any] struct { + name string + + store *storage[V] + states []cacheState + + ttl time.Duration + jitter time.Duration + negativeTTL time.Duration +} + +type cacheState struct { + // mu serializes invalidation with singleflight registration and cache + // publication for keys routed to this state segment. + mu sync.RWMutex + + generation uint64 + group *singleflight.Group +} + +type invalidationTarget struct { + key string + index int +} + +// New creates a bounded in-process cache. +func New[V any](config Config) (*Cache[V], error) { + if err := config.validate(); err != nil { + return nil, err + } + + store := newStorage[V](config.MaxEntries, config.Segments) + + return &Cache[V]{ + name: config.Name, + store: store, + states: newCacheStates(len(store.segments)), + ttl: config.TTL, + jitter: config.Jitter, + negativeTTL: config.NegativeTTL, + }, nil +} + +// Name returns the configured cache name. +func (cache *Cache[V]) Name() string { + if cache == nil { + return "" + } + + return cache.name +} + +// GetOrLoad returns a cached value or executes loader on a cache miss. +// +// Concurrent misses for the same key share the loader started by the first +// caller. Each caller may stop waiting through its own context. +// +// The context of the caller that starts the shared load controls the loader. +// +// found=false represents a negative result. Negative results are cached only +// when Config.NegativeTTL is greater than zero. Loader errors are never cached. +func (cache *Cache[V]) GetOrLoad( + ctx context.Context, + key string, + loader Loader[V], +) (V, bool, error) { + var zero V + + if !cache.initialized() { + return zero, false, errors.New("xpg/cache: cache is not initialized") + } + + if loader == nil { + return zero, false, errors.New("xpg/cache: loader is nil") + } + + index := cache.store.segmentIndex(key) + + if cached, ok := cache.store.getAt(index, key, time.Now()); ok { + return cached.value, cached.found, nil + } + + state := &cache.states[index] + + // Registration and generation selection must be atomic with respect to + // invalidation for this state segment. DoChan only registers or starts the + // shared call; the loader itself executes outside state.mu. + state.mu.RLock() + + generation := state.generation + group := state.group + + resultChannel := group.DoChan( + key, + func() (any, error) { + // Another caller may have populated the cache between the initial + // lookup and this call becoming the singleflight owner. + if cached, ok := cache.store.getAt(index, key, time.Now()); ok { + return loadResult[V]{ + value: cached.value, + found: cached.found, + }, nil + } + + value, found, err := loader(ctx) + if err != nil { + return nil, err + } + + if !found { + value = zero + } + + loaded := loadResult[V]{ + value: value, + found: found, + } + + // Generation validation and publication are atomic with respect to + // invalidation for this state segment. A pre-invalidation load may + // still return to callers that already joined it, but cannot + // repopulate the cache after the barrier. + state.mu.RLock() + + if state.generation == generation { + cache.storeLoaded(index, key, loaded) + } + + state.mu.RUnlock() + + return loaded, nil + }, + ) + + state.mu.RUnlock() + + select { + case <-ctx.Done(): + return zero, false, ctx.Err() + + case result := <-resultChannel: + if result.Err != nil { + return zero, false, result.Err + } + + loaded, ok := result.Val.(loadResult[V]) + if !ok { + return zero, false, errors.New("xpg/cache: unexpected singleflight result type") + } + + return loaded.value, loaded.found, nil + } +} + +// Invalidate removes keys from the cache and prevents loads registered before +// the invalidation from repopulating them. +// +// Already running loaders are not canceled. Callers already waiting for such a +// loader may still receive its result. +func (cache *Cache[V]) Invalidate(keys ...string) { + if !cache.initialized() || len(keys) == 0 { + return + } + + // Keep the common single-key path allocation-free. + if len(keys) == 1 { + cache.invalidateOne(keys[0]) + return + } + + targets := make([]invalidationTarget, len(keys)) + + for index, key := range keys { + targets[index] = invalidationTarget{ + key: key, + index: cache.store.segmentIndex(key), + } + } + + // Every invalidation path acquires state locks in ascending segment order. + // This keeps multi-key invalidation and InvalidateAll deadlock-free. + slices.SortFunc( + targets, + func(left, right invalidationTarget) int { + return cmp.Compare(left.index, right.index) + }, + ) + + previous := -1 + + for _, target := range targets { + if target.index == previous { + continue + } + + cache.states[target.index].mu.Lock() + previous = target.index + } + + // Once every affected state is locked, advance each generation. Loads that + // registered before this barrier may finish for existing waiters, but they + // cannot publish into any affected segment afterward. + previous = -1 + + for _, target := range targets { + if target.index == previous { + continue + } + + cache.states[target.index].generation++ + previous = target.index + } + + for _, target := range targets { + state := &cache.states[target.index] + + // Forget ensures callers registered after this invalidation cannot join + // the pre-invalidation flight for key. + state.group.Forget(target.key) + cache.store.deleteAt(target.index, target.key) + } + + previous = -1 + + for index := len(targets) - 1; index >= 0; index-- { + target := targets[index] + if target.index == previous { + continue + } + + cache.states[target.index].mu.Unlock() + previous = target.index + } +} + +// InvalidateAll removes every cached entry and detaches future callers from all +// currently running singleflight calls. +// +// Existing loaders continue for callers that already joined them, but their +// results cannot repopulate the cache. +func (cache *Cache[V]) InvalidateAll() { + if !cache.initialized() { + return + } + + // Lock every state in a stable order. InvalidateAll is intentionally a + // cache-wide barrier and is expected to be rare compared with key-scoped + // invalidation. + for index := range cache.states { + cache.states[index].mu.Lock() + } + + for index := range cache.states { + state := &cache.states[index] + + state.generation++ + + // singleflight.Group has no ForgetAll operation. Existing callers retain + // the old group while future callers use this new group. + state.group = &singleflight.Group{} + } + + cache.store.deleteAll() + + for index := len(cache.states) - 1; index >= 0; index-- { + cache.states[index].mu.Unlock() + } +} + +func (cache *Cache[V]) invalidateOne( + key string, +) { + index := cache.store.segmentIndex(key) + state := &cache.states[index] + + state.mu.Lock() + + state.generation++ + state.group.Forget(key) + cache.store.deleteAt(index, key) + + state.mu.Unlock() +} + +func (cache *Cache[V]) storeLoaded( + index int, + key string, + loaded loadResult[V], +) { + now := time.Now() + + switch { + case loaded.found: + cache.store.setAt( + index, + key, + cachedValue[V]{ + value: loaded.value, + found: true, + }, + now.Add(cache.effectiveTTL()), + ) + + case cache.negativeTTL > 0: + cache.store.setAt( + index, + key, + cachedValue[V]{ + found: false, + }, + now.Add(cache.negativeTTL), + ) + } +} + +func (cache *Cache[V]) effectiveTTL() time.Duration { + if cache.jitter == 0 { + return cache.ttl + } + + return cache.ttl + time.Duration( + rand.Int64N( + int64(cache.jitter), + ), + ) +} + +func (cache *Cache[V]) initialized() bool { + return cache != nil && + cache.store != nil && + len(cache.states) == len(cache.store.segments) +} + +func newCacheStates(count int) []cacheState { + states := make([]cacheState, count) + + for index := range states { + states[index].group = &singleflight.Group{} + } + + return states +} diff --git a/cache/config.go b/cache/config.go new file mode 100644 index 0000000..129019d --- /dev/null +++ b/cache/config.go @@ -0,0 +1,98 @@ +package cache + +import ( + "errors" + "strings" + "time" +) + +const ( + maxDuration = time.Duration(1<<63 - 1) + defaultMaxEntries = 10_000 +) + +// Config configures a bounded in-process cache. +type Config struct { + // Name identifies the cache for diagnostics and observability. + Name string + + // MaxEntries is the maximum total entry budget of the cache. + // + // The budget is distributed across storage segments. Because each segment + // enforces its capacity independently, the number of simultaneously resident + // entries may be slightly lower than MaxEntries depending on key distribution. + // + // Zero uses the default. + MaxEntries int + + // Segments is the number of independent storage segments used to reduce lock + // contention. + // + // Zero uses the default segment count. + Segments int + + // TTL is the lifetime of positive entries. + TTL time.Duration + + // Jitter adds a random duration in [0, Jitter) to positive-entry TTLs. + // It can be used to spread expiration of entries loaded around the same time. + Jitter time.Duration + + // NegativeTTL is the lifetime of cached negative results. + // Zero disables negative caching. + NegativeTTL time.Duration +} + +func (config Config) validate() error { + name := strings.TrimSpace(config.Name) + + if name == "" { + return errors.New( + "xpg/cache: name must not be blank", + ) + } + + if name != config.Name { + return errors.New( + "xpg/cache: name must not contain surrounding whitespace", + ) + } + + if config.MaxEntries < 0 { + return errors.New( + "xpg/cache: max entries must not be negative", + ) + } + + if config.Segments < 0 { + return errors.New( + "xpg/cache: segments must not be negative", + ) + } + + if config.TTL <= 0 { + return errors.New( + "xpg/cache: ttl must be greater than zero", + ) + } + + if config.Jitter < 0 { + return errors.New( + "xpg/cache: jitter must not be negative", + ) + } + + if config.NegativeTTL < 0 { + return errors.New( + "xpg/cache: negative ttl must not be negative", + ) + } + + if config.Jitter > maxDuration-config.TTL { + return errors.New( + "xpg/cache: ttl and jitter overflow time.Duration", + ) + } + + return nil +} diff --git a/cache/doc.go b/cache/doc.go new file mode 100644 index 0000000..75f3bc1 --- /dev/null +++ b/cache/doc.go @@ -0,0 +1,9 @@ +// Package cache provides bounded in-process read-through caching. +// +// Cache entries use absolute TTL expiration, optional TTL jitter, LRU +// eviction, negative caching, duplicate load suppression, and explicit +// invalidation. +// +// The cache is local to one application process. It does not provide +// distributed cache coherence between application instances. +package cache diff --git a/cache/loader.go b/cache/loader.go new file mode 100644 index 0000000..244ded0 --- /dev/null +++ b/cache/loader.go @@ -0,0 +1,21 @@ +package cache + +import "context" + +// Loader loads one value. +// +// found=false represents a successful negative result. Negative results are +// cached only when Config.NegativeTTL is greater than zero. +// +// Loader errors are never cached. +type Loader[V any] func(ctx context.Context) (value V, found bool, err error) + +type loadResult[V any] struct { + value V + found bool +} + +type cachedValue[V any] struct { + value V + found bool +} diff --git a/cache/storage.go b/cache/storage.go new file mode 100644 index 0000000..a2f195f --- /dev/null +++ b/cache/storage.go @@ -0,0 +1,330 @@ +package cache + +import ( + "hash/maphash" + "sync" + "time" +) + +const defaultStorageSegmentCount = 32 + +type entry[V any] struct { + key string + cached cachedValue[V] + + expiresAt time.Time + + previous *entry[V] + next *entry[V] +} + +// storage routes keys across independent storage segments. +// +// Each segment owns its map, LRU list, and mutex. Segment capacities sum +// exactly to MaxEntries. +type storage[V any] struct { + seed maphash.Seed + + segments []storageSegment[V] + mask uint64 +} + +type storageSegment[V any] struct { + mu sync.Mutex + + entries map[string]*entry[V] + + head *entry[V] + tail *entry[V] + + maxEntries int +} + +func newStorage[V any]( + maxEntries, + segmentCount int, +) *storage[V] { + if maxEntries == 0 { + maxEntries = defaultMaxEntries + } + + if segmentCount == 0 { + segmentCount = defaultStorageSegmentCount + } + + return newStorageWithSegments[V](maxEntries, segmentCount) +} + +func newStorageWithSegments[V any](maxEntries, segmentCount int) *storage[V] { + segments := make([]storageSegment[V], segmentCount) + + baseCapacity := maxEntries / segmentCount + extraCapacity := maxEntries % segmentCount + + for index := range segments { + capacity := baseCapacity + + if index < extraCapacity { + capacity++ + } + + segments[index] = storageSegment[V]{ + entries: make(map[string]*entry[V], capacity), + maxEntries: capacity, + } + } + + mask := uint64(0) + + if segmentCount&(segmentCount-1) == 0 { + mask = uint64(segmentCount - 1) + } + + return &storage[V]{ + seed: maphash.MakeSeed(), + segments: segments, + mask: mask, + } +} + +func (storage *storage[V]) get( + key string, + now time.Time, +) (cachedValue[V], bool) { + return storage.getAt( + storage.segmentIndex(key), + key, + now, + ) +} + +func (storage *storage[V]) getAt( + index int, + key string, + now time.Time, +) (cachedValue[V], bool) { + return storage.segments[index].get( + key, + now, + ) +} + +func (storage *storage[V]) set( + key string, + value cachedValue[V], + expiresAt time.Time, +) { + storage.setAt( + storage.segmentIndex(key), + key, + value, + expiresAt, + ) +} + +func (storage *storage[V]) setAt( + index int, + key string, + value cachedValue[V], + expiresAt time.Time, +) { + storage.segments[index].set( + key, + value, + expiresAt, + ) +} + +func (storage *storage[V]) delete(key string) { + storage.deleteAt( + storage.segmentIndex(key), + key, + ) +} + +func (storage *storage[V]) deleteAt( + index int, + key string, +) { + storage.segments[index].delete(key) +} + +func (storage *storage[V]) deleteAll() { + for index := range storage.segments { + storage.segments[index].deleteAll() + } +} + +func (storage *storage[V]) segmentIndex(key string) int { + if len(storage.segments) == 1 { + return 0 + } + + hash := maphash.String(storage.seed, key) + + if storage.mask != 0 { + return int(hash & storage.mask) + } + + return int( + hash % uint64(len(storage.segments)), + ) +} + +func (segment *storageSegment[V]) get( + key string, + now time.Time, +) (cachedValue[V], bool) { + segment.mu.Lock() + defer segment.mu.Unlock() + + item, ok := segment.entries[key] + if !ok { + var zero cachedValue[V] + + return zero, false + } + + // TTL controls logical validity. Expired entries are removed lazily when + // accessed; there is no background or opportunistic expiration sweep. + if !now.Before(item.expiresAt) { + segment.removeLocked(item) + + var zero cachedValue[V] + + return zero, false + } + + // A hit affects LRU recency but never extends the entry TTL. + segment.moveToFrontLocked(item) + + return item.cached, true +} + +func (segment *storageSegment[V]) set( + key string, + value cachedValue[V], + expiresAt time.Time, +) { + // A zero-capacity segment is possible when the caller explicitly chooses + // more segments than MaxEntries. Such a segment simply stores nothing. + if segment.maxEntries == 0 { + return + } + + segment.mu.Lock() + defer segment.mu.Unlock() + + if item, ok := segment.entries[key]; ok { + item.cached = value + item.expiresAt = expiresAt + + segment.moveToFrontLocked(item) + + return + } + + // Once the segment reaches capacity, reuse its LRU victim instead of + // allocating another entry. + if len(segment.entries) >= segment.maxEntries { + item := segment.tail + + delete(segment.entries, item.key) + + item.key = key + item.cached = value + item.expiresAt = expiresAt + + segment.entries[key] = item + + segment.moveToFrontLocked(item) + + return + } + + item := &entry[V]{ + key: key, + cached: value, + expiresAt: expiresAt, + } + + segment.entries[key] = item + segment.pushFrontLocked(item) +} + +func (segment *storageSegment[V]) delete( + key string, +) { + segment.mu.Lock() + defer segment.mu.Unlock() + + item, ok := segment.entries[key] + if !ok { + return + } + + segment.removeLocked(item) +} + +func (segment *storageSegment[V]) deleteAll() { + segment.mu.Lock() + defer segment.mu.Unlock() + + clear(segment.entries) + + segment.head = nil + segment.tail = nil +} + +func (segment *storageSegment[V]) removeLocked( + item *entry[V], +) { + delete(segment.entries, item.key) + + segment.unlinkLocked(item) +} + +func (segment *storageSegment[V]) pushFrontLocked( + item *entry[V], +) { + item.previous = nil + item.next = segment.head + + if segment.head != nil { + segment.head.previous = item + } else { + segment.tail = item + } + + segment.head = item +} + +func (segment *storageSegment[V]) moveToFrontLocked( + item *entry[V], +) { + if segment.head == item { + return + } + + segment.unlinkLocked(item) + segment.pushFrontLocked(item) +} + +func (segment *storageSegment[V]) unlinkLocked( + item *entry[V], +) { + if item.previous != nil { + item.previous.next = item.next + } else { + segment.head = item.next + } + + if item.next != nil { + item.next.previous = item.previous + } else { + segment.tail = item.previous + } + + item.previous = nil + item.next = nil +} diff --git a/go.mod b/go.mod index a572d11..1beb438 100644 --- a/go.mod +++ b/go.mod @@ -2,12 +2,14 @@ module github.com/mkbeh/xpg go 1.26 -require github.com/jackc/pgx/v5 v5.10.0 +require ( + github.com/jackc/pgx/v5 v5.10.0 + golang.org/x/sync v0.22.0 +) require ( github.com/jackc/pgpassfile v1.0.0 // indirect github.com/jackc/pgservicefile v0.0.0-20240606120523-5a60cdf6a761 // indirect github.com/jackc/puddle/v2 v2.2.2 // indirect - golang.org/x/sync v0.22.0 // indirect - golang.org/x/text v0.40.0 // indirect + golang.org/x/text v0.41.0 // indirect ) diff --git a/go.sum b/go.sum index 583cb35..31890d5 100644 --- a/go.sum +++ b/go.sum @@ -18,8 +18,8 @@ github.com/stretchr/testify v1.11.1 h1:7s2iGBzp5EwR7/aIZr8ao5+dra3wiQyKjjFuvgVKu github.com/stretchr/testify v1.11.1/go.mod h1:wZwfW3scLgRK+23gO65QZefKpKQRnfz6sD981Nm4B6U= golang.org/x/sync v0.22.0 h1:SZjpbeLmrCk4xhRSZFNZW5gFUeCeFgjekvI/+gfScek= golang.org/x/sync v0.22.0/go.mod h1:9xrNwdLfx4jkKbNva9FpL6vEN7evnE43NNNJQ2LF3+0= -golang.org/x/text v0.40.0 h1:Ub2Z6/xjgF1WrYQz2nuITOEegKFtiIy+rieRJ5lHZKs= -golang.org/x/text v0.40.0/go.mod h1:hpnzDAfGV753zIKo+wk3u1bVKCGPbrnF7+7LBF/UHVY= +golang.org/x/text v0.41.0 h1:vz/seA0lnX87Othu2f/0L24RcgrXD9/YFTSuGjj3rH8= +golang.org/x/text v0.41.0/go.mod h1:jvf1O8ajNzZqhSrQBPbutR/EB83Cc0CFrezNQIwbb5M= gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0= gopkg.in/yaml.v3 v3.0.0-20200313102051-9f266ea9e77c/go.mod h1:K4uyk7z7BCEPqu6E+C64Yfv1cQ7kz7rIZviUmN+EgEM= gopkg.in/yaml.v3 v3.0.1 h1:fxVm/GzAzEWqLHuvctI91KS9hhNmmWOoWu0XTYJS7CA= From a7d5b3358ad8253d11f2c4f71c4e46b9247f87fc Mon Sep 17 00:00:00 2001 From: mkbeh Date: Thu, 13 Aug 2026 18:24:41 +0300 Subject: [PATCH 14/41] feat(cache): add cache statistics and metrics --- cache/cache.go | 159 +++++++----- cache/config.go | 7 +- cache/doc.go | 4 + cache/metrics.go | 22 ++ cache/stats.go | 189 +++++++++++++++ cache/storage.go | 123 ++++++---- metrics.go | 10 +- metrics/otel/cache.go | 451 +++++++++++++++++++++++++++++++++++ metrics/otel/metrics.go | 408 +------------------------------ metrics/otel/options.go | 10 +- metrics/otel/options_test.go | 62 +++-- metrics/otel/pool.go | 429 +++++++++++++++++++++++++++++++++ options.go | 4 +- pool.go | 4 +- 14 files changed, 1342 insertions(+), 540 deletions(-) create mode 100644 cache/metrics.go create mode 100644 cache/stats.go create mode 100644 metrics/otel/cache.go create mode 100644 metrics/otel/pool.go diff --git a/cache/cache.go b/cache/cache.go index 64a49c5..45abf1a 100644 --- a/cache/cache.go +++ b/cache/cache.go @@ -4,6 +4,7 @@ import ( "cmp" "context" "errors" + "fmt" "math/rand/v2" "slices" "sync" @@ -12,14 +13,15 @@ import ( "golang.org/x/sync/singleflight" ) -// Cache is a concurrency-safe bounded in-process cache. -// -// Cache must be created with New and must not be copied after first use. type Cache[V any] struct { name string store *storage[V] states []cacheState + stats *statsCollector + + metrics MetricsRegistration + closeOnce sync.Once ttl time.Duration jitter time.Duration @@ -27,8 +29,6 @@ type Cache[V any] struct { } type cacheState struct { - // mu serializes invalidation with singleflight registration and cache - // publication for keys routed to this state segment. mu sync.RWMutex generation uint64 @@ -40,7 +40,6 @@ type invalidationTarget struct { index int } -// New creates a bounded in-process cache. func New[V any](config Config) (*Cache[V], error) { if err := config.validate(); err != nil { return nil, err @@ -48,17 +47,25 @@ func New[V any](config Config) (*Cache[V], error) { store := newStorage[V](config.MaxEntries, config.Segments) - return &Cache[V]{ - name: config.Name, - store: store, - states: newCacheStates(len(store.segments)), + cache := &Cache[V]{ + name: config.Name, + + store: store, + states: newCacheStates(len(store.segments)), + stats: newStatsCollector(len(store.segments)), + ttl: config.TTL, jitter: config.Jitter, negativeTTL: config.NegativeTTL, - }, nil + } + + if err := cache.registerMetrics(config.Metrics); err != nil { + return nil, fmt.Errorf("xpg/cache: register metrics: %w", err) + } + + return cache, nil } -// Name returns the configured cache name. func (cache *Cache[V]) Name() string { if cache == nil { return "" @@ -67,15 +74,20 @@ func (cache *Cache[V]) Name() string { return cache.name } -// GetOrLoad returns a cached value or executes loader on a cache miss. -// -// Concurrent misses for the same key share the loader started by the first -// caller. Each caller may stop waiting through its own context. -// -// The context of the caller that starts the shared load controls the loader. -// -// found=false represents a negative result. Negative results are cached only -// when Config.NegativeTTL is greater than zero. Loader errors are never cached. +func (cache *Cache[V]) Close() { + if cache == nil { + return + } + + cache.closeOnce.Do( + func() { + if cache.metrics != nil { + cache.metrics.Close() + } + }, + ) +} + func (cache *Cache[V]) GetOrLoad( ctx context.Context, key string, @@ -92,8 +104,9 @@ func (cache *Cache[V]) GetOrLoad( } index := cache.store.segmentIndex(key) + stats := cache.stats.shard(index) - if cached, ok := cache.store.getAt(index, key, time.Now()); ok { + if cached, ok := cache.store.lookupAt(index, key, time.Now(), stats); ok { return cached.value, cached.found, nil } @@ -112,20 +125,29 @@ func (cache *Cache[V]) GetOrLoad( func() (any, error) { // Another caller may have populated the cache between the initial // lookup and this call becoming the singleflight owner. - if cached, ok := cache.store.getAt(index, key, time.Now()); ok { + if cached, ok := cache.store.getAt(index, key, time.Now(), cache.stats.shard(index)); ok { return loadResult[V]{ value: cached.value, found: cached.found, }, nil } + startedAt := time.Now() + value, found, err := loader(ctx) + + finishedAt := time.Now() + + cache.stats.recordLoad(index, found, err, finishedAt.Sub(startedAt)) + if err != nil { return nil, err } if !found { - value = zero + var zeroValue V + + value = zeroValue } loaded := loadResult[V]{ @@ -137,13 +159,15 @@ func (cache *Cache[V]) GetOrLoad( // invalidation for this state segment. A pre-invalidation load may // still return to callers that already joined it, but cannot // repopulate the cache after the barrier. - state.mu.RLock() + publishState := &cache.states[index] - if state.generation == generation { - cache.storeLoaded(index, key, loaded) + publishState.mu.RLock() + + if publishState.generation == generation { + cache.storeLoaded(index, key, loaded, finishedAt) } - state.mu.RUnlock() + publishState.mu.RUnlock() return loaded, nil }, @@ -156,6 +180,10 @@ func (cache *Cache[V]) GetOrLoad( return zero, false, ctx.Err() case result := <-resultChannel: + if result.Shared { + cache.stats.recordShared(index) + } + if result.Err != nil { return zero, false, result.Err } @@ -169,19 +197,18 @@ func (cache *Cache[V]) GetOrLoad( } } -// Invalidate removes keys from the cache and prevents loads registered before -// the invalidation from repopulating them. -// -// Already running loaders are not canceled. Callers already waiting for such a -// loader may still receive its result. func (cache *Cache[V]) Invalidate(keys ...string) { - if !cache.initialized() || len(keys) == 0 { + if !cache.initialized() || + len(keys) == 0 { return } - // Keep the common single-key path allocation-free. if len(keys) == 1 { - cache.invalidateOne(keys[0]) + index := cache.store.segmentIndex(keys[0]) + + cache.invalidateOne(index, keys[0]) + cache.stats.recordInvalidation(index) + return } @@ -194,6 +221,11 @@ func (cache *Cache[V]) Invalidate(keys ...string) { } } + // Keep the statistics shard based on the caller's first key rather than + // the sorted target order. One non-empty Invalidate call contributes one + // operation regardless of how many keys it contains. + statsIndex := targets[0].index + // Every invalidation path acquires state locks in ascending segment order. // This keeps multi-key invalidation and InvalidateAll deadlock-free. slices.SortFunc( @@ -211,6 +243,7 @@ func (cache *Cache[V]) Invalidate(keys ...string) { } cache.states[target.index].mu.Lock() + previous = target.index } @@ -225,15 +258,15 @@ func (cache *Cache[V]) Invalidate(keys ...string) { } cache.states[target.index].generation++ + previous = target.index } for _, target := range targets { state := &cache.states[target.index] - // Forget ensures callers registered after this invalidation cannot join - // the pre-invalidation flight for key. state.group.Forget(target.key) + cache.store.deleteAt(target.index, target.key) } @@ -241,28 +274,24 @@ func (cache *Cache[V]) Invalidate(keys ...string) { for index := len(targets) - 1; index >= 0; index-- { target := targets[index] + if target.index == previous { continue } cache.states[target.index].mu.Unlock() + previous = target.index } + + cache.stats.recordInvalidation(statsIndex) } -// InvalidateAll removes every cached entry and detaches future callers from all -// currently running singleflight calls. -// -// Existing loaders continue for callers that already joined them, but their -// results cannot repopulate the cache. func (cache *Cache[V]) InvalidateAll() { if !cache.initialized() { return } - // Lock every state in a stable order. InvalidateAll is intentionally a - // cache-wide barrier and is expected to be rare compared with key-scoped - // invalidation. for index := range cache.states { cache.states[index].mu.Lock() } @@ -271,9 +300,6 @@ func (cache *Cache[V]) InvalidateAll() { state := &cache.states[index] state.generation++ - - // singleflight.Group has no ForgetAll operation. Existing callers retain - // the old group while future callers use this new group. state.group = &singleflight.Group{} } @@ -282,18 +308,21 @@ func (cache *Cache[V]) InvalidateAll() { for index := len(cache.states) - 1; index >= 0; index-- { cache.states[index].mu.Unlock() } + + cache.stats.recordInvalidateAll() } func (cache *Cache[V]) invalidateOne( + index int, key string, ) { - index := cache.store.segmentIndex(key) state := &cache.states[index] state.mu.Lock() state.generation++ state.group.Forget(key) + cache.store.deleteAt(index, key) state.mu.Unlock() @@ -303,9 +332,8 @@ func (cache *Cache[V]) storeLoaded( index int, key string, loaded loadResult[V], + now time.Time, ) { - now := time.Now() - switch { case loaded.found: cache.store.setAt( @@ -316,6 +344,7 @@ func (cache *Cache[V]) storeLoaded( found: true, }, now.Add(cache.effectiveTTL()), + cache.stats.shard(index), ) case cache.negativeTTL > 0: @@ -326,6 +355,7 @@ func (cache *Cache[V]) storeLoaded( found: false, }, now.Add(cache.negativeTTL), + cache.stats.shard(index), ) } } @@ -336,16 +366,33 @@ func (cache *Cache[V]) effectiveTTL() time.Duration { } return cache.ttl + time.Duration( - rand.Int64N( - int64(cache.jitter), - ), + rand.Int64N(int64(cache.jitter)), ) } +func (cache *Cache[V]) registerMetrics(metrics Metrics) error { + if metrics == nil { + return nil + } + + registration, err := metrics.RegisterCache(cache) + if err != nil { + return err + } + + cache.metrics = registration + + return nil +} + func (cache *Cache[V]) initialized() bool { return cache != nil && cache.store != nil && - len(cache.states) == len(cache.store.segments) + cache.stats != nil && + len(cache.states) == + len(cache.store.segments) && + len(cache.stats.shards) == + len(cache.store.segments) } func newCacheStates(count int) []cacheState { diff --git a/cache/config.go b/cache/config.go index 129019d..12d51a1 100644 --- a/cache/config.go +++ b/cache/config.go @@ -2,12 +2,13 @@ package cache import ( "errors" + "math" "strings" "time" ) const ( - maxDuration = time.Duration(1<<63 - 1) + maxDuration = time.Duration(math.MaxInt64) defaultMaxEntries = 10_000 ) @@ -41,6 +42,10 @@ type Config struct { // NegativeTTL is the lifetime of cached negative results. // Zero disables negative caching. NegativeTTL time.Duration + + // Metrics optionally registers cache statistics during New. + // The registration is released when Cache.Close is called. + Metrics Metrics } func (config Config) validate() error { diff --git a/cache/doc.go b/cache/doc.go index 75f3bc1..af8342b 100644 --- a/cache/doc.go +++ b/cache/doc.go @@ -4,6 +4,10 @@ // eviction, negative caching, duplicate load suppression, and explicit // invalidation. // +// Cache statistics are collected locally and exposed through Cache.Stats. +// Optional metrics integrations register during New and observe those snapshots +// without adding telemetry calls to the cache request path. +// // The cache is local to one application process. It does not provide // distributed cache coherence between application instances. package cache diff --git a/cache/metrics.go b/cache/metrics.go new file mode 100644 index 0000000..05145e3 --- /dev/null +++ b/cache/metrics.go @@ -0,0 +1,22 @@ +package cache + +// StatsProvider exposes one cache's identity and statistics to a metrics +// implementation. Cache implements StatsProvider. +type StatsProvider interface { + Name() string + Stats() Stats +} + +// Metrics registers observability for a Cache. +// +// Implementations are expected to be immutable and safe to reuse for multiple +// caches. RegisterCache is called after the cache has been fully initialized. +type Metrics interface { + RegisterCache(cache StatsProvider) (MetricsRegistration, error) +} + +// MetricsRegistration owns a metrics registration associated with one Cache. +// Close is called once when the Cache is closed. +type MetricsRegistration interface { + Close() +} diff --git a/cache/stats.go b/cache/stats.go new file mode 100644 index 0000000..fa18cbd --- /dev/null +++ b/cache/stats.go @@ -0,0 +1,189 @@ +package cache + +import ( + "sync/atomic" + "time" +) + +// Stats is a detached snapshot of cache statistics. +// +// The snapshot is assembled from independent cache segments. Concurrent cache +// activity may continue while Stats is being collected, so individual fields +// are not guaranteed to represent one globally atomic instant. +type Stats struct { + // Current state. + + // EntryCount is the number of entries currently resident in storage. + // Because expiration is lazy, this count may include expired entries that + // have not been accessed or evicted yet. + EntryCount int64 + + // MaxEntries is the configured total entry budget after applying defaults. + MaxEntries int64 + + // SegmentCount is the number of independent storage and coordination + // segments. + SegmentCount int64 + + // Lookup lifecycle. + + // HitCount is the cumulative number of positive cache hits. + HitCount int64 + + // NegativeHitCount is the cumulative number of cached negative hits. + NegativeHitCount int64 + + // MissCount is the cumulative number of cache misses. An expired entry + // observed by a caller is counted as a miss. + MissCount int64 + + // Load lifecycle. + + // LoadFoundCount is the cumulative number of actual loader invocations that + // returned found=true. + LoadFoundCount int64 + + // LoadNotFoundCount is the cumulative number of actual loader invocations + // that returned found=false without an error. + LoadNotFoundCount int64 + + // LoadErrorCount is the cumulative number of actual loader invocations that + // returned an error. + LoadErrorCount int64 + + // LoadDuration is the cumulative duration of actual loader invocations, + // including successful, negative, and failed loads. + LoadDuration time.Duration + + // SharedCount is the cumulative number of callers that received a + // singleflight result shared with at least one other caller. + SharedCount int64 + + // Invalidation lifecycle. + + // InvalidationCount is the cumulative number of non-empty Invalidate calls. + InvalidationCount int64 + + // InvalidateAllCount is the cumulative number of InvalidateAll calls. + InvalidateAllCount int64 + + // Storage lifecycle. + + // EvictionCount is the cumulative number of entries evicted because a + // storage segment reached capacity. + EvictionCount int64 + + // ExpirationCount is the cumulative number of expired entries removed + // lazily while looking up a key. + ExpirationCount int64 +} + +// statsCollector owns all mutable cache statistics. +// +// Each shard corresponds to one storage and coordination segment. +type statsCollector struct { + shards []statsShard + + // InvalidateAll is cache-wide and expected to be rare, so keeping this + // counter global does not introduce meaningful contention. + invalidateAllCount atomic.Int64 +} + +type statsShard struct { + // Protected by the corresponding storage segment mutex. + hitCount int64 + negativeHitCount int64 + missCount int64 + evictionCount int64 + expirationCount int64 + + // Updated outside a suitable existing lock. + loadFoundCount atomic.Int64 + loadNotFoundCount atomic.Int64 + loadErrorCount atomic.Int64 + loadDurationNanos atomic.Int64 + + // These events are observed without a suitable existing lock. + // Counters are sharded to avoid a single global hot cache line. + sharedCount atomic.Int64 + invalidationCount atomic.Int64 +} + +func newStatsCollector(segmentCount int) *statsCollector { + return &statsCollector{ + shards: make([]statsShard, segmentCount), + } +} + +// Stats returns a detached snapshot of the current cache statistics. +func (cache *Cache[V]) Stats() Stats { + if !cache.initialized() { + return Stats{} + } + + snapshot := Stats{ + MaxEntries: int64(cache.store.maxEntries), + SegmentCount: int64(len(cache.store.segments)), + InvalidateAllCount: cache.stats.invalidateAllCount.Load(), + } + + for index := range cache.store.segments { + segment := &cache.store.segments[index] + shard := cache.stats.shard(index) + + // Storage counters share the storage segment mutex with the operations + // that update them. + segment.mu.Lock() + + snapshot.EntryCount += int64(len(segment.entries)) + snapshot.HitCount += shard.hitCount + snapshot.NegativeHitCount += shard.negativeHitCount + snapshot.MissCount += shard.missCount + snapshot.EvictionCount += shard.evictionCount + snapshot.ExpirationCount += shard.expirationCount + + segment.mu.Unlock() + + snapshot.LoadFoundCount += shard.loadFoundCount.Load() + snapshot.LoadNotFoundCount += shard.loadNotFoundCount.Load() + snapshot.LoadErrorCount += shard.loadErrorCount.Load() + snapshot.LoadDuration += time.Duration(shard.loadDurationNanos.Load()) + snapshot.SharedCount += shard.sharedCount.Load() + snapshot.InvalidationCount += shard.invalidationCount.Load() + } + + return snapshot +} + +func (stats *statsCollector) shard(index int) *statsShard { + return &stats.shards[index] +} + +func (stats *statsCollector) recordLoad(index int, found bool, err error, duration time.Duration) { + shard := stats.shard(index) + + shard.loadDurationNanos.Add(duration.Nanoseconds()) + + switch { + case err != nil: + shard.loadErrorCount.Add(1) + + case found: + shard.loadFoundCount.Add(1) + + default: + shard.loadNotFoundCount.Add(1) + } +} + +func (stats *statsCollector) recordShared(index int) { + stats.shard(index).sharedCount.Add(1) +} + +func (stats *statsCollector) recordInvalidation(index int) { + stats.shard(index).invalidationCount.Add(1) +} + +func (stats *statsCollector) recordInvalidateAll() { + stats.invalidateAllCount.Add(1) +} diff --git a/cache/storage.go b/cache/storage.go index a2f195f..700a330 100644 --- a/cache/storage.go +++ b/cache/storage.go @@ -25,8 +25,9 @@ type entry[V any] struct { type storage[V any] struct { seed maphash.Seed - segments []storageSegment[V] - mask uint64 + segments []storageSegment[V] + mask uint64 + maxEntries int } type storageSegment[V any] struct { @@ -40,10 +41,7 @@ type storageSegment[V any] struct { maxEntries int } -func newStorage[V any]( - maxEntries, - segmentCount int, -) *storage[V] { +func newStorage[V any](maxEntries, segmentCount int) *storage[V] { if maxEntries == 0 { maxEntries = defaultMaxEntries } @@ -81,9 +79,10 @@ func newStorageWithSegments[V any](maxEntries, segmentCount int) *storage[V] { } return &storage[V]{ - seed: maphash.MakeSeed(), - segments: segments, - mask: mask, + seed: maphash.MakeSeed(), + segments: segments, + mask: mask, + maxEntries: maxEntries, } } @@ -91,22 +90,25 @@ func (storage *storage[V]) get( key string, now time.Time, ) (cachedValue[V], bool) { - return storage.getAt( - storage.segmentIndex(key), - key, - now, - ) + return storage.getAt(storage.segmentIndex(key), key, now, nil) +} + +func (storage *storage[V]) lookupAt( + index int, + key string, + now time.Time, + stats *statsShard, +) (cachedValue[V], bool) { + return storage.segments[index].lookup(key, now, stats) } func (storage *storage[V]) getAt( index int, key string, now time.Time, + stats *statsShard, ) (cachedValue[V], bool) { - return storage.segments[index].get( - key, - now, - ) + return storage.segments[index].get(key, now, stats) } func (storage *storage[V]) set( @@ -114,12 +116,7 @@ func (storage *storage[V]) set( value cachedValue[V], expiresAt time.Time, ) { - storage.setAt( - storage.segmentIndex(key), - key, - value, - expiresAt, - ) + storage.setAt(storage.segmentIndex(key), key, value, expiresAt, nil) } func (storage *storage[V]) setAt( @@ -127,19 +124,13 @@ func (storage *storage[V]) setAt( key string, value cachedValue[V], expiresAt time.Time, + stats *statsShard, ) { - storage.segments[index].set( - key, - value, - expiresAt, - ) + storage.segments[index].set(key, value, expiresAt, stats) } func (storage *storage[V]) delete(key string) { - storage.deleteAt( - storage.segmentIndex(key), - key, - ) + storage.deleteAt(storage.segmentIndex(key), key) } func (storage *storage[V]) deleteAt( @@ -171,13 +162,50 @@ func (storage *storage[V]) segmentIndex(key string) int { ) } +func (segment *storageSegment[V]) lookup( + key string, + now time.Time, + stats *statsShard, +) (cachedValue[V], bool) { + segment.mu.Lock() + defer segment.mu.Unlock() + + cached, ok := segment.getLocked(key, now, stats) + if !ok { + if stats != nil { + stats.missCount++ + } + + return cached, false + } + + if stats != nil { + if cached.found { + stats.hitCount++ + } else { + stats.negativeHitCount++ + } + } + + return cached, true +} + func (segment *storageSegment[V]) get( key string, now time.Time, + stats *statsShard, ) (cachedValue[V], bool) { segment.mu.Lock() defer segment.mu.Unlock() + return segment.getLocked(key, now, stats) +} + +func (segment *storageSegment[V]) getLocked( + key string, + now time.Time, + stats *statsShard, +) (cachedValue[V], bool) { item, ok := segment.entries[key] if !ok { var zero cachedValue[V] @@ -190,6 +218,10 @@ func (segment *storageSegment[V]) get( if !now.Before(item.expiresAt) { segment.removeLocked(item) + if stats != nil { + stats.expirationCount++ + } + var zero cachedValue[V] return zero, false @@ -205,6 +237,7 @@ func (segment *storageSegment[V]) set( key string, value cachedValue[V], expiresAt time.Time, + stats *statsShard, ) { // A zero-capacity segment is possible when the caller explicitly chooses // more segments than MaxEntries. Such a segment simply stores nothing. @@ -227,6 +260,10 @@ func (segment *storageSegment[V]) set( // Once the segment reaches capacity, reuse its LRU victim instead of // allocating another entry. if len(segment.entries) >= segment.maxEntries { + if stats != nil { + stats.evictionCount++ + } + item := segment.tail delete(segment.entries, item.key) @@ -252,9 +289,7 @@ func (segment *storageSegment[V]) set( segment.pushFrontLocked(item) } -func (segment *storageSegment[V]) delete( - key string, -) { +func (segment *storageSegment[V]) delete(key string) { segment.mu.Lock() defer segment.mu.Unlock() @@ -276,17 +311,13 @@ func (segment *storageSegment[V]) deleteAll() { segment.tail = nil } -func (segment *storageSegment[V]) removeLocked( - item *entry[V], -) { +func (segment *storageSegment[V]) removeLocked(item *entry[V]) { delete(segment.entries, item.key) segment.unlinkLocked(item) } -func (segment *storageSegment[V]) pushFrontLocked( - item *entry[V], -) { +func (segment *storageSegment[V]) pushFrontLocked(item *entry[V]) { item.previous = nil item.next = segment.head @@ -299,9 +330,7 @@ func (segment *storageSegment[V]) pushFrontLocked( segment.head = item } -func (segment *storageSegment[V]) moveToFrontLocked( - item *entry[V], -) { +func (segment *storageSegment[V]) moveToFrontLocked(item *entry[V]) { if segment.head == item { return } @@ -310,9 +339,7 @@ func (segment *storageSegment[V]) moveToFrontLocked( segment.pushFrontLocked(item) } -func (segment *storageSegment[V]) unlinkLocked( - item *entry[V], -) { +func (segment *storageSegment[V]) unlinkLocked(item *entry[V]) { if item.previous != nil { item.previous.next = item.next } else { diff --git a/metrics.go b/metrics.go index c100f03..be69822 100644 --- a/metrics.go +++ b/metrics.go @@ -1,15 +1,15 @@ package xpg -// PoolMetrics registers metrics for a Pool. +// Metrics registers metrics for a Pool. // // Implementations are expected to be immutable and safe to reuse for multiple // pools. Register is called after the underlying pgxpool.Pool has been created. -type PoolMetrics interface { - Register(pool *Pool) (PoolMetricsRegistration, error) +type Metrics interface { + Register(pool *Pool) (MetricsRegistration, error) } -// PoolMetricsRegistration owns a metrics registration associated with one +// MetricsRegistration owns a metrics registration associated with one // Pool. Close is called once before the underlying pgxpool.Pool is closed. -type PoolMetricsRegistration interface { +type MetricsRegistration interface { Close() } diff --git a/metrics/otel/cache.go b/metrics/otel/cache.go new file mode 100644 index 0000000..3b7e8e9 --- /dev/null +++ b/metrics/otel/cache.go @@ -0,0 +1,451 @@ +package xpgotel + +import ( + "context" + "errors" + "fmt" + "slices" + + xpgcache "github.com/mkbeh/xpg/cache" + "go.opentelemetry.io/otel" + "go.opentelemetry.io/otel/attribute" + "go.opentelemetry.io/otel/metric" +) + +const ( + cacheEntryCountMetricName = "xpg.cache.entry.count" + cacheEntryMaxMetricName = "xpg.cache.entry.max" + cacheSegmentCountMetricName = "xpg.cache.segment.count" + cacheLookupCountMetricName = "xpg.cache.lookup.count" + cacheLoadCountMetricName = "xpg.cache.load.count" + cacheLoadTimeMetricName = "xpg.cache.load.time" + cacheSharedCountMetricName = "xpg.cache.singleflight.shared.count" + cacheInvalidationMetricName = "xpg.cache.invalidation.count" + cacheEvictionCountMetricName = "xpg.cache.entry.eviction.count" + cacheExpirationCountMetricName = "xpg.cache.entry.expiration.count" +) + +const ( + cacheNameAttribute = "xpg.cache.name" + cacheLookupResultAttribute = "xpg.cache.lookup.result" + cacheLoadResultAttribute = "xpg.cache.load.result" + cacheInvalidationScopeAttribute = "xpg.cache.invalidation.scope" +) + +const ( + cacheLookupResultHit = "hit" + cacheLookupResultNegativeHit = "negative_hit" + cacheLookupResultMiss = "miss" + + cacheLoadResultFound = "found" + cacheLoadResultNotFound = "not_found" + cacheLoadResultError = "error" + + cacheInvalidationScopeKeys = "keys" + cacheInvalidationScopeAll = "all" +) + +var _ xpgcache.Metrics = (*Metrics)(nil) + +type cacheMetricInstruments struct { + entryCount metric.Int64ObservableGauge + entryMax metric.Int64ObservableGauge + segmentCount metric.Int64ObservableGauge + lookupCount metric.Int64ObservableCounter + loadCount metric.Int64ObservableCounter + loadTime metric.Float64ObservableCounter + sharedCount metric.Int64ObservableCounter + invalidationCount metric.Int64ObservableCounter + evictionCount metric.Int64ObservableCounter + expirationCount metric.Int64ObservableCounter +} + +type cacheMetricAttributes struct { + base metric.ObserveOption + + hit metric.ObserveOption + negativeHit metric.ObserveOption + miss metric.ObserveOption + + loadFound metric.ObserveOption + loadNotFound metric.ObserveOption + loadError metric.ObserveOption + + invalidateKeys metric.ObserveOption + invalidateAll metric.ObserveOption +} + +// RegisterCache registers OpenTelemetry metrics for one cache. +func (m *Metrics) RegisterCache(cache xpgcache.StatsProvider) (xpgcache.MetricsRegistration, error) { + if m == nil { + return nil, errors.New("xpg/otel: metrics is nil") + } + + if cache == nil { + return nil, errors.New("xpg/otel: cache is nil") + } + + name := cache.Name() + if name == "" { + return nil, errors.New("xpg/otel: cache name is blank") + } + + provider := m.meterProvider + if provider == nil { + provider = otel.GetMeterProvider() + } + + return registerCacheMetrics(cache, name, provider) +} + +func registerCacheMetrics( + cache xpgcache.StatsProvider, + name string, + provider metric.MeterProvider, +) (xpgcache.MetricsRegistration, error) { + meter := provider.Meter(instrumentationName) + + instruments, err := newCacheMetricInstruments(meter) + if err != nil { + return nil, err + } + + attributes := newCacheMetricAttributes(name) + + registration, err := meter.RegisterCallback( + func(_ context.Context, observer metric.Observer) error { + instruments.observe(observer, cache.Stats(), attributes) + + return nil + }, + instruments.observables()..., + ) + if err != nil { + return nil, fmt.Errorf("xpg/otel: register cache metrics callback: %w", err) + } + + return &metricsRegistration{ + registration: registration, + }, nil +} + +func (instruments cacheMetricInstruments) observe( + observer metric.Observer, + stats xpgcache.Stats, + attributes cacheMetricAttributes, +) { + observer.ObserveInt64( + instruments.entryCount, + stats.EntryCount, + attributes.base, + ) + observer.ObserveInt64( + instruments.entryMax, + stats.MaxEntries, + attributes.base, + ) + observer.ObserveInt64( + instruments.segmentCount, + stats.SegmentCount, + attributes.base, + ) + observer.ObserveInt64( + instruments.lookupCount, + stats.HitCount, + attributes.hit, + ) + observer.ObserveInt64( + instruments.lookupCount, + stats.NegativeHitCount, + attributes.negativeHit, + ) + observer.ObserveInt64( + instruments.lookupCount, + stats.MissCount, + attributes.miss, + ) + observer.ObserveInt64( + instruments.loadCount, + stats.LoadFoundCount, + attributes.loadFound, + ) + observer.ObserveInt64( + instruments.loadCount, + stats.LoadNotFoundCount, + attributes.loadNotFound, + ) + observer.ObserveInt64( + instruments.loadCount, + stats.LoadErrorCount, + attributes.loadError, + ) + observer.ObserveFloat64( + instruments.loadTime, + stats.LoadDuration.Seconds(), + attributes.base, + ) + observer.ObserveInt64( + instruments.sharedCount, + stats.SharedCount, + attributes.base, + ) + observer.ObserveInt64( + instruments.invalidationCount, + stats.InvalidationCount, + attributes.invalidateKeys, + ) + observer.ObserveInt64( + instruments.invalidationCount, + stats.InvalidateAllCount, + attributes.invalidateAll, + ) + observer.ObserveInt64( + instruments.evictionCount, + stats.EvictionCount, + attributes.base, + ) + observer.ObserveInt64( + instruments.expirationCount, + stats.ExpirationCount, + attributes.base, + ) +} + +func (instruments cacheMetricInstruments) observables() []metric.Observable { + return []metric.Observable{ + instruments.entryCount, + instruments.entryMax, + instruments.segmentCount, + instruments.lookupCount, + instruments.loadCount, + instruments.loadTime, + instruments.sharedCount, + instruments.invalidationCount, + instruments.evictionCount, + instruments.expirationCount, + } +} + +func newCacheMetricInstruments(meter metric.Meter) (cacheMetricInstruments, error) { + var instruments cacheMetricInstruments + + var err error + + instruments.entryCount, err = meter.Int64ObservableGauge( + cacheEntryCountMetricName, + metric.WithDescription( + "The number of entries currently resident in cache storage.", + ), + metric.WithUnit("{entry}"), + ) + if err != nil { + return cacheMetricInstruments{}, fmt.Errorf( + "xpg/otel: create %s: %w", + cacheEntryCountMetricName, + err, + ) + } + + instruments.entryMax, err = meter.Int64ObservableGauge( + cacheEntryMaxMetricName, + metric.WithDescription( + "The configured total cache entry budget.", + ), + metric.WithUnit("{entry}"), + ) + if err != nil { + return cacheMetricInstruments{}, fmt.Errorf( + "xpg/otel: create %s: %w", + cacheEntryMaxMetricName, + err, + ) + } + + instruments.segmentCount, err = meter.Int64ObservableGauge( + cacheSegmentCountMetricName, + metric.WithDescription( + "The number of independent cache storage segments.", + ), + metric.WithUnit("{segment}"), + ) + if err != nil { + return cacheMetricInstruments{}, fmt.Errorf( + "xpg/otel: create %s: %w", + cacheSegmentCountMetricName, + err, + ) + } + + instruments.lookupCount, err = meter.Int64ObservableCounter( + cacheLookupCountMetricName, + metric.WithDescription( + "The cumulative number of cache lookups by result.", + ), + metric.WithUnit("{lookup}"), + ) + if err != nil { + return cacheMetricInstruments{}, fmt.Errorf( + "xpg/otel: create %s: %w", + cacheLookupCountMetricName, + err, + ) + } + + instruments.loadCount, err = meter.Int64ObservableCounter( + cacheLoadCountMetricName, + metric.WithDescription( + "The cumulative number of actual cache loader invocations by result.", + ), + metric.WithUnit("{load}"), + ) + if err != nil { + return cacheMetricInstruments{}, fmt.Errorf( + "xpg/otel: create %s: %w", + cacheLoadCountMetricName, + err, + ) + } + + instruments.loadTime, err = meter.Float64ObservableCounter( + cacheLoadTimeMetricName, + metric.WithDescription( + "The cumulative time spent in actual cache loader invocations.", + ), + metric.WithUnit("s"), + ) + if err != nil { + return cacheMetricInstruments{}, fmt.Errorf( + "xpg/otel: create %s: %w", + cacheLoadTimeMetricName, + err, + ) + } + + instruments.sharedCount, err = meter.Int64ObservableCounter( + cacheSharedCountMetricName, + metric.WithDescription( + "The cumulative number of cache callers that received a shared singleflight result.", + ), + metric.WithUnit("{request}"), + ) + if err != nil { + return cacheMetricInstruments{}, fmt.Errorf( + "xpg/otel: create %s: %w", + cacheSharedCountMetricName, + err, + ) + } + + instruments.invalidationCount, err = meter.Int64ObservableCounter( + cacheInvalidationMetricName, + metric.WithDescription( + "The cumulative number of explicit cache invalidation operations by scope.", + ), + metric.WithUnit("{operation}"), + ) + if err != nil { + return cacheMetricInstruments{}, fmt.Errorf( + "xpg/otel: create %s: %w", + cacheInvalidationMetricName, + err, + ) + } + + instruments.evictionCount, err = meter.Int64ObservableCounter( + cacheEvictionCountMetricName, + metric.WithDescription( + "The cumulative number of cache entries evicted because a storage segment reached capacity.", + ), + metric.WithUnit("{entry}"), + ) + if err != nil { + return cacheMetricInstruments{}, fmt.Errorf( + "xpg/otel: create %s: %w", + cacheEvictionCountMetricName, + err, + ) + } + + instruments.expirationCount, err = meter.Int64ObservableCounter( + cacheExpirationCountMetricName, + metric.WithDescription( + "The cumulative number of expired cache entries removed during lookup.", + ), + metric.WithUnit("{entry}"), + ) + if err != nil { + return cacheMetricInstruments{}, fmt.Errorf( + "xpg/otel: create %s: %w", + cacheExpirationCountMetricName, + err, + ) + } + + return instruments, nil +} + +func newCacheMetricAttributes(name string) cacheMetricAttributes { + base := []attribute.KeyValue{ + attribute.String(cacheNameAttribute, name), + } + + option := func(extra ...attribute.KeyValue) metric.ObserveOption { + return metric.WithAttributeSet( + attribute.NewSet( + slices.Concat(base, extra)..., + ), + ) + } + + return cacheMetricAttributes{ + base: option(), + hit: option( + attribute.String( + cacheLookupResultAttribute, + cacheLookupResultHit, + ), + ), + negativeHit: option( + attribute.String( + cacheLookupResultAttribute, + cacheLookupResultNegativeHit, + ), + ), + miss: option( + attribute.String( + cacheLookupResultAttribute, + cacheLookupResultMiss, + ), + ), + loadFound: option( + attribute.String( + cacheLoadResultAttribute, + cacheLoadResultFound, + ), + ), + loadNotFound: option( + attribute.String( + cacheLoadResultAttribute, + cacheLoadResultNotFound, + ), + ), + loadError: option( + attribute.String( + cacheLoadResultAttribute, + cacheLoadResultError, + ), + ), + invalidateKeys: option( + attribute.String( + cacheInvalidationScopeAttribute, + cacheInvalidationScopeKeys, + ), + ), + invalidateAll: option( + attribute.String( + cacheInvalidationScopeAttribute, + cacheInvalidationScopeAll, + ), + ), + } +} diff --git a/metrics/otel/metrics.go b/metrics/otel/metrics.go index 706599b..0cff385 100644 --- a/metrics/otel/metrics.go +++ b/metrics/otel/metrics.go @@ -1,127 +1,33 @@ package xpgotel import ( - "context" - "errors" "fmt" - "slices" "sync" - "github.com/mkbeh/xpg" "go.opentelemetry.io/otel" - "go.opentelemetry.io/otel/attribute" "go.opentelemetry.io/otel/metric" ) const instrumentationName = "github.com/mkbeh/xpg/otel" -const ( - connectionCountMetricName = "db.client.connection.count" - connectionMaxMetricName = "db.client.connection.max" - connectionConstructingMetricName = "xpg.pool.connection.constructing" - connectionAcquireCountMetricName = "xpg.pool.connection.acquire.count" - connectionAcquireTimeMetricName = "xpg.pool.connection.acquire.time" - connectionAcquireCanceledCountMetricName = "xpg.pool.connection.acquire.canceled.count" - connectionAcquireEmptyCountMetricName = "xpg.pool.connection.acquire.empty.count" - connectionAcquireEmptyWaitTimeMetricName = "xpg.pool.connection.acquire.empty.wait_time" - connectionCreateCountMetricName = "xpg.pool.connection.create.count" - connectionDestroyCountMetricName = "xpg.pool.connection.destroy.count" -) - -const ( - dbSystemNameAttribute = "db.system.name" - poolNameAttribute = "db.client.connection.pool.name" - connectionStateAttribute = "db.client.connection.state" - destroyReasonAttribute = "xpg.pool.connection.destroy.reason" -) - -const ( - dbSystemPostgreSQL = "postgresql" - - connectionStateIdle = "idle" - connectionStateUsed = "used" - - destroyReasonIdleTimeout = "idle_timeout" - destroyReasonLifetime = "lifetime" -) - -type poolMetricInstruments struct { - connectionCount metric.Int64ObservableUpDownCounter - connectionMax metric.Int64ObservableUpDownCounter - constructingConnections metric.Int64ObservableGauge - acquireCount metric.Int64ObservableCounter - acquireTime metric.Float64ObservableCounter - canceledAcquireCount metric.Int64ObservableCounter - emptyAcquireCount metric.Int64ObservableCounter - emptyAcquireWaitTime metric.Float64ObservableCounter - createdConnections metric.Int64ObservableCounter - destroyedConnections metric.Int64ObservableCounter -} - -type poolMetricAttributes struct { - base metric.ObserveOption - idle metric.ObserveOption - used metric.ObserveOption - destroyedIdle metric.ObserveOption - destroyedLifetime metric.ObserveOption -} - -// Metrics exports xpg pool statistics through OpenTelemetry. +// Metrics exports xpg statistics through OpenTelemetry. +// +// Metrics is immutable after construction and may be reused for multiple +// pool and cache registrations. type Metrics struct { meterProvider metric.MeterProvider } -type poolMetrics struct { +// metricsRegistration owns one OpenTelemetry callback registration. +// +// The same implementation is used by pool and cache metrics because both +// registrations have identical lifecycle semantics. +type metricsRegistration struct { registration metric.Registration closeOnce sync.Once } -// Register registers metrics for one xpg Pool. -func (m *Metrics) Register( - pool *xpg.Pool, -) (xpg.PoolMetricsRegistration, error) { - if m == nil { - return nil, errors.New("xpg/otel: metrics is nil") - } - - provider := m.meterProvider - if provider == nil { - provider = otel.GetMeterProvider() - } - - return registerPoolMetrics(pool, provider) -} - -func registerPoolMetrics( - pool *xpg.Pool, - provider metric.MeterProvider, -) (xpg.PoolMetricsRegistration, error) { - meter := provider.Meter(instrumentationName) - - instruments, err := newPoolMetricInstruments(meter) - if err != nil { - return nil, err - } - - attributes := newPoolMetricAttributes(pool.Name(), pool.Labels()) - - registration, err := meter.RegisterCallback( - func(_ context.Context, observer metric.Observer) error { - instruments.observe(observer, pool.Stats(), attributes) - return nil - }, - instruments.observables()..., - ) - if err != nil { - return nil, fmt.Errorf("xpg/otel: register pool metrics callback: %w", err) - } - - return &poolMetrics{ - registration: registration, - }, nil -} - -func (m *poolMetrics) Close() { +func (m *metricsRegistration) Close() { if m == nil || m.registration == nil { return } @@ -129,300 +35,8 @@ func (m *poolMetrics) Close() { m.closeOnce.Do(func() { if err := m.registration.Unregister(); err != nil { otel.Handle( - fmt.Errorf("xpg/otel: unregister pool metrics: %w", err), + fmt.Errorf("xpg/otel: unregister metrics: %w", err), ) } }) } - -func (i poolMetricInstruments) observe( - observer metric.Observer, - stats xpg.PoolStats, - attributes poolMetricAttributes, -) { - observer.ObserveInt64( - i.connectionCount, - int64(stats.IdleConns), - attributes.idle, - ) - observer.ObserveInt64( - i.connectionCount, - int64(stats.AcquiredConns), - attributes.used, - ) - observer.ObserveInt64( - i.connectionMax, - int64(stats.MaxConns), - attributes.base, - ) - observer.ObserveInt64( - i.constructingConnections, - int64(stats.ConstructingConns), - attributes.base, - ) - observer.ObserveInt64( - i.acquireCount, - stats.AcquireCount, - attributes.base, - ) - observer.ObserveFloat64( - i.acquireTime, - stats.AcquireDuration.Seconds(), - attributes.base, - ) - observer.ObserveInt64( - i.canceledAcquireCount, - stats.CanceledAcquireCount, - attributes.base, - ) - observer.ObserveInt64( - i.emptyAcquireCount, - stats.EmptyAcquireCount, - attributes.base, - ) - observer.ObserveFloat64( - i.emptyAcquireWaitTime, - stats.EmptyAcquireWaitTime.Seconds(), - attributes.base, - ) - observer.ObserveInt64( - i.createdConnections, - stats.NewConnsCount, - attributes.base, - ) - observer.ObserveInt64( - i.destroyedConnections, - stats.MaxIdleDestroyCount, - attributes.destroyedIdle, - ) - observer.ObserveInt64( - i.destroyedConnections, - stats.MaxLifetimeDestroyCount, - attributes.destroyedLifetime, - ) -} - -func (i poolMetricInstruments) observables() []metric.Observable { - return []metric.Observable{ - i.connectionCount, - i.connectionMax, - i.constructingConnections, - i.acquireCount, - i.acquireTime, - i.canceledAcquireCount, - i.emptyAcquireCount, - i.emptyAcquireWaitTime, - i.createdConnections, - i.destroyedConnections, - } -} - -func newPoolMetricInstruments(meter metric.Meter) (poolMetricInstruments, error) { - var instruments poolMetricInstruments - - var err error - - instruments.connectionCount, err = meter.Int64ObservableUpDownCounter( - connectionCountMetricName, - metric.WithDescription( - "The number of connections currently used or idle in the pool.", - ), - metric.WithUnit("{connection}"), - ) - if err != nil { - return poolMetricInstruments{}, fmt.Errorf( - "xpg/otel: create %s: %w", - connectionCountMetricName, - err, - ) - } - - instruments.connectionMax, err = meter.Int64ObservableUpDownCounter( - connectionMaxMetricName, - metric.WithDescription( - "The maximum number of open connections allowed by the pool.", - ), - metric.WithUnit("{connection}"), - ) - if err != nil { - return poolMetricInstruments{}, fmt.Errorf( - "xpg/otel: create %s: %w", - connectionMaxMetricName, - err, - ) - } - - instruments.constructingConnections, err = meter.Int64ObservableGauge( - connectionConstructingMetricName, - metric.WithDescription( - "The number of connections currently being created by the pool.", - ), - metric.WithUnit("{connection}"), - ) - if err != nil { - return poolMetricInstruments{}, fmt.Errorf( - "xpg/otel: create %s: %w", - connectionConstructingMetricName, - err, - ) - } - - instruments.acquireCount, err = meter.Int64ObservableCounter( - connectionAcquireCountMetricName, - metric.WithDescription( - "The cumulative number of successful connection acquires.", - ), - metric.WithUnit("{request}"), - ) - if err != nil { - return poolMetricInstruments{}, fmt.Errorf( - "xpg/otel: create %s: %w", - connectionAcquireCountMetricName, - err, - ) - } - - instruments.acquireTime, err = meter.Float64ObservableCounter( - connectionAcquireTimeMetricName, - metric.WithDescription( - "The cumulative time spent acquiring connections.", - ), - metric.WithUnit("s"), - ) - if err != nil { - return poolMetricInstruments{}, fmt.Errorf( - "xpg/otel: create %s: %w", - connectionAcquireTimeMetricName, - err, - ) - } - - instruments.canceledAcquireCount, err = meter.Int64ObservableCounter( - connectionAcquireCanceledCountMetricName, - metric.WithDescription( - "The cumulative number of connection acquires canceled by context.", - ), - metric.WithUnit("{request}"), - ) - if err != nil { - return poolMetricInstruments{}, fmt.Errorf( - "xpg/otel: create %s: %w", - connectionAcquireCanceledCountMetricName, - err, - ) - } - - instruments.emptyAcquireCount, err = meter.Int64ObservableCounter( - connectionAcquireEmptyCountMetricName, - metric.WithDescription( - "The cumulative number of successful acquires that waited because the pool was empty.", - ), - metric.WithUnit("{request}"), - ) - if err != nil { - return poolMetricInstruments{}, fmt.Errorf( - "xpg/otel: create %s: %w", - connectionAcquireEmptyCountMetricName, - err, - ) - } - - instruments.emptyAcquireWaitTime, err = meter.Float64ObservableCounter( - connectionAcquireEmptyWaitTimeMetricName, - metric.WithDescription( - "The cumulative time spent waiting for a connection while the pool was empty.", - ), - metric.WithUnit("s"), - ) - if err != nil { - return poolMetricInstruments{}, fmt.Errorf( - "xpg/otel: create %s: %w", - connectionAcquireEmptyWaitTimeMetricName, - err, - ) - } - - instruments.createdConnections, err = meter.Int64ObservableCounter( - connectionCreateCountMetricName, - metric.WithDescription( - "The cumulative number of connections opened by the pool.", - ), - metric.WithUnit("{connection}"), - ) - if err != nil { - return poolMetricInstruments{}, fmt.Errorf( - "xpg/otel: create %s: %w", - connectionCreateCountMetricName, - err, - ) - } - - instruments.destroyedConnections, err = meter.Int64ObservableCounter( - connectionDestroyCountMetricName, - metric.WithDescription( - "The cumulative number of connections destroyed by pool lifecycle limits.", - ), - metric.WithUnit("{connection}"), - ) - if err != nil { - return poolMetricInstruments{}, fmt.Errorf( - "xpg/otel: create %s: %w", - connectionDestroyCountMetricName, - err, - ) - } - - return instruments, nil -} - -func newPoolMetricAttributes(name string, labels map[string]string) poolMetricAttributes { - var base []attribute.KeyValue - - for key, value := range labels { - base = append(base, attribute.String(key, value)) - } - - // System attributes are appended last, so xpg-controlled values win when - // user labels contain duplicate keys. - base = append( - base, - attribute.String(dbSystemNameAttribute, dbSystemPostgreSQL), - attribute.String(poolNameAttribute, name), - ) - - option := func(extra ...attribute.KeyValue) metric.ObserveOption { - return metric.WithAttributeSet( - attribute.NewSet( - slices.Concat(base, extra)..., - ), - ) - } - - return poolMetricAttributes{ - base: option(), - idle: option( - attribute.String( - connectionStateAttribute, - connectionStateIdle, - ), - ), - used: option( - attribute.String( - connectionStateAttribute, - connectionStateUsed, - ), - ), - destroyedIdle: option( - attribute.String( - destroyReasonAttribute, - destroyReasonIdleTimeout, - ), - ), - destroyedLifetime: option( - attribute.String( - destroyReasonAttribute, - destroyReasonLifetime, - ), - ), - } -} diff --git a/metrics/otel/options.go b/metrics/otel/options.go index dac9ab0..fb8f27e 100644 --- a/metrics/otel/options.go +++ b/metrics/otel/options.go @@ -4,7 +4,7 @@ import ( "go.opentelemetry.io/otel/metric" ) -// MetricsOption configures OpenTelemetry pool metrics. +// MetricsOption configures OpenTelemetry metrics. // // The interface is sealed so options can only be created by this package. type MetricsOption interface { @@ -21,10 +21,10 @@ type metricsSettings struct { meterProvider metric.MeterProvider } -// NewMetrics creates an OpenTelemetry pool metrics implementation. +// NewMetrics creates an OpenTelemetry metrics implementation. // // By default, metrics use the global OpenTelemetry MeterProvider. The returned -// value is immutable and may be reused for multiple pools. +// value is immutable and may be reused for multiple pools and caches. func NewMetrics(options ...MetricsOption) *Metrics { settings := metricsSettings{} @@ -41,10 +41,10 @@ func NewMetrics(options ...MetricsOption) *Metrics { } } -// WithMeterProvider configures the MeterProvider used for pool metrics. +// WithMeterProvider configures the MeterProvider used for metrics. // // The caller owns the provider and must shut it down after all instrumented -// pools have been closed. +// pools and caches have been closed. func WithMeterProvider(provider metric.MeterProvider) MetricsOption { return metricsOptionFunc(func(settings *metricsSettings) { if provider != nil { diff --git a/metrics/otel/options_test.go b/metrics/otel/options_test.go index da0d53e..1467a3a 100644 --- a/metrics/otel/options_test.go +++ b/metrics/otel/options_test.go @@ -2,56 +2,70 @@ package xpgotel import ( "context" - "strings" "testing" + "time" "github.com/jackc/pgx/v5/pgxpool" "github.com/mkbeh/xpg" + xpgcache "github.com/mkbeh/xpg/cache" "go.opentelemetry.io/otel/metric/noop" ) -func TestWithMetrics(t *testing.T) { +func TestMetricsRegistration(t *testing.T) { t.Parallel() - config, err := pgxpool.ParseConfig("") + metrics := NewMetrics( + WithMeterProvider(noop.NewMeterProvider()), + ) + + poolConfig, err := pgxpool.ParseConfig("") if err != nil { t.Fatalf("parse pool config: %v", err) } pool, err := xpg.New( context.Background(), - config, - xpg.WithName("test"), - WithMetrics( - WithMeterProvider(noop.NewMeterProvider()), - ), + poolConfig, + xpg.WithName("test-pool"), + xpg.WithMetrics(metrics), ) if err != nil { t.Fatalf("create pool: %v", err) } - pool.Close() pool.Close() + + cache, err := xpgcache.New[int]( + xpgcache.Config{ + Name: "test-cache", + TTL: time.Minute, + Metrics: metrics, + }, + ) + if err != nil { + t.Fatalf("create cache: %v", err) + } + cache.Close() + cache.Close() } -func TestWithMetricsNilProvider(t *testing.T) { +func TestWithMeterProviderNilUsesGlobalProvider(t *testing.T) { t.Parallel() - config, err := pgxpool.ParseConfig("") - if err != nil { - t.Fatalf("parse pool config: %v", err) - } + metrics := NewMetrics( + WithMeterProvider(nil), + ) - pool, err := xpg.New( - context.Background(), - config, - WithMetrics(WithMeterProvider(nil)), + cache, err := xpgcache.New[int]( + xpgcache.Config{ + Name: "test-cache", + TTL: time.Minute, + Metrics: metrics, + }, ) - if pool != nil { - pool.Close() - t.Fatal("New returned a pool with a nil MeterProvider") - } - if err == nil || !strings.Contains(err.Error(), "meter provider is nil") { - t.Fatalf("unexpected error: %v", err) + if err != nil { + t.Fatalf("create cache: %v", err) } + + cache.Close() } diff --git a/metrics/otel/pool.go b/metrics/otel/pool.go new file mode 100644 index 0000000..9b610ae --- /dev/null +++ b/metrics/otel/pool.go @@ -0,0 +1,429 @@ +package xpgotel + +import ( + "context" + "errors" + "fmt" + "slices" + + "github.com/mkbeh/xpg" + "go.opentelemetry.io/otel" + "go.opentelemetry.io/otel/attribute" + "go.opentelemetry.io/otel/metric" +) + +const ( + connectionCountMetricName = "db.client.connection.count" + connectionMaxMetricName = "db.client.connection.max" + connectionConstructingMetricName = "xpg.pool.connection.constructing" + connectionAcquireCountMetricName = "xpg.pool.connection.acquire.count" + connectionAcquireTimeMetricName = "xpg.pool.connection.acquire.time" + connectionAcquireCanceledCountMetricName = "xpg.pool.connection.acquire.canceled.count" + connectionAcquireEmptyCountMetricName = "xpg.pool.connection.acquire.empty.count" + connectionAcquireEmptyWaitTimeMetricName = "xpg.pool.connection.acquire.empty.wait_time" + connectionCreateCountMetricName = "xpg.pool.connection.create.count" + connectionDestroyCountMetricName = "xpg.pool.connection.destroy.count" +) + +const ( + dbSystemNameAttribute = "db.system.name" + poolNameAttribute = "db.client.connection.pool.name" + connectionStateAttribute = "db.client.connection.state" + destroyReasonAttribute = "xpg.pool.connection.destroy.reason" +) + +const ( + dbSystemPostgreSQL = "postgresql" + + connectionStateIdle = "idle" + connectionStateUsed = "used" + + destroyReasonIdleTimeout = "idle_timeout" + destroyReasonLifetime = "lifetime" +) + +var _ xpg.Metrics = (*Metrics)(nil) + +type poolMetricInstruments struct { + connectionCount metric.Int64ObservableUpDownCounter + connectionMax metric.Int64ObservableUpDownCounter + constructingConnections metric.Int64ObservableGauge + acquireCount metric.Int64ObservableCounter + acquireTime metric.Float64ObservableCounter + canceledAcquireCount metric.Int64ObservableCounter + emptyAcquireCount metric.Int64ObservableCounter + emptyAcquireWaitTime metric.Float64ObservableCounter + createdConnections metric.Int64ObservableCounter + destroyedConnections metric.Int64ObservableCounter +} + +type poolMetricAttributes struct { + base metric.ObserveOption + idle metric.ObserveOption + used metric.ObserveOption + destroyedIdle metric.ObserveOption + destroyedLifetime metric.ObserveOption +} + +var _ xpg.Metrics = (*Metrics)(nil) + +// Register registers metrics for one xpg Pool. +func (m *Metrics) Register(pool *xpg.Pool) (xpg.MetricsRegistration, error) { + if m == nil { + return nil, errors.New("xpg/otel: metrics is nil") + } + + provider := m.meterProvider + if provider == nil { + provider = otel.GetMeterProvider() + } + + return registerPoolMetrics(pool, provider) +} + +func registerPoolMetrics(pool *xpg.Pool, provider metric.MeterProvider) (xpg.MetricsRegistration, error) { + meter := provider.Meter(instrumentationName) + + instruments, err := newPoolMetricInstruments(meter) + if err != nil { + return nil, err + } + + attributes := newPoolMetricAttributes(pool.Name(), pool.Labels()) + + registration, err := meter.RegisterCallback( + func(_ context.Context, observer metric.Observer) error { + instruments.observe( + observer, + pool.Stats(), + attributes, + ) + + return nil + }, + instruments.observables()..., + ) + if err != nil { + return nil, fmt.Errorf("xpg/otel: register pool metrics callback: %w", err) + } + + return &metricsRegistration{ + registration: registration, + }, nil +} + +func (i poolMetricInstruments) observe( + observer metric.Observer, + stats xpg.PoolStats, + attributes poolMetricAttributes, +) { + observer.ObserveInt64( + i.connectionCount, + int64(stats.IdleConns), + attributes.idle, + ) + + observer.ObserveInt64( + i.connectionCount, + int64(stats.AcquiredConns), + attributes.used, + ) + + observer.ObserveInt64( + i.connectionMax, + int64(stats.MaxConns), + attributes.base, + ) + + observer.ObserveInt64( + i.constructingConnections, + int64(stats.ConstructingConns), + attributes.base, + ) + + observer.ObserveInt64( + i.acquireCount, + stats.AcquireCount, + attributes.base, + ) + + observer.ObserveFloat64( + i.acquireTime, + stats.AcquireDuration.Seconds(), + attributes.base, + ) + + observer.ObserveInt64( + i.canceledAcquireCount, + stats.CanceledAcquireCount, + attributes.base, + ) + + observer.ObserveInt64( + i.emptyAcquireCount, + stats.EmptyAcquireCount, + attributes.base, + ) + + observer.ObserveFloat64( + i.emptyAcquireWaitTime, + stats.EmptyAcquireWaitTime.Seconds(), + attributes.base, + ) + + observer.ObserveInt64( + i.createdConnections, + stats.NewConnsCount, + attributes.base, + ) + + observer.ObserveInt64( + i.destroyedConnections, + stats.MaxIdleDestroyCount, + attributes.destroyedIdle, + ) + + observer.ObserveInt64( + i.destroyedConnections, + stats.MaxLifetimeDestroyCount, + attributes.destroyedLifetime, + ) +} + +func (i poolMetricInstruments) observables() []metric.Observable { + return []metric.Observable{ + i.connectionCount, + i.connectionMax, + i.constructingConnections, + i.acquireCount, + i.acquireTime, + i.canceledAcquireCount, + i.emptyAcquireCount, + i.emptyAcquireWaitTime, + i.createdConnections, + i.destroyedConnections, + } +} + +func newPoolMetricInstruments(meter metric.Meter) (poolMetricInstruments, error) { + var instruments poolMetricInstruments + + var err error + + instruments.connectionCount, err = meter.Int64ObservableUpDownCounter( + connectionCountMetricName, + metric.WithDescription( + "The number of connections currently used or idle in the pool.", + ), + metric.WithUnit("{connection}"), + ) + if err != nil { + return poolMetricInstruments{}, fmt.Errorf( + "xpg/otel: create %s: %w", + connectionCountMetricName, + err, + ) + } + + instruments.connectionMax, err = meter.Int64ObservableUpDownCounter( + connectionMaxMetricName, + metric.WithDescription( + "The maximum number of open connections allowed by the pool.", + ), + metric.WithUnit("{connection}"), + ) + if err != nil { + return poolMetricInstruments{}, fmt.Errorf( + "xpg/otel: create %s: %w", + connectionMaxMetricName, + err, + ) + } + + instruments.constructingConnections, err = meter.Int64ObservableGauge( + connectionConstructingMetricName, + metric.WithDescription( + "The number of connections currently being created by the pool.", + ), + metric.WithUnit("{connection}"), + ) + if err != nil { + return poolMetricInstruments{}, fmt.Errorf( + "xpg/otel: create %s: %w", + connectionConstructingMetricName, + err, + ) + } + + instruments.acquireCount, err = meter.Int64ObservableCounter( + connectionAcquireCountMetricName, + metric.WithDescription( + "The cumulative number of successful connection acquires.", + ), + metric.WithUnit("{request}"), + ) + if err != nil { + return poolMetricInstruments{}, fmt.Errorf( + "xpg/otel: create %s: %w", + connectionAcquireCountMetricName, + err, + ) + } + + instruments.acquireTime, err = meter.Float64ObservableCounter( + connectionAcquireTimeMetricName, + metric.WithDescription( + "The cumulative time spent acquiring connections.", + ), + metric.WithUnit("s"), + ) + if err != nil { + return poolMetricInstruments{}, fmt.Errorf( + "xpg/otel: create %s: %w", + connectionAcquireTimeMetricName, + err, + ) + } + + instruments.canceledAcquireCount, err = meter.Int64ObservableCounter( + connectionAcquireCanceledCountMetricName, + metric.WithDescription( + "The cumulative number of connection acquires canceled by context.", + ), + metric.WithUnit("{request}"), + ) + if err != nil { + return poolMetricInstruments{}, fmt.Errorf( + "xpg/otel: create %s: %w", + connectionAcquireCanceledCountMetricName, + err, + ) + } + + instruments.emptyAcquireCount, err = meter.Int64ObservableCounter( + connectionAcquireEmptyCountMetricName, + metric.WithDescription( + "The cumulative number of successful acquires that waited because the pool was empty.", + ), + metric.WithUnit("{request}"), + ) + if err != nil { + return poolMetricInstruments{}, fmt.Errorf( + "xpg/otel: create %s: %w", + connectionAcquireEmptyCountMetricName, + err, + ) + } + + instruments.emptyAcquireWaitTime, err = meter.Float64ObservableCounter( + connectionAcquireEmptyWaitTimeMetricName, + metric.WithDescription( + "The cumulative time spent waiting for a connection while the pool was empty.", + ), + metric.WithUnit("s"), + ) + if err != nil { + return poolMetricInstruments{}, fmt.Errorf( + "xpg/otel: create %s: %w", + connectionAcquireEmptyWaitTimeMetricName, + err, + ) + } + + instruments.createdConnections, err = meter.Int64ObservableCounter( + connectionCreateCountMetricName, + metric.WithDescription( + "The cumulative number of connections opened by the pool.", + ), + metric.WithUnit("{connection}"), + ) + if err != nil { + return poolMetricInstruments{}, fmt.Errorf( + "xpg/otel: create %s: %w", + connectionCreateCountMetricName, + err, + ) + } + + instruments.destroyedConnections, err = meter.Int64ObservableCounter( + connectionDestroyCountMetricName, + metric.WithDescription( + "The cumulative number of connections destroyed by pool lifecycle limits.", + ), + metric.WithUnit("{connection}"), + ) + if err != nil { + return poolMetricInstruments{}, fmt.Errorf( + "xpg/otel: create %s: %w", + connectionDestroyCountMetricName, + err, + ) + } + + return instruments, nil +} + +func newPoolMetricAttributes(name string, labels map[string]string) poolMetricAttributes { + var base []attribute.KeyValue + + for key, value := range labels { + base = append( + base, + attribute.String(key, value), + ) + } + + // System attributes are appended last, so xpg-controlled values win when + // user labels contain duplicate keys. + base = append( + base, + attribute.String( + dbSystemNameAttribute, + dbSystemPostgreSQL, + ), + attribute.String( + poolNameAttribute, + name, + ), + ) + + option := func(extra ...attribute.KeyValue) metric.ObserveOption { + return metric.WithAttributeSet( + attribute.NewSet( + slices.Concat(base, extra)..., + ), + ) + } + + return poolMetricAttributes{ + base: option(), + + idle: option( + attribute.String( + connectionStateAttribute, + connectionStateIdle, + ), + ), + + used: option( + attribute.String( + connectionStateAttribute, + connectionStateUsed, + ), + ), + + destroyedIdle: option( + attribute.String( + destroyReasonAttribute, + destroyReasonIdleTimeout, + ), + ), + + destroyedLifetime: option( + attribute.String( + destroyReasonAttribute, + destroyReasonLifetime, + ), + ), + } +} diff --git a/options.go b/options.go index 164522b..4bea3ff 100644 --- a/options.go +++ b/options.go @@ -25,7 +25,7 @@ func (option optionFunc) apply(settings *settings) error { type settings struct { name string labels map[string]string - metrics PoolMetrics + metrics Metrics } func (s settings) poolName(host string, port uint16, database string) string { @@ -122,7 +122,7 @@ func WithLabel(key, value string) Option { // // Metrics are registered during New and unregistered automatically when the // Pool is closed. -func WithMetrics(metrics PoolMetrics) Option { +func WithMetrics(metrics Metrics) Option { return optionFunc(func(settings *settings) error { if metrics == nil { return errors.New("pool metrics is nil") diff --git a/pool.go b/pool.go index 3ed4223..759dfeb 100644 --- a/pool.go +++ b/pool.go @@ -14,7 +14,7 @@ import ( // Pool is a concurrency-safe PostgreSQL connection pool backed by pgxpool. type Pool struct { pool *pgxpool.Pool - metrics PoolMetricsRegistration + metrics MetricsRegistration name string labels map[string]string @@ -127,7 +127,7 @@ func (p *Pool) CopyFrom(ctx context.Context, tableName pgx.Identifier, columnNam return p.pool.CopyFrom(ctx, tableName, columnNames, rowSrc) } -func (p *Pool) registerMetrics(metrics PoolMetrics) error { +func (p *Pool) registerMetrics(metrics Metrics) error { if metrics == nil { return nil } From 435d3d0dd81b4ebf539170ca72d2ca659a1cf935 Mon Sep 17 00:00:00 2001 From: mkbeh Date: Thu, 13 Aug 2026 20:03:47 +0300 Subject: [PATCH 15/41] refactor(cache): refine invalidation metrics --- cache/cache.go | 32 ++++++++++++++----------- cache/stats.go | 55 ++++++++++++++++++++++++++----------------- cache/storage.go | 43 +++++++++++++-------------------- metrics/otel/cache.go | 8 +++---- 4 files changed, 72 insertions(+), 66 deletions(-) diff --git a/cache/cache.go b/cache/cache.go index 45abf1a..c27e582 100644 --- a/cache/cache.go +++ b/cache/cache.go @@ -126,10 +126,7 @@ func (cache *Cache[V]) GetOrLoad( // Another caller may have populated the cache between the initial // lookup and this call becoming the singleflight owner. if cached, ok := cache.store.getAt(index, key, time.Now(), cache.stats.shard(index)); ok { - return loadResult[V]{ - value: cached.value, - found: cached.found, - }, nil + return loadResult[V](cached), nil } startedAt := time.Now() @@ -206,8 +203,9 @@ func (cache *Cache[V]) Invalidate(keys ...string) { if len(keys) == 1 { index := cache.store.segmentIndex(keys[0]) - cache.invalidateOne(index, keys[0]) - cache.stats.recordInvalidation(index) + if cache.invalidateOne(index, keys[0]) { + cache.stats.recordKeyInvalidation(index, 1) + } return } @@ -222,8 +220,8 @@ func (cache *Cache[V]) Invalidate(keys ...string) { } // Keep the statistics shard based on the caller's first key rather than - // the sorted target order. One non-empty Invalidate call contributes one - // operation regardless of how many keys it contains. + // the sorted target order. The shard records the total number of resident + // entries actually removed by this batch. statsIndex := targets[0].index // Every invalidation path acquires state locks in ascending segment order. @@ -262,12 +260,16 @@ func (cache *Cache[V]) Invalidate(keys ...string) { previous = target.index } + var invalidated int64 + for _, target := range targets { state := &cache.states[target.index] state.group.Forget(target.key) - cache.store.deleteAt(target.index, target.key) + if cache.store.deleteAt(target.index, target.key) { + invalidated++ + } } previous = -1 @@ -284,7 +286,7 @@ func (cache *Cache[V]) Invalidate(keys ...string) { previous = target.index } - cache.stats.recordInvalidation(statsIndex) + cache.stats.recordKeyInvalidation(statsIndex, invalidated) } func (cache *Cache[V]) InvalidateAll() { @@ -303,19 +305,19 @@ func (cache *Cache[V]) InvalidateAll() { state.group = &singleflight.Group{} } - cache.store.deleteAll() + invalidated := cache.store.deleteAll() for index := len(cache.states) - 1; index >= 0; index-- { cache.states[index].mu.Unlock() } - cache.stats.recordInvalidateAll() + cache.stats.recordAllInvalidation(invalidated) } func (cache *Cache[V]) invalidateOne( index int, key string, -) { +) bool { state := &cache.states[index] state.mu.Lock() @@ -323,9 +325,11 @@ func (cache *Cache[V]) invalidateOne( state.generation++ state.group.Forget(key) - cache.store.deleteAt(index, key) + removed := cache.store.deleteAt(index, key) state.mu.Unlock() + + return removed } func (cache *Cache[V]) storeLoaded( diff --git a/cache/stats.go b/cache/stats.go index fa18cbd..4939f48 100644 --- a/cache/stats.go +++ b/cache/stats.go @@ -61,11 +61,17 @@ type Stats struct { // Invalidation lifecycle. - // InvalidationCount is the cumulative number of non-empty Invalidate calls. - InvalidationCount int64 - - // InvalidateAllCount is the cumulative number of InvalidateAll calls. - InvalidateAllCount int64 + // InvalidatedKeyCount is the cumulative number of resident entries removed by + // Invalidate. Missing keys and duplicate keys that were already removed do not + // increase the count. + InvalidatedKeyCount int64 + + // InvalidatedAllCount is the cumulative number of resident entries removed by + // InvalidateAll. + // + // Because expiration is lazy, this may include physically resident entries + // whose TTL had already expired but which had not yet been accessed. + InvalidatedAllCount int64 // Storage lifecycle. @@ -82,11 +88,8 @@ type Stats struct { // // Each shard corresponds to one storage and coordination segment. type statsCollector struct { - shards []statsShard - - // InvalidateAll is cache-wide and expected to be rare, so keeping this - // counter global does not introduce meaningful contention. - invalidateAllCount atomic.Int64 + shards []statsShard + invalidatedAllCount atomic.Int64 } type statsShard struct { @@ -103,10 +106,10 @@ type statsShard struct { loadErrorCount atomic.Int64 loadDurationNanos atomic.Int64 - // These events are observed without a suitable existing lock. - // Counters are sharded to avoid a single global hot cache line. - sharedCount atomic.Int64 - invalidationCount atomic.Int64 + sharedCount atomic.Int64 + + // Sharded to avoid a global hot cache line. + invalidatedKeyCount atomic.Int64 } func newStatsCollector(segmentCount int) *statsCollector { @@ -122,9 +125,9 @@ func (cache *Cache[V]) Stats() Stats { } snapshot := Stats{ - MaxEntries: int64(cache.store.maxEntries), - SegmentCount: int64(len(cache.store.segments)), - InvalidateAllCount: cache.stats.invalidateAllCount.Load(), + MaxEntries: int64(cache.store.maxEntries), + SegmentCount: int64(len(cache.store.segments)), + InvalidatedAllCount: cache.stats.invalidatedAllCount.Load(), } for index := range cache.store.segments { @@ -149,7 +152,7 @@ func (cache *Cache[V]) Stats() Stats { snapshot.LoadErrorCount += shard.loadErrorCount.Load() snapshot.LoadDuration += time.Duration(shard.loadDurationNanos.Load()) snapshot.SharedCount += shard.sharedCount.Load() - snapshot.InvalidationCount += shard.invalidationCount.Load() + snapshot.InvalidatedKeyCount += shard.invalidatedKeyCount.Load() } return snapshot @@ -180,10 +183,18 @@ func (stats *statsCollector) recordShared(index int) { stats.shard(index).sharedCount.Add(1) } -func (stats *statsCollector) recordInvalidation(index int) { - stats.shard(index).invalidationCount.Add(1) +func (stats *statsCollector) recordKeyInvalidation(index int, count int64) { + if count <= 0 { + return + } + + stats.shard(index).invalidatedKeyCount.Add(count) } -func (stats *statsCollector) recordInvalidateAll() { - stats.invalidateAllCount.Add(1) +func (stats *statsCollector) recordAllInvalidation(count int64) { + if count <= 0 { + return + } + + stats.invalidatedAllCount.Add(count) } diff --git a/cache/storage.go b/cache/storage.go index 700a330..c53c779 100644 --- a/cache/storage.go +++ b/cache/storage.go @@ -86,13 +86,6 @@ func newStorageWithSegments[V any](maxEntries, segmentCount int) *storage[V] { } } -func (storage *storage[V]) get( - key string, - now time.Time, -) (cachedValue[V], bool) { - return storage.getAt(storage.segmentIndex(key), key, now, nil) -} - func (storage *storage[V]) lookupAt( index int, key string, @@ -111,14 +104,6 @@ func (storage *storage[V]) getAt( return storage.segments[index].get(key, now, stats) } -func (storage *storage[V]) set( - key string, - value cachedValue[V], - expiresAt time.Time, -) { - storage.setAt(storage.segmentIndex(key), key, value, expiresAt, nil) -} - func (storage *storage[V]) setAt( index int, key string, @@ -129,21 +114,21 @@ func (storage *storage[V]) setAt( storage.segments[index].set(key, value, expiresAt, stats) } -func (storage *storage[V]) delete(key string) { - storage.deleteAt(storage.segmentIndex(key), key) -} - func (storage *storage[V]) deleteAt( index int, key string, -) { - storage.segments[index].delete(key) +) bool { + return storage.segments[index].delete(key) } -func (storage *storage[V]) deleteAll() { +func (storage *storage[V]) deleteAll() int64 { + var deleted int64 + for index := range storage.segments { - storage.segments[index].deleteAll() + deleted += storage.segments[index].deleteAll() } + + return deleted } func (storage *storage[V]) segmentIndex(key string) int { @@ -289,26 +274,32 @@ func (segment *storageSegment[V]) set( segment.pushFrontLocked(item) } -func (segment *storageSegment[V]) delete(key string) { +func (segment *storageSegment[V]) delete(key string) bool { segment.mu.Lock() defer segment.mu.Unlock() item, ok := segment.entries[key] if !ok { - return + return false } segment.removeLocked(item) + + return true } -func (segment *storageSegment[V]) deleteAll() { +func (segment *storageSegment[V]) deleteAll() int64 { segment.mu.Lock() defer segment.mu.Unlock() + deleted := int64(len(segment.entries)) + clear(segment.entries) segment.head = nil segment.tail = nil + + return deleted } func (segment *storageSegment[V]) removeLocked(item *entry[V]) { diff --git a/metrics/otel/cache.go b/metrics/otel/cache.go index 3b7e8e9..73ee273 100644 --- a/metrics/otel/cache.go +++ b/metrics/otel/cache.go @@ -191,12 +191,12 @@ func (instruments cacheMetricInstruments) observe( ) observer.ObserveInt64( instruments.invalidationCount, - stats.InvalidationCount, + stats.InvalidatedKeyCount, attributes.invalidateKeys, ) observer.ObserveInt64( instruments.invalidationCount, - stats.InvalidateAllCount, + stats.InvalidatedAllCount, attributes.invalidateAll, ) observer.ObserveInt64( @@ -339,9 +339,9 @@ func newCacheMetricInstruments(meter metric.Meter) (cacheMetricInstruments, erro instruments.invalidationCount, err = meter.Int64ObservableCounter( cacheInvalidationMetricName, metric.WithDescription( - "The cumulative number of explicit cache invalidation operations by scope.", + "The cumulative number of resident cache entries removed by explicit invalidation, by scope.", ), - metric.WithUnit("{operation}"), + metric.WithUnit("{entry}"), ) if err != nil { return cacheMetricInstruments{}, fmt.Errorf( From 4b8178b629be73971c113d2d9dbc3c0ef42b7da6 Mon Sep 17 00:00:00 2001 From: mkbeh Date: Sun, 23 Aug 2026 13:54:52 +0300 Subject: [PATCH 16/41] refactor: remove cache and move OpenTelemetry integration --- .github/dependabot.yml | 24 + .github/scripts/create-go-workspace.sh | 87 ++++ .github/workflows/codecov.yml | 60 +++ .github/workflows/lint.yml | 93 ++++ .github/workflows/test.yml | 107 +++++ Taskfile.yml | 144 +++++- cache/cache.go | 410 ---------------- cache/config.go | 103 ---- cache/doc.go | 13 - cache/loader.go | 21 - cache/metrics.go | 22 - cache/stats.go | 200 -------- cache/storage.go | 348 -------------- examples/otel/go.mod | 1 + examples/otel/main.go | 4 +- {metrics/otel => extra/otelxpg}/go.mod | 2 +- {metrics/otel => extra/otelxpg}/metrics.go | 11 +- {metrics/otel => extra/otelxpg}/options.go | 6 +- .../otel => extra/otelxpg}/options_test.go | 59 +-- {metrics/otel => extra/otelxpg}/pool.go | 26 +- go.mod | 6 +- metrics/otel/cache.go | 451 ------------------ 22 files changed, 548 insertions(+), 1650 deletions(-) create mode 100644 .github/dependabot.yml create mode 100755 .github/scripts/create-go-workspace.sh create mode 100644 .github/workflows/codecov.yml create mode 100644 .github/workflows/lint.yml create mode 100644 .github/workflows/test.yml delete mode 100644 cache/cache.go delete mode 100644 cache/config.go delete mode 100644 cache/doc.go delete mode 100644 cache/loader.go delete mode 100644 cache/metrics.go delete mode 100644 cache/stats.go delete mode 100644 cache/storage.go rename {metrics/otel => extra/otelxpg}/go.mod (75%) rename {metrics/otel => extra/otelxpg}/metrics.go (69%) rename {metrics/otel => extra/otelxpg}/options.go (90%) rename {metrics/otel => extra/otelxpg}/options_test.go (59%) rename {metrics/otel => extra/otelxpg}/pool.go (95%) delete mode 100644 metrics/otel/cache.go diff --git a/.github/dependabot.yml b/.github/dependabot.yml new file mode 100644 index 0000000..a4f8b0f --- /dev/null +++ b/.github/dependabot.yml @@ -0,0 +1,24 @@ +version: 2 + +updates: + - package-ecosystem: "gomod" + directories: + - "/" + - "/extra/*" + - "/examples/*" + schedule: + interval: "weekly" + ignore: + # Local modules are resolved through go.work during repository builds. + - dependency-name: "github.com/mkbeh/xpg" + - dependency-name: "github.com/mkbeh/xpg/*" + groups: + go-dependencies: + patterns: + - "*" + group-by: dependency-name + + - package-ecosystem: "github-actions" + directory: "/" + schedule: + interval: "weekly" diff --git a/.github/scripts/create-go-workspace.sh b/.github/scripts/create-go-workspace.sh new file mode 100755 index 0000000..acca8ac --- /dev/null +++ b/.github/scripts/create-go-workspace.sh @@ -0,0 +1,87 @@ +#!/usr/bin/env bash +set -euo pipefail + +if [[ -n "${GITHUB_WORKSPACE:-}" ]]; then + root="${GITHUB_WORKSPACE}" +else + root="$(git rev-parse --show-toplevel 2>/dev/null || pwd)" +fi + +include_examples=false + +for arg in "$@"; do + case "${arg}" in + --with-examples) + include_examples=true + ;; + *) + echo "unknown argument: ${arg}" >&2 + exit 2 + ;; + esac +done + +cd "${root}" + +rm -f go.work go.work.sum + +modules=(.) + +while IFS= read -r -d '' mod; do + modules+=("$(dirname "${mod}")") +done < <(find ./extra -name go.mod -type f -print0 2>/dev/null | sort -z) + +if [[ "${include_examples}" == "true" ]]; then + while IFS= read -r -d '' mod; do + modules+=("$(dirname "${mod}")") + done < <(find ./examples -name go.mod -type f -print0 2>/dev/null | sort -z) +fi + +# The workspace is ephemeral in CI and local-only for development. +GOWORK=off go work init "${modules[@]}" + +# Local modules may already require release versions such as v0.1.0 while the +# corresponding tags do not exist yet. Bind those exact requirements to the +# current checkout so all Go tooling, including go/packages-based linters, +# resolves the unpublished modules locally. +while read -r module version; do + [[ -n "${module}" && -n "${version}" ]] || continue + + case "${module}" in + github.com/mkbeh/xpg) + local_dir="." + ;; + github.com/mkbeh/xpg/*) + local_dir="./${module#github.com/mkbeh/xpg/}" + ;; + *) + continue + ;; + esac + + if [[ -f "${root}/${local_dir#./}/go.mod" ]]; then + GOWORK="${root}/go.work" go work edit \ + -replace="${module}@${version}=${local_dir}" + fi +done < <( + { + printf '%s\0' ./go.mod + + find ./extra ./examples \ + -name go.mod -type f -print0 2>/dev/null || true + } | + while IFS= read -r -d '' mod; do + awk ' + $1 ~ /^github\.com\/mkbeh\/xpg(\/.*)?$/ && $2 ~ /^v[0-9]/ { + print $1, $2 + } + ' "${mod}" + done | + sort -u +) + +echo "Created temporary Go workspace at ${root}/go.work" + +if [[ "${WORKSPACE_DEBUG:-false}" == "true" ]]; then + GOWORK="${root}/go.work" go work edit -json +fi diff --git a/.github/workflows/codecov.yml b/.github/workflows/codecov.yml new file mode 100644 index 0000000..d0c3d86 --- /dev/null +++ b/.github/workflows/codecov.yml @@ -0,0 +1,60 @@ +name: Coverage + +on: + push: + branches: + - main + pull_request: + branches: + - main + workflow_dispatch: + +concurrency: + group: ${{ github.workflow }}-${{ github.ref }} + cancel-in-progress: true + +permissions: + contents: read + +env: + GOTOOLCHAIN: local + CODECOV_TOKEN: ${{ secrets.CODECOV_TOKEN }} + +jobs: + coverage: + name: Coverage + runs-on: ubuntu-latest + timeout-minutes: 15 + + steps: + - name: Check out repository + uses: actions/checkout@v7 + + - name: Set up Go + uses: actions/setup-go@v7 + with: + go-version-file: go.mod + check-latest: true + cache-dependency-path: | + go.sum + extra/**/go.mod + + - name: Set up Task + uses: go-task/setup-task@v2 + + - name: Create temporary Go workspace + run: bash .github/scripts/create-go-workspace.sh + + - name: Run coverage + env: + GOWORK: ${{ github.workspace }}/go.work + run: task coverage + + - name: Upload coverage to Codecov + if: env.CODECOV_TOKEN != '' + uses: codecov/codecov-action@v7 + with: + token: ${{ env.CODECOV_TOKEN }} + files: ./coverage.out + disable_search: true + fail_ci_if_error: true diff --git a/.github/workflows/lint.yml b/.github/workflows/lint.yml new file mode 100644 index 0000000..e02f372 --- /dev/null +++ b/.github/workflows/lint.yml @@ -0,0 +1,93 @@ +name: Lint + +on: + push: + branches: + - main + pull_request: + branches: + - main + workflow_dispatch: + +concurrency: + group: ${{ github.workflow }}-${{ github.ref }} + cancel-in-progress: true + +permissions: + contents: read + +env: + GOTOOLCHAIN: local + +jobs: + lint: + name: Lint (${{ matrix.module }}) + runs-on: ubuntu-latest + timeout-minutes: 15 + + strategy: + fail-fast: false + matrix: + module: + - . + - extra/otelxpg + - examples/advisory + - examples/basic + - examples/cluster + - examples/otel + - examples/shard + - examples/transactions + + steps: + - name: Check out repository + uses: actions/checkout@v7 + + - name: Set up Go + uses: actions/setup-go@v7 + with: + go-version-file: go.mod + check-latest: true + cache-dependency-path: | + go.sum + extra/**/go.mod + examples/**/go.mod + + - name: Create temporary Go workspace + run: bash .github/scripts/create-go-workspace.sh --with-examples + + - name: Verify workspace resolution + env: + GOWORK: ${{ github.workspace }}/go.work + run: | + set -euo pipefail + + test "$(go env GOWORK)" = "${GITHUB_WORKSPACE}/go.work" + + go list -m \ + -f '{{if .Main}}{{.Path}} => {{.Dir}}{{end}}' \ + all | + sed '/^$/d' + + - name: Run golangci-lint + uses: golangci/golangci-lint-action@v9 + env: + GOWORK: ${{ github.workspace }}/go.work + with: + version: v2.12.2 + working-directory: ${{ matrix.module }} + args: >- + --config=${{ github.workspace }}/.golangci.yml + --timeout=5m + --modules-download-mode=readonly + + result: + name: Lint + if: always() + needs: + - lint + runs-on: ubuntu-latest + timeout-minutes: 5 + + steps: + - name: Check lint jobs + run: test "${{ needs.lint.result }}" = "success" diff --git a/.github/workflows/test.yml b/.github/workflows/test.yml new file mode 100644 index 0000000..70af09a --- /dev/null +++ b/.github/workflows/test.yml @@ -0,0 +1,107 @@ +name: Test + +on: + push: + branches: + - main + pull_request: + branches: + - main + workflow_dispatch: + +concurrency: + group: ${{ github.workflow }}-${{ github.ref }} + cancel-in-progress: true + +permissions: + contents: read + +env: + GOTOOLCHAIN: local + +jobs: + test: + name: Tests + runs-on: ubuntu-latest + timeout-minutes: 15 + + steps: + - name: Check out repository + uses: actions/checkout@v7 + + - name: Set up Go + uses: actions/setup-go@v7 + with: + go-version-file: go.mod + check-latest: true + cache-dependency-path: | + go.sum + extra/**/go.mod + examples/**/go.mod + + - name: Set up Task + uses: go-task/setup-task@v2 + + - name: Create temporary Go workspace + run: bash .github/scripts/create-go-workspace.sh --with-examples + + - name: Verify workspace resolution + env: + GOWORK: ${{ github.workspace }}/go.work + run: | + set -euo pipefail + + test "$(go env GOWORK)" = "${GITHUB_WORKSPACE}/go.work" + + go list -m \ + -f '{{if .Main}}{{.Path}} => {{.Dir}}{{end}}' \ + all | + sed '/^$/d' + + - name: Run tests + env: + GOWORK: ${{ github.workspace }}/go.work + run: task test + + race: + name: Race + runs-on: ubuntu-latest + timeout-minutes: 20 + + steps: + - name: Check out repository + uses: actions/checkout@v7 + + - name: Set up Go + uses: actions/setup-go@v7 + with: + go-version-file: go.mod + check-latest: true + cache-dependency-path: | + go.sum + extra/**/go.mod + examples/**/go.mod + + - name: Set up Task + uses: go-task/setup-task@v2 + + - name: Create temporary Go workspace + run: bash .github/scripts/create-go-workspace.sh --with-examples + + - name: Verify workspace resolution + env: + GOWORK: ${{ github.workspace }}/go.work + run: | + set -euo pipefail + + test "$(go env GOWORK)" = "${GITHUB_WORKSPACE}/go.work" + + go list -m \ + -f '{{if .Main}}{{.Path}} => {{.Dir}}{{end}}' \ + all | + sed '/^$/d' + + - name: Run race tests + env: + GOWORK: ${{ github.workspace }}/go.work + run: task test-race diff --git a/Taskfile.yml b/Taskfile.yml index e1d4358..12b46ca 100755 --- a/Taskfile.yml +++ b/Taskfile.yml @@ -9,47 +9,163 @@ vars: GOCACHE: sh: go env GOCACHE - PWD_DIR: - sh: pwd + GOSUMCACHE: + sh: echo "$(go env GOPATH)/pkg/sumdb" + + # Root module + all local submodules that are part of the project. + # Examples are included because they depend on unpublished local modules. + MODULE_DIRS: + sh: | + { + printf '.\n' + + find ./extra ./examples -name go.mod -print | + while IFS= read -r mod; do + dirname "$mod" + done + } | sort -u + + # Coverage is aggregated for production modules only. Example modules are + # still linted and tested, but do not affect the library coverage metric. + COVERAGE_MODULE_DIRS: + sh: | + { + printf '.\n' + + find ./extra -name go.mod -print | + while IFS= read -r mod; do + dirname "$mod" + done + } | sort -u | tr '\n' ' ' env: - GOWORK: "off" + # The workspace resolves unpublished local modules, including examples, + # against the current checkout instead of released module versions. + GOWORK: "{{.ROOT_DIR}}/go.work" tasks: + init-workspace: + desc: "Initialize or reset the local Go workspace" + cmds: + - bash .github/scripts/create-go-workspace.sh --with-examples + run: - desc: "Run all checks" + desc: "Run lint and race tests for all modules" cmds: - task: lint - task: test-race lint: - desc: "Run golangci-lint" + desc: "Lint all modules" cmds: - - task: _run-linter + - for: { var: MODULE_DIRS } + task: _run-linter + vars: + WORK_DIR: "/app/{{.ITEM}}" test: - desc: "Run tests" + desc: "Run tests for all modules" cmds: - - go test -v -count=1 ./... + - for: { var: MODULE_DIRS } + task: _run-tests + vars: + WORK_DIR: "{{.ITEM}}" test-race: - desc: "Run tests with the race detector" + desc: "Run tests with the race detector for all modules" + cmds: + - for: { var: MODULE_DIRS } + task: _run-race-tests + vars: + WORK_DIR: "{{.ITEM}}" + + coverage: + desc: "Run tests and print total coverage for production modules" + cmds: + - | + set -eu + + root="{{.ROOT_DIR}}" + output="$root/coverage.out" + profiles="$(mktemp -d)" + + cleanup() { + rm -rf "$profiles" + } + trap cleanup EXIT + + printf 'mode: atomic\n' > "$output" + + profile_index=0 + + for module in {{.COVERAGE_MODULE_DIRS}}; do + echo "==> $module" + + profile="$profiles/$profile_index.out" + profile_index=$((profile_index + 1)) + + ( + cd "$root/$module" + + GOWORK="$root/go.work" go test \ + -mod=readonly \ + -count=1 \ + -covermode=atomic \ + -coverprofile="$profile" \ + ./... + ) + + tail -n +2 "$profile" >> "$output" + done + + GOWORK="$root/go.work" go tool cover -func="$output" | grep total + + coverage:html: + desc: "Generate HTML coverage report" + deps: + - coverage + cmds: + - go tool cover -html=coverage.out -o coverage.html + + _run-tests: + internal: true + requires: + vars: + - WORK_DIR + dir: "{{.ROOT_DIR}}/{{.WORK_DIR}}" + cmds: + - echo "==> {{.WORK_DIR}}" + - go test -mod=readonly -v -count=1 ./... + + _run-race-tests: + internal: true + requires: + vars: + - WORK_DIR + dir: "{{.ROOT_DIR}}/{{.WORK_DIR}}" cmds: - - go test -race -v -count=1 ./... + - echo "==> {{.WORK_DIR}}" + - go test -mod=readonly -race -v -count=1 ./... _run-linter: internal: true + requires: + vars: + - WORK_DIR cmds: + - echo "==> {{.WORK_DIR}}" + - mkdir -p "{{.GOSUMCACHE}}" - > docker run --rm -u $(id -u):$(id -g) - -v {{.PWD_DIR}}:/app + -v {{.ROOT_DIR}}:/app -v {{.GOMODCACHE}}:/go/pkg/mod + -v {{.GOSUMCACHE}}:/go/pkg/sumdb -v {{.GOCACHE}}:/go/build-cache -e GOMODCACHE=/go/pkg/mod -e GOCACHE=/go/build-cache -e GOLANGCI_LINT_CACHE=/go/build-cache/golangci-lint - -e GOWORK=off - -w /app + -e GOWORK=/app/go.work + -w {{.WORK_DIR}} golangci/golangci-lint:{{.LINTER_VER}} - golangci-lint run --timeout 5m \ No newline at end of file + golangci-lint run --timeout 5m --modules-download-mode readonly diff --git a/cache/cache.go b/cache/cache.go deleted file mode 100644 index c27e582..0000000 --- a/cache/cache.go +++ /dev/null @@ -1,410 +0,0 @@ -package cache - -import ( - "cmp" - "context" - "errors" - "fmt" - "math/rand/v2" - "slices" - "sync" - "time" - - "golang.org/x/sync/singleflight" -) - -type Cache[V any] struct { - name string - - store *storage[V] - states []cacheState - stats *statsCollector - - metrics MetricsRegistration - closeOnce sync.Once - - ttl time.Duration - jitter time.Duration - negativeTTL time.Duration -} - -type cacheState struct { - mu sync.RWMutex - - generation uint64 - group *singleflight.Group -} - -type invalidationTarget struct { - key string - index int -} - -func New[V any](config Config) (*Cache[V], error) { - if err := config.validate(); err != nil { - return nil, err - } - - store := newStorage[V](config.MaxEntries, config.Segments) - - cache := &Cache[V]{ - name: config.Name, - - store: store, - states: newCacheStates(len(store.segments)), - stats: newStatsCollector(len(store.segments)), - - ttl: config.TTL, - jitter: config.Jitter, - negativeTTL: config.NegativeTTL, - } - - if err := cache.registerMetrics(config.Metrics); err != nil { - return nil, fmt.Errorf("xpg/cache: register metrics: %w", err) - } - - return cache, nil -} - -func (cache *Cache[V]) Name() string { - if cache == nil { - return "" - } - - return cache.name -} - -func (cache *Cache[V]) Close() { - if cache == nil { - return - } - - cache.closeOnce.Do( - func() { - if cache.metrics != nil { - cache.metrics.Close() - } - }, - ) -} - -func (cache *Cache[V]) GetOrLoad( - ctx context.Context, - key string, - loader Loader[V], -) (V, bool, error) { - var zero V - - if !cache.initialized() { - return zero, false, errors.New("xpg/cache: cache is not initialized") - } - - if loader == nil { - return zero, false, errors.New("xpg/cache: loader is nil") - } - - index := cache.store.segmentIndex(key) - stats := cache.stats.shard(index) - - if cached, ok := cache.store.lookupAt(index, key, time.Now(), stats); ok { - return cached.value, cached.found, nil - } - - state := &cache.states[index] - - // Registration and generation selection must be atomic with respect to - // invalidation for this state segment. DoChan only registers or starts the - // shared call; the loader itself executes outside state.mu. - state.mu.RLock() - - generation := state.generation - group := state.group - - resultChannel := group.DoChan( - key, - func() (any, error) { - // Another caller may have populated the cache between the initial - // lookup and this call becoming the singleflight owner. - if cached, ok := cache.store.getAt(index, key, time.Now(), cache.stats.shard(index)); ok { - return loadResult[V](cached), nil - } - - startedAt := time.Now() - - value, found, err := loader(ctx) - - finishedAt := time.Now() - - cache.stats.recordLoad(index, found, err, finishedAt.Sub(startedAt)) - - if err != nil { - return nil, err - } - - if !found { - var zeroValue V - - value = zeroValue - } - - loaded := loadResult[V]{ - value: value, - found: found, - } - - // Generation validation and publication are atomic with respect to - // invalidation for this state segment. A pre-invalidation load may - // still return to callers that already joined it, but cannot - // repopulate the cache after the barrier. - publishState := &cache.states[index] - - publishState.mu.RLock() - - if publishState.generation == generation { - cache.storeLoaded(index, key, loaded, finishedAt) - } - - publishState.mu.RUnlock() - - return loaded, nil - }, - ) - - state.mu.RUnlock() - - select { - case <-ctx.Done(): - return zero, false, ctx.Err() - - case result := <-resultChannel: - if result.Shared { - cache.stats.recordShared(index) - } - - if result.Err != nil { - return zero, false, result.Err - } - - loaded, ok := result.Val.(loadResult[V]) - if !ok { - return zero, false, errors.New("xpg/cache: unexpected singleflight result type") - } - - return loaded.value, loaded.found, nil - } -} - -func (cache *Cache[V]) Invalidate(keys ...string) { - if !cache.initialized() || - len(keys) == 0 { - return - } - - if len(keys) == 1 { - index := cache.store.segmentIndex(keys[0]) - - if cache.invalidateOne(index, keys[0]) { - cache.stats.recordKeyInvalidation(index, 1) - } - - return - } - - targets := make([]invalidationTarget, len(keys)) - - for index, key := range keys { - targets[index] = invalidationTarget{ - key: key, - index: cache.store.segmentIndex(key), - } - } - - // Keep the statistics shard based on the caller's first key rather than - // the sorted target order. The shard records the total number of resident - // entries actually removed by this batch. - statsIndex := targets[0].index - - // Every invalidation path acquires state locks in ascending segment order. - // This keeps multi-key invalidation and InvalidateAll deadlock-free. - slices.SortFunc( - targets, - func(left, right invalidationTarget) int { - return cmp.Compare(left.index, right.index) - }, - ) - - previous := -1 - - for _, target := range targets { - if target.index == previous { - continue - } - - cache.states[target.index].mu.Lock() - - previous = target.index - } - - // Once every affected state is locked, advance each generation. Loads that - // registered before this barrier may finish for existing waiters, but they - // cannot publish into any affected segment afterward. - previous = -1 - - for _, target := range targets { - if target.index == previous { - continue - } - - cache.states[target.index].generation++ - - previous = target.index - } - - var invalidated int64 - - for _, target := range targets { - state := &cache.states[target.index] - - state.group.Forget(target.key) - - if cache.store.deleteAt(target.index, target.key) { - invalidated++ - } - } - - previous = -1 - - for index := len(targets) - 1; index >= 0; index-- { - target := targets[index] - - if target.index == previous { - continue - } - - cache.states[target.index].mu.Unlock() - - previous = target.index - } - - cache.stats.recordKeyInvalidation(statsIndex, invalidated) -} - -func (cache *Cache[V]) InvalidateAll() { - if !cache.initialized() { - return - } - - for index := range cache.states { - cache.states[index].mu.Lock() - } - - for index := range cache.states { - state := &cache.states[index] - - state.generation++ - state.group = &singleflight.Group{} - } - - invalidated := cache.store.deleteAll() - - for index := len(cache.states) - 1; index >= 0; index-- { - cache.states[index].mu.Unlock() - } - - cache.stats.recordAllInvalidation(invalidated) -} - -func (cache *Cache[V]) invalidateOne( - index int, - key string, -) bool { - state := &cache.states[index] - - state.mu.Lock() - - state.generation++ - state.group.Forget(key) - - removed := cache.store.deleteAt(index, key) - - state.mu.Unlock() - - return removed -} - -func (cache *Cache[V]) storeLoaded( - index int, - key string, - loaded loadResult[V], - now time.Time, -) { - switch { - case loaded.found: - cache.store.setAt( - index, - key, - cachedValue[V]{ - value: loaded.value, - found: true, - }, - now.Add(cache.effectiveTTL()), - cache.stats.shard(index), - ) - - case cache.negativeTTL > 0: - cache.store.setAt( - index, - key, - cachedValue[V]{ - found: false, - }, - now.Add(cache.negativeTTL), - cache.stats.shard(index), - ) - } -} - -func (cache *Cache[V]) effectiveTTL() time.Duration { - if cache.jitter == 0 { - return cache.ttl - } - - return cache.ttl + time.Duration( - rand.Int64N(int64(cache.jitter)), - ) -} - -func (cache *Cache[V]) registerMetrics(metrics Metrics) error { - if metrics == nil { - return nil - } - - registration, err := metrics.RegisterCache(cache) - if err != nil { - return err - } - - cache.metrics = registration - - return nil -} - -func (cache *Cache[V]) initialized() bool { - return cache != nil && - cache.store != nil && - cache.stats != nil && - len(cache.states) == - len(cache.store.segments) && - len(cache.stats.shards) == - len(cache.store.segments) -} - -func newCacheStates(count int) []cacheState { - states := make([]cacheState, count) - - for index := range states { - states[index].group = &singleflight.Group{} - } - - return states -} diff --git a/cache/config.go b/cache/config.go deleted file mode 100644 index 12d51a1..0000000 --- a/cache/config.go +++ /dev/null @@ -1,103 +0,0 @@ -package cache - -import ( - "errors" - "math" - "strings" - "time" -) - -const ( - maxDuration = time.Duration(math.MaxInt64) - defaultMaxEntries = 10_000 -) - -// Config configures a bounded in-process cache. -type Config struct { - // Name identifies the cache for diagnostics and observability. - Name string - - // MaxEntries is the maximum total entry budget of the cache. - // - // The budget is distributed across storage segments. Because each segment - // enforces its capacity independently, the number of simultaneously resident - // entries may be slightly lower than MaxEntries depending on key distribution. - // - // Zero uses the default. - MaxEntries int - - // Segments is the number of independent storage segments used to reduce lock - // contention. - // - // Zero uses the default segment count. - Segments int - - // TTL is the lifetime of positive entries. - TTL time.Duration - - // Jitter adds a random duration in [0, Jitter) to positive-entry TTLs. - // It can be used to spread expiration of entries loaded around the same time. - Jitter time.Duration - - // NegativeTTL is the lifetime of cached negative results. - // Zero disables negative caching. - NegativeTTL time.Duration - - // Metrics optionally registers cache statistics during New. - // The registration is released when Cache.Close is called. - Metrics Metrics -} - -func (config Config) validate() error { - name := strings.TrimSpace(config.Name) - - if name == "" { - return errors.New( - "xpg/cache: name must not be blank", - ) - } - - if name != config.Name { - return errors.New( - "xpg/cache: name must not contain surrounding whitespace", - ) - } - - if config.MaxEntries < 0 { - return errors.New( - "xpg/cache: max entries must not be negative", - ) - } - - if config.Segments < 0 { - return errors.New( - "xpg/cache: segments must not be negative", - ) - } - - if config.TTL <= 0 { - return errors.New( - "xpg/cache: ttl must be greater than zero", - ) - } - - if config.Jitter < 0 { - return errors.New( - "xpg/cache: jitter must not be negative", - ) - } - - if config.NegativeTTL < 0 { - return errors.New( - "xpg/cache: negative ttl must not be negative", - ) - } - - if config.Jitter > maxDuration-config.TTL { - return errors.New( - "xpg/cache: ttl and jitter overflow time.Duration", - ) - } - - return nil -} diff --git a/cache/doc.go b/cache/doc.go deleted file mode 100644 index af8342b..0000000 --- a/cache/doc.go +++ /dev/null @@ -1,13 +0,0 @@ -// Package cache provides bounded in-process read-through caching. -// -// Cache entries use absolute TTL expiration, optional TTL jitter, LRU -// eviction, negative caching, duplicate load suppression, and explicit -// invalidation. -// -// Cache statistics are collected locally and exposed through Cache.Stats. -// Optional metrics integrations register during New and observe those snapshots -// without adding telemetry calls to the cache request path. -// -// The cache is local to one application process. It does not provide -// distributed cache coherence between application instances. -package cache diff --git a/cache/loader.go b/cache/loader.go deleted file mode 100644 index 244ded0..0000000 --- a/cache/loader.go +++ /dev/null @@ -1,21 +0,0 @@ -package cache - -import "context" - -// Loader loads one value. -// -// found=false represents a successful negative result. Negative results are -// cached only when Config.NegativeTTL is greater than zero. -// -// Loader errors are never cached. -type Loader[V any] func(ctx context.Context) (value V, found bool, err error) - -type loadResult[V any] struct { - value V - found bool -} - -type cachedValue[V any] struct { - value V - found bool -} diff --git a/cache/metrics.go b/cache/metrics.go deleted file mode 100644 index 05145e3..0000000 --- a/cache/metrics.go +++ /dev/null @@ -1,22 +0,0 @@ -package cache - -// StatsProvider exposes one cache's identity and statistics to a metrics -// implementation. Cache implements StatsProvider. -type StatsProvider interface { - Name() string - Stats() Stats -} - -// Metrics registers observability for a Cache. -// -// Implementations are expected to be immutable and safe to reuse for multiple -// caches. RegisterCache is called after the cache has been fully initialized. -type Metrics interface { - RegisterCache(cache StatsProvider) (MetricsRegistration, error) -} - -// MetricsRegistration owns a metrics registration associated with one Cache. -// Close is called once when the Cache is closed. -type MetricsRegistration interface { - Close() -} diff --git a/cache/stats.go b/cache/stats.go deleted file mode 100644 index 4939f48..0000000 --- a/cache/stats.go +++ /dev/null @@ -1,200 +0,0 @@ -package cache - -import ( - "sync/atomic" - "time" -) - -// Stats is a detached snapshot of cache statistics. -// -// The snapshot is assembled from independent cache segments. Concurrent cache -// activity may continue while Stats is being collected, so individual fields -// are not guaranteed to represent one globally atomic instant. -type Stats struct { - // Current state. - - // EntryCount is the number of entries currently resident in storage. - // Because expiration is lazy, this count may include expired entries that - // have not been accessed or evicted yet. - EntryCount int64 - - // MaxEntries is the configured total entry budget after applying defaults. - MaxEntries int64 - - // SegmentCount is the number of independent storage and coordination - // segments. - SegmentCount int64 - - // Lookup lifecycle. - - // HitCount is the cumulative number of positive cache hits. - HitCount int64 - - // NegativeHitCount is the cumulative number of cached negative hits. - NegativeHitCount int64 - - // MissCount is the cumulative number of cache misses. An expired entry - // observed by a caller is counted as a miss. - MissCount int64 - - // Load lifecycle. - - // LoadFoundCount is the cumulative number of actual loader invocations that - // returned found=true. - LoadFoundCount int64 - - // LoadNotFoundCount is the cumulative number of actual loader invocations - // that returned found=false without an error. - LoadNotFoundCount int64 - - // LoadErrorCount is the cumulative number of actual loader invocations that - // returned an error. - LoadErrorCount int64 - - // LoadDuration is the cumulative duration of actual loader invocations, - // including successful, negative, and failed loads. - LoadDuration time.Duration - - // SharedCount is the cumulative number of callers that received a - // singleflight result shared with at least one other caller. - SharedCount int64 - - // Invalidation lifecycle. - - // InvalidatedKeyCount is the cumulative number of resident entries removed by - // Invalidate. Missing keys and duplicate keys that were already removed do not - // increase the count. - InvalidatedKeyCount int64 - - // InvalidatedAllCount is the cumulative number of resident entries removed by - // InvalidateAll. - // - // Because expiration is lazy, this may include physically resident entries - // whose TTL had already expired but which had not yet been accessed. - InvalidatedAllCount int64 - - // Storage lifecycle. - - // EvictionCount is the cumulative number of entries evicted because a - // storage segment reached capacity. - EvictionCount int64 - - // ExpirationCount is the cumulative number of expired entries removed - // lazily while looking up a key. - ExpirationCount int64 -} - -// statsCollector owns all mutable cache statistics. -// -// Each shard corresponds to one storage and coordination segment. -type statsCollector struct { - shards []statsShard - invalidatedAllCount atomic.Int64 -} - -type statsShard struct { - // Protected by the corresponding storage segment mutex. - hitCount int64 - negativeHitCount int64 - missCount int64 - evictionCount int64 - expirationCount int64 - - // Updated outside a suitable existing lock. - loadFoundCount atomic.Int64 - loadNotFoundCount atomic.Int64 - loadErrorCount atomic.Int64 - loadDurationNanos atomic.Int64 - - sharedCount atomic.Int64 - - // Sharded to avoid a global hot cache line. - invalidatedKeyCount atomic.Int64 -} - -func newStatsCollector(segmentCount int) *statsCollector { - return &statsCollector{ - shards: make([]statsShard, segmentCount), - } -} - -// Stats returns a detached snapshot of the current cache statistics. -func (cache *Cache[V]) Stats() Stats { - if !cache.initialized() { - return Stats{} - } - - snapshot := Stats{ - MaxEntries: int64(cache.store.maxEntries), - SegmentCount: int64(len(cache.store.segments)), - InvalidatedAllCount: cache.stats.invalidatedAllCount.Load(), - } - - for index := range cache.store.segments { - segment := &cache.store.segments[index] - shard := cache.stats.shard(index) - - // Storage counters share the storage segment mutex with the operations - // that update them. - segment.mu.Lock() - - snapshot.EntryCount += int64(len(segment.entries)) - snapshot.HitCount += shard.hitCount - snapshot.NegativeHitCount += shard.negativeHitCount - snapshot.MissCount += shard.missCount - snapshot.EvictionCount += shard.evictionCount - snapshot.ExpirationCount += shard.expirationCount - - segment.mu.Unlock() - - snapshot.LoadFoundCount += shard.loadFoundCount.Load() - snapshot.LoadNotFoundCount += shard.loadNotFoundCount.Load() - snapshot.LoadErrorCount += shard.loadErrorCount.Load() - snapshot.LoadDuration += time.Duration(shard.loadDurationNanos.Load()) - snapshot.SharedCount += shard.sharedCount.Load() - snapshot.InvalidatedKeyCount += shard.invalidatedKeyCount.Load() - } - - return snapshot -} - -func (stats *statsCollector) shard(index int) *statsShard { - return &stats.shards[index] -} - -func (stats *statsCollector) recordLoad(index int, found bool, err error, duration time.Duration) { - shard := stats.shard(index) - - shard.loadDurationNanos.Add(duration.Nanoseconds()) - - switch { - case err != nil: - shard.loadErrorCount.Add(1) - - case found: - shard.loadFoundCount.Add(1) - - default: - shard.loadNotFoundCount.Add(1) - } -} - -func (stats *statsCollector) recordShared(index int) { - stats.shard(index).sharedCount.Add(1) -} - -func (stats *statsCollector) recordKeyInvalidation(index int, count int64) { - if count <= 0 { - return - } - - stats.shard(index).invalidatedKeyCount.Add(count) -} - -func (stats *statsCollector) recordAllInvalidation(count int64) { - if count <= 0 { - return - } - - stats.invalidatedAllCount.Add(count) -} diff --git a/cache/storage.go b/cache/storage.go deleted file mode 100644 index c53c779..0000000 --- a/cache/storage.go +++ /dev/null @@ -1,348 +0,0 @@ -package cache - -import ( - "hash/maphash" - "sync" - "time" -) - -const defaultStorageSegmentCount = 32 - -type entry[V any] struct { - key string - cached cachedValue[V] - - expiresAt time.Time - - previous *entry[V] - next *entry[V] -} - -// storage routes keys across independent storage segments. -// -// Each segment owns its map, LRU list, and mutex. Segment capacities sum -// exactly to MaxEntries. -type storage[V any] struct { - seed maphash.Seed - - segments []storageSegment[V] - mask uint64 - maxEntries int -} - -type storageSegment[V any] struct { - mu sync.Mutex - - entries map[string]*entry[V] - - head *entry[V] - tail *entry[V] - - maxEntries int -} - -func newStorage[V any](maxEntries, segmentCount int) *storage[V] { - if maxEntries == 0 { - maxEntries = defaultMaxEntries - } - - if segmentCount == 0 { - segmentCount = defaultStorageSegmentCount - } - - return newStorageWithSegments[V](maxEntries, segmentCount) -} - -func newStorageWithSegments[V any](maxEntries, segmentCount int) *storage[V] { - segments := make([]storageSegment[V], segmentCount) - - baseCapacity := maxEntries / segmentCount - extraCapacity := maxEntries % segmentCount - - for index := range segments { - capacity := baseCapacity - - if index < extraCapacity { - capacity++ - } - - segments[index] = storageSegment[V]{ - entries: make(map[string]*entry[V], capacity), - maxEntries: capacity, - } - } - - mask := uint64(0) - - if segmentCount&(segmentCount-1) == 0 { - mask = uint64(segmentCount - 1) - } - - return &storage[V]{ - seed: maphash.MakeSeed(), - segments: segments, - mask: mask, - maxEntries: maxEntries, - } -} - -func (storage *storage[V]) lookupAt( - index int, - key string, - now time.Time, - stats *statsShard, -) (cachedValue[V], bool) { - return storage.segments[index].lookup(key, now, stats) -} - -func (storage *storage[V]) getAt( - index int, - key string, - now time.Time, - stats *statsShard, -) (cachedValue[V], bool) { - return storage.segments[index].get(key, now, stats) -} - -func (storage *storage[V]) setAt( - index int, - key string, - value cachedValue[V], - expiresAt time.Time, - stats *statsShard, -) { - storage.segments[index].set(key, value, expiresAt, stats) -} - -func (storage *storage[V]) deleteAt( - index int, - key string, -) bool { - return storage.segments[index].delete(key) -} - -func (storage *storage[V]) deleteAll() int64 { - var deleted int64 - - for index := range storage.segments { - deleted += storage.segments[index].deleteAll() - } - - return deleted -} - -func (storage *storage[V]) segmentIndex(key string) int { - if len(storage.segments) == 1 { - return 0 - } - - hash := maphash.String(storage.seed, key) - - if storage.mask != 0 { - return int(hash & storage.mask) - } - - return int( - hash % uint64(len(storage.segments)), - ) -} - -func (segment *storageSegment[V]) lookup( - key string, - now time.Time, - stats *statsShard, -) (cachedValue[V], bool) { - segment.mu.Lock() - defer segment.mu.Unlock() - - cached, ok := segment.getLocked(key, now, stats) - if !ok { - if stats != nil { - stats.missCount++ - } - - return cached, false - } - - if stats != nil { - if cached.found { - stats.hitCount++ - } else { - stats.negativeHitCount++ - } - } - - return cached, true -} - -func (segment *storageSegment[V]) get( - key string, - now time.Time, - stats *statsShard, -) (cachedValue[V], bool) { - segment.mu.Lock() - defer segment.mu.Unlock() - - return segment.getLocked(key, now, stats) -} - -func (segment *storageSegment[V]) getLocked( - key string, - now time.Time, - stats *statsShard, -) (cachedValue[V], bool) { - item, ok := segment.entries[key] - if !ok { - var zero cachedValue[V] - - return zero, false - } - - // TTL controls logical validity. Expired entries are removed lazily when - // accessed; there is no background or opportunistic expiration sweep. - if !now.Before(item.expiresAt) { - segment.removeLocked(item) - - if stats != nil { - stats.expirationCount++ - } - - var zero cachedValue[V] - - return zero, false - } - - // A hit affects LRU recency but never extends the entry TTL. - segment.moveToFrontLocked(item) - - return item.cached, true -} - -func (segment *storageSegment[V]) set( - key string, - value cachedValue[V], - expiresAt time.Time, - stats *statsShard, -) { - // A zero-capacity segment is possible when the caller explicitly chooses - // more segments than MaxEntries. Such a segment simply stores nothing. - if segment.maxEntries == 0 { - return - } - - segment.mu.Lock() - defer segment.mu.Unlock() - - if item, ok := segment.entries[key]; ok { - item.cached = value - item.expiresAt = expiresAt - - segment.moveToFrontLocked(item) - - return - } - - // Once the segment reaches capacity, reuse its LRU victim instead of - // allocating another entry. - if len(segment.entries) >= segment.maxEntries { - if stats != nil { - stats.evictionCount++ - } - - item := segment.tail - - delete(segment.entries, item.key) - - item.key = key - item.cached = value - item.expiresAt = expiresAt - - segment.entries[key] = item - - segment.moveToFrontLocked(item) - - return - } - - item := &entry[V]{ - key: key, - cached: value, - expiresAt: expiresAt, - } - - segment.entries[key] = item - segment.pushFrontLocked(item) -} - -func (segment *storageSegment[V]) delete(key string) bool { - segment.mu.Lock() - defer segment.mu.Unlock() - - item, ok := segment.entries[key] - if !ok { - return false - } - - segment.removeLocked(item) - - return true -} - -func (segment *storageSegment[V]) deleteAll() int64 { - segment.mu.Lock() - defer segment.mu.Unlock() - - deleted := int64(len(segment.entries)) - - clear(segment.entries) - - segment.head = nil - segment.tail = nil - - return deleted -} - -func (segment *storageSegment[V]) removeLocked(item *entry[V]) { - delete(segment.entries, item.key) - - segment.unlinkLocked(item) -} - -func (segment *storageSegment[V]) pushFrontLocked(item *entry[V]) { - item.previous = nil - item.next = segment.head - - if segment.head != nil { - segment.head.previous = item - } else { - segment.tail = item - } - - segment.head = item -} - -func (segment *storageSegment[V]) moveToFrontLocked(item *entry[V]) { - if segment.head == item { - return - } - - segment.unlinkLocked(item) - segment.pushFrontLocked(item) -} - -func (segment *storageSegment[V]) unlinkLocked(item *entry[V]) { - if item.previous != nil { - item.previous.next = item.next - } else { - segment.head = item.next - } - - if item.next != nil { - item.next.previous = item.previous - } else { - segment.tail = item.previous - } - - item.previous = nil - item.next = nil -} diff --git a/examples/otel/go.mod b/examples/otel/go.mod index a0aabe9..77210a8 100644 --- a/examples/otel/go.mod +++ b/examples/otel/go.mod @@ -5,6 +5,7 @@ go 1.26 require ( github.com/jackc/pgx/v5 v5.10.0 github.com/mkbeh/xpg v0.2.0 + github.com/mkbeh/xpg/extra/otelxpg v0.1.0 github.com/prometheus/client_golang v1.24.1 go.opentelemetry.io/otel v1.44.0 go.opentelemetry.io/otel/exporters/prometheus v0.67.0 diff --git a/examples/otel/main.go b/examples/otel/main.go index 13f0f6f..7660f94 100644 --- a/examples/otel/main.go +++ b/examples/otel/main.go @@ -13,7 +13,7 @@ import ( "github.com/jackc/pgx/v5/pgxpool" "github.com/mkbeh/xpg" - xpgotel "github.com/mkbeh/xpg/metrics/otel" + "github.com/mkbeh/xpg/extra/otelxpg" ) const ( @@ -59,7 +59,7 @@ func run(ctx context.Context) (runErr error) { xpg.WithName("otel-example"), xpg.WithLabel("xpg.pool.role", "primary"), xpg.WithMetrics( - xpgotel.NewMetrics(), + otelxpg.NewMetrics(), ), ) if err != nil { diff --git a/metrics/otel/go.mod b/extra/otelxpg/go.mod similarity index 75% rename from metrics/otel/go.mod rename to extra/otelxpg/go.mod index 966131f..702e6e3 100644 --- a/metrics/otel/go.mod +++ b/extra/otelxpg/go.mod @@ -1,4 +1,4 @@ -module github.com/mkbeh/xpg/metrics/otel +module github.com/mkbeh/xpg/extra/otelxpg go 1.26 diff --git a/metrics/otel/metrics.go b/extra/otelxpg/metrics.go similarity index 69% rename from metrics/otel/metrics.go rename to extra/otelxpg/metrics.go index 0cff385..7ae85ec 100644 --- a/metrics/otel/metrics.go +++ b/extra/otelxpg/metrics.go @@ -1,4 +1,4 @@ -package xpgotel +package otelxpg import ( "fmt" @@ -8,20 +8,17 @@ import ( "go.opentelemetry.io/otel/metric" ) -const instrumentationName = "github.com/mkbeh/xpg/otel" +const instrumentationName = "github.com/mkbeh/xpg/extra/otelxpg" // Metrics exports xpg statistics through OpenTelemetry. // // Metrics is immutable after construction and may be reused for multiple -// pool and cache registrations. +// pool registrations. type Metrics struct { meterProvider metric.MeterProvider } // metricsRegistration owns one OpenTelemetry callback registration. -// -// The same implementation is used by pool and cache metrics because both -// registrations have identical lifecycle semantics. type metricsRegistration struct { registration metric.Registration closeOnce sync.Once @@ -35,7 +32,7 @@ func (m *metricsRegistration) Close() { m.closeOnce.Do(func() { if err := m.registration.Unregister(); err != nil { otel.Handle( - fmt.Errorf("xpg/otel: unregister metrics: %w", err), + fmt.Errorf("otelxpg: unregister metrics: %w", err), ) } }) diff --git a/metrics/otel/options.go b/extra/otelxpg/options.go similarity index 90% rename from metrics/otel/options.go rename to extra/otelxpg/options.go index fb8f27e..b35f0f6 100644 --- a/metrics/otel/options.go +++ b/extra/otelxpg/options.go @@ -1,4 +1,4 @@ -package xpgotel +package otelxpg import ( "go.opentelemetry.io/otel/metric" @@ -24,7 +24,7 @@ type metricsSettings struct { // NewMetrics creates an OpenTelemetry metrics implementation. // // By default, metrics use the global OpenTelemetry MeterProvider. The returned -// value is immutable and may be reused for multiple pools and caches. +// value is immutable and may be reused for multiple pools. func NewMetrics(options ...MetricsOption) *Metrics { settings := metricsSettings{} @@ -44,7 +44,7 @@ func NewMetrics(options ...MetricsOption) *Metrics { // WithMeterProvider configures the MeterProvider used for metrics. // // The caller owns the provider and must shut it down after all instrumented -// pools and caches have been closed. +// pools have been closed. func WithMeterProvider(provider metric.MeterProvider) MetricsOption { return metricsOptionFunc(func(settings *metricsSettings) { if provider != nil { diff --git a/metrics/otel/options_test.go b/extra/otelxpg/options_test.go similarity index 59% rename from metrics/otel/options_test.go rename to extra/otelxpg/options_test.go index 1467a3a..47b213d 100644 --- a/metrics/otel/options_test.go +++ b/extra/otelxpg/options_test.go @@ -1,13 +1,11 @@ -package xpgotel +package otelxpg import ( "context" "testing" - "time" "github.com/jackc/pgx/v5/pgxpool" "github.com/mkbeh/xpg" - xpgcache "github.com/mkbeh/xpg/cache" "go.opentelemetry.io/otel/metric/noop" ) @@ -18,35 +16,9 @@ func TestMetricsRegistration(t *testing.T) { WithMeterProvider(noop.NewMeterProvider()), ) - poolConfig, err := pgxpool.ParseConfig("") - if err != nil { - t.Fatalf("parse pool config: %v", err) - } - - pool, err := xpg.New( - context.Background(), - poolConfig, - xpg.WithName("test-pool"), - xpg.WithMetrics(metrics), - ) - if err != nil { - t.Fatalf("create pool: %v", err) - } + pool := newTestPool(t, metrics) pool.Close() pool.Close() - - cache, err := xpgcache.New[int]( - xpgcache.Config{ - Name: "test-cache", - TTL: time.Minute, - Metrics: metrics, - }, - ) - if err != nil { - t.Fatalf("create cache: %v", err) - } - cache.Close() - cache.Close() } func TestWithMeterProviderNilUsesGlobalProvider(t *testing.T) { @@ -56,16 +28,27 @@ func TestWithMeterProviderNilUsesGlobalProvider(t *testing.T) { WithMeterProvider(nil), ) - cache, err := xpgcache.New[int]( - xpgcache.Config{ - Name: "test-cache", - TTL: time.Minute, - Metrics: metrics, - }, + pool := newTestPool(t, metrics) + pool.Close() +} + +func newTestPool(t *testing.T, metrics xpg.Metrics) *xpg.Pool { + t.Helper() + + poolConfig, err := pgxpool.ParseConfig("") + if err != nil { + t.Fatalf("parse pool config: %v", err) + } + + pool, err := xpg.New( + context.Background(), + poolConfig, + xpg.WithName("test-pool"), + xpg.WithMetrics(metrics), ) if err != nil { - t.Fatalf("create cache: %v", err) + t.Fatalf("create pool: %v", err) } - cache.Close() + return pool } diff --git a/metrics/otel/pool.go b/extra/otelxpg/pool.go similarity index 95% rename from metrics/otel/pool.go rename to extra/otelxpg/pool.go index 9b610ae..f1bc29c 100644 --- a/metrics/otel/pool.go +++ b/extra/otelxpg/pool.go @@ -1,4 +1,4 @@ -package xpgotel +package otelxpg import ( "context" @@ -70,7 +70,7 @@ var _ xpg.Metrics = (*Metrics)(nil) // Register registers metrics for one xpg Pool. func (m *Metrics) Register(pool *xpg.Pool) (xpg.MetricsRegistration, error) { if m == nil { - return nil, errors.New("xpg/otel: metrics is nil") + return nil, errors.New("otelxpg: metrics is nil") } provider := m.meterProvider @@ -104,7 +104,7 @@ func registerPoolMetrics(pool *xpg.Pool, provider metric.MeterProvider) (xpg.Met instruments.observables()..., ) if err != nil { - return nil, fmt.Errorf("xpg/otel: register pool metrics callback: %w", err) + return nil, fmt.Errorf("otelxpg: register pool metrics callback: %w", err) } return &metricsRegistration{ @@ -219,7 +219,7 @@ func newPoolMetricInstruments(meter metric.Meter) (poolMetricInstruments, error) ) if err != nil { return poolMetricInstruments{}, fmt.Errorf( - "xpg/otel: create %s: %w", + "otelxpg: create %s: %w", connectionCountMetricName, err, ) @@ -234,7 +234,7 @@ func newPoolMetricInstruments(meter metric.Meter) (poolMetricInstruments, error) ) if err != nil { return poolMetricInstruments{}, fmt.Errorf( - "xpg/otel: create %s: %w", + "otelxpg: create %s: %w", connectionMaxMetricName, err, ) @@ -249,7 +249,7 @@ func newPoolMetricInstruments(meter metric.Meter) (poolMetricInstruments, error) ) if err != nil { return poolMetricInstruments{}, fmt.Errorf( - "xpg/otel: create %s: %w", + "otelxpg: create %s: %w", connectionConstructingMetricName, err, ) @@ -264,7 +264,7 @@ func newPoolMetricInstruments(meter metric.Meter) (poolMetricInstruments, error) ) if err != nil { return poolMetricInstruments{}, fmt.Errorf( - "xpg/otel: create %s: %w", + "otelxpg: create %s: %w", connectionAcquireCountMetricName, err, ) @@ -279,7 +279,7 @@ func newPoolMetricInstruments(meter metric.Meter) (poolMetricInstruments, error) ) if err != nil { return poolMetricInstruments{}, fmt.Errorf( - "xpg/otel: create %s: %w", + "otelxpg: create %s: %w", connectionAcquireTimeMetricName, err, ) @@ -294,7 +294,7 @@ func newPoolMetricInstruments(meter metric.Meter) (poolMetricInstruments, error) ) if err != nil { return poolMetricInstruments{}, fmt.Errorf( - "xpg/otel: create %s: %w", + "otelxpg: create %s: %w", connectionAcquireCanceledCountMetricName, err, ) @@ -309,7 +309,7 @@ func newPoolMetricInstruments(meter metric.Meter) (poolMetricInstruments, error) ) if err != nil { return poolMetricInstruments{}, fmt.Errorf( - "xpg/otel: create %s: %w", + "otelxpg: create %s: %w", connectionAcquireEmptyCountMetricName, err, ) @@ -324,7 +324,7 @@ func newPoolMetricInstruments(meter metric.Meter) (poolMetricInstruments, error) ) if err != nil { return poolMetricInstruments{}, fmt.Errorf( - "xpg/otel: create %s: %w", + "otelxpg: create %s: %w", connectionAcquireEmptyWaitTimeMetricName, err, ) @@ -339,7 +339,7 @@ func newPoolMetricInstruments(meter metric.Meter) (poolMetricInstruments, error) ) if err != nil { return poolMetricInstruments{}, fmt.Errorf( - "xpg/otel: create %s: %w", + "otelxpg: create %s: %w", connectionCreateCountMetricName, err, ) @@ -354,7 +354,7 @@ func newPoolMetricInstruments(meter metric.Meter) (poolMetricInstruments, error) ) if err != nil { return poolMetricInstruments{}, fmt.Errorf( - "xpg/otel: create %s: %w", + "otelxpg: create %s: %w", connectionDestroyCountMetricName, err, ) diff --git a/go.mod b/go.mod index 1beb438..dd70cea 100644 --- a/go.mod +++ b/go.mod @@ -2,14 +2,12 @@ module github.com/mkbeh/xpg go 1.26 -require ( - github.com/jackc/pgx/v5 v5.10.0 - golang.org/x/sync v0.22.0 -) +require github.com/jackc/pgx/v5 v5.10.0 require ( github.com/jackc/pgpassfile v1.0.0 // indirect github.com/jackc/pgservicefile v0.0.0-20240606120523-5a60cdf6a761 // indirect github.com/jackc/puddle/v2 v2.2.2 // indirect + golang.org/x/sync v0.22.0 // indirect golang.org/x/text v0.41.0 // indirect ) diff --git a/metrics/otel/cache.go b/metrics/otel/cache.go deleted file mode 100644 index 73ee273..0000000 --- a/metrics/otel/cache.go +++ /dev/null @@ -1,451 +0,0 @@ -package xpgotel - -import ( - "context" - "errors" - "fmt" - "slices" - - xpgcache "github.com/mkbeh/xpg/cache" - "go.opentelemetry.io/otel" - "go.opentelemetry.io/otel/attribute" - "go.opentelemetry.io/otel/metric" -) - -const ( - cacheEntryCountMetricName = "xpg.cache.entry.count" - cacheEntryMaxMetricName = "xpg.cache.entry.max" - cacheSegmentCountMetricName = "xpg.cache.segment.count" - cacheLookupCountMetricName = "xpg.cache.lookup.count" - cacheLoadCountMetricName = "xpg.cache.load.count" - cacheLoadTimeMetricName = "xpg.cache.load.time" - cacheSharedCountMetricName = "xpg.cache.singleflight.shared.count" - cacheInvalidationMetricName = "xpg.cache.invalidation.count" - cacheEvictionCountMetricName = "xpg.cache.entry.eviction.count" - cacheExpirationCountMetricName = "xpg.cache.entry.expiration.count" -) - -const ( - cacheNameAttribute = "xpg.cache.name" - cacheLookupResultAttribute = "xpg.cache.lookup.result" - cacheLoadResultAttribute = "xpg.cache.load.result" - cacheInvalidationScopeAttribute = "xpg.cache.invalidation.scope" -) - -const ( - cacheLookupResultHit = "hit" - cacheLookupResultNegativeHit = "negative_hit" - cacheLookupResultMiss = "miss" - - cacheLoadResultFound = "found" - cacheLoadResultNotFound = "not_found" - cacheLoadResultError = "error" - - cacheInvalidationScopeKeys = "keys" - cacheInvalidationScopeAll = "all" -) - -var _ xpgcache.Metrics = (*Metrics)(nil) - -type cacheMetricInstruments struct { - entryCount metric.Int64ObservableGauge - entryMax metric.Int64ObservableGauge - segmentCount metric.Int64ObservableGauge - lookupCount metric.Int64ObservableCounter - loadCount metric.Int64ObservableCounter - loadTime metric.Float64ObservableCounter - sharedCount metric.Int64ObservableCounter - invalidationCount metric.Int64ObservableCounter - evictionCount metric.Int64ObservableCounter - expirationCount metric.Int64ObservableCounter -} - -type cacheMetricAttributes struct { - base metric.ObserveOption - - hit metric.ObserveOption - negativeHit metric.ObserveOption - miss metric.ObserveOption - - loadFound metric.ObserveOption - loadNotFound metric.ObserveOption - loadError metric.ObserveOption - - invalidateKeys metric.ObserveOption - invalidateAll metric.ObserveOption -} - -// RegisterCache registers OpenTelemetry metrics for one cache. -func (m *Metrics) RegisterCache(cache xpgcache.StatsProvider) (xpgcache.MetricsRegistration, error) { - if m == nil { - return nil, errors.New("xpg/otel: metrics is nil") - } - - if cache == nil { - return nil, errors.New("xpg/otel: cache is nil") - } - - name := cache.Name() - if name == "" { - return nil, errors.New("xpg/otel: cache name is blank") - } - - provider := m.meterProvider - if provider == nil { - provider = otel.GetMeterProvider() - } - - return registerCacheMetrics(cache, name, provider) -} - -func registerCacheMetrics( - cache xpgcache.StatsProvider, - name string, - provider metric.MeterProvider, -) (xpgcache.MetricsRegistration, error) { - meter := provider.Meter(instrumentationName) - - instruments, err := newCacheMetricInstruments(meter) - if err != nil { - return nil, err - } - - attributes := newCacheMetricAttributes(name) - - registration, err := meter.RegisterCallback( - func(_ context.Context, observer metric.Observer) error { - instruments.observe(observer, cache.Stats(), attributes) - - return nil - }, - instruments.observables()..., - ) - if err != nil { - return nil, fmt.Errorf("xpg/otel: register cache metrics callback: %w", err) - } - - return &metricsRegistration{ - registration: registration, - }, nil -} - -func (instruments cacheMetricInstruments) observe( - observer metric.Observer, - stats xpgcache.Stats, - attributes cacheMetricAttributes, -) { - observer.ObserveInt64( - instruments.entryCount, - stats.EntryCount, - attributes.base, - ) - observer.ObserveInt64( - instruments.entryMax, - stats.MaxEntries, - attributes.base, - ) - observer.ObserveInt64( - instruments.segmentCount, - stats.SegmentCount, - attributes.base, - ) - observer.ObserveInt64( - instruments.lookupCount, - stats.HitCount, - attributes.hit, - ) - observer.ObserveInt64( - instruments.lookupCount, - stats.NegativeHitCount, - attributes.negativeHit, - ) - observer.ObserveInt64( - instruments.lookupCount, - stats.MissCount, - attributes.miss, - ) - observer.ObserveInt64( - instruments.loadCount, - stats.LoadFoundCount, - attributes.loadFound, - ) - observer.ObserveInt64( - instruments.loadCount, - stats.LoadNotFoundCount, - attributes.loadNotFound, - ) - observer.ObserveInt64( - instruments.loadCount, - stats.LoadErrorCount, - attributes.loadError, - ) - observer.ObserveFloat64( - instruments.loadTime, - stats.LoadDuration.Seconds(), - attributes.base, - ) - observer.ObserveInt64( - instruments.sharedCount, - stats.SharedCount, - attributes.base, - ) - observer.ObserveInt64( - instruments.invalidationCount, - stats.InvalidatedKeyCount, - attributes.invalidateKeys, - ) - observer.ObserveInt64( - instruments.invalidationCount, - stats.InvalidatedAllCount, - attributes.invalidateAll, - ) - observer.ObserveInt64( - instruments.evictionCount, - stats.EvictionCount, - attributes.base, - ) - observer.ObserveInt64( - instruments.expirationCount, - stats.ExpirationCount, - attributes.base, - ) -} - -func (instruments cacheMetricInstruments) observables() []metric.Observable { - return []metric.Observable{ - instruments.entryCount, - instruments.entryMax, - instruments.segmentCount, - instruments.lookupCount, - instruments.loadCount, - instruments.loadTime, - instruments.sharedCount, - instruments.invalidationCount, - instruments.evictionCount, - instruments.expirationCount, - } -} - -func newCacheMetricInstruments(meter metric.Meter) (cacheMetricInstruments, error) { - var instruments cacheMetricInstruments - - var err error - - instruments.entryCount, err = meter.Int64ObservableGauge( - cacheEntryCountMetricName, - metric.WithDescription( - "The number of entries currently resident in cache storage.", - ), - metric.WithUnit("{entry}"), - ) - if err != nil { - return cacheMetricInstruments{}, fmt.Errorf( - "xpg/otel: create %s: %w", - cacheEntryCountMetricName, - err, - ) - } - - instruments.entryMax, err = meter.Int64ObservableGauge( - cacheEntryMaxMetricName, - metric.WithDescription( - "The configured total cache entry budget.", - ), - metric.WithUnit("{entry}"), - ) - if err != nil { - return cacheMetricInstruments{}, fmt.Errorf( - "xpg/otel: create %s: %w", - cacheEntryMaxMetricName, - err, - ) - } - - instruments.segmentCount, err = meter.Int64ObservableGauge( - cacheSegmentCountMetricName, - metric.WithDescription( - "The number of independent cache storage segments.", - ), - metric.WithUnit("{segment}"), - ) - if err != nil { - return cacheMetricInstruments{}, fmt.Errorf( - "xpg/otel: create %s: %w", - cacheSegmentCountMetricName, - err, - ) - } - - instruments.lookupCount, err = meter.Int64ObservableCounter( - cacheLookupCountMetricName, - metric.WithDescription( - "The cumulative number of cache lookups by result.", - ), - metric.WithUnit("{lookup}"), - ) - if err != nil { - return cacheMetricInstruments{}, fmt.Errorf( - "xpg/otel: create %s: %w", - cacheLookupCountMetricName, - err, - ) - } - - instruments.loadCount, err = meter.Int64ObservableCounter( - cacheLoadCountMetricName, - metric.WithDescription( - "The cumulative number of actual cache loader invocations by result.", - ), - metric.WithUnit("{load}"), - ) - if err != nil { - return cacheMetricInstruments{}, fmt.Errorf( - "xpg/otel: create %s: %w", - cacheLoadCountMetricName, - err, - ) - } - - instruments.loadTime, err = meter.Float64ObservableCounter( - cacheLoadTimeMetricName, - metric.WithDescription( - "The cumulative time spent in actual cache loader invocations.", - ), - metric.WithUnit("s"), - ) - if err != nil { - return cacheMetricInstruments{}, fmt.Errorf( - "xpg/otel: create %s: %w", - cacheLoadTimeMetricName, - err, - ) - } - - instruments.sharedCount, err = meter.Int64ObservableCounter( - cacheSharedCountMetricName, - metric.WithDescription( - "The cumulative number of cache callers that received a shared singleflight result.", - ), - metric.WithUnit("{request}"), - ) - if err != nil { - return cacheMetricInstruments{}, fmt.Errorf( - "xpg/otel: create %s: %w", - cacheSharedCountMetricName, - err, - ) - } - - instruments.invalidationCount, err = meter.Int64ObservableCounter( - cacheInvalidationMetricName, - metric.WithDescription( - "The cumulative number of resident cache entries removed by explicit invalidation, by scope.", - ), - metric.WithUnit("{entry}"), - ) - if err != nil { - return cacheMetricInstruments{}, fmt.Errorf( - "xpg/otel: create %s: %w", - cacheInvalidationMetricName, - err, - ) - } - - instruments.evictionCount, err = meter.Int64ObservableCounter( - cacheEvictionCountMetricName, - metric.WithDescription( - "The cumulative number of cache entries evicted because a storage segment reached capacity.", - ), - metric.WithUnit("{entry}"), - ) - if err != nil { - return cacheMetricInstruments{}, fmt.Errorf( - "xpg/otel: create %s: %w", - cacheEvictionCountMetricName, - err, - ) - } - - instruments.expirationCount, err = meter.Int64ObservableCounter( - cacheExpirationCountMetricName, - metric.WithDescription( - "The cumulative number of expired cache entries removed during lookup.", - ), - metric.WithUnit("{entry}"), - ) - if err != nil { - return cacheMetricInstruments{}, fmt.Errorf( - "xpg/otel: create %s: %w", - cacheExpirationCountMetricName, - err, - ) - } - - return instruments, nil -} - -func newCacheMetricAttributes(name string) cacheMetricAttributes { - base := []attribute.KeyValue{ - attribute.String(cacheNameAttribute, name), - } - - option := func(extra ...attribute.KeyValue) metric.ObserveOption { - return metric.WithAttributeSet( - attribute.NewSet( - slices.Concat(base, extra)..., - ), - ) - } - - return cacheMetricAttributes{ - base: option(), - hit: option( - attribute.String( - cacheLookupResultAttribute, - cacheLookupResultHit, - ), - ), - negativeHit: option( - attribute.String( - cacheLookupResultAttribute, - cacheLookupResultNegativeHit, - ), - ), - miss: option( - attribute.String( - cacheLookupResultAttribute, - cacheLookupResultMiss, - ), - ), - loadFound: option( - attribute.String( - cacheLoadResultAttribute, - cacheLoadResultFound, - ), - ), - loadNotFound: option( - attribute.String( - cacheLoadResultAttribute, - cacheLoadResultNotFound, - ), - ), - loadError: option( - attribute.String( - cacheLoadResultAttribute, - cacheLoadResultError, - ), - ), - invalidateKeys: option( - attribute.String( - cacheInvalidationScopeAttribute, - cacheInvalidationScopeKeys, - ), - ), - invalidateAll: option( - attribute.String( - cacheInvalidationScopeAttribute, - cacheInvalidationScopeAll, - ), - ), - } -} From ad8c4b223899f58317c15a94e671feeec217d925 Mon Sep 17 00:00:00 2001 From: mkbeh Date: Sun, 23 Aug 2026 14:10:04 +0300 Subject: [PATCH 17/41] build: upgrade to go1.27 --- README.md | 2 +- Taskfile.yml | 2 +- examples/advisory/go.mod | 2 +- examples/basic/go.mod | 2 +- examples/cluster/go.mod | 2 +- examples/otel/go.mod | 2 +- examples/shard/go.mod | 2 +- examples/transactions/go.mod | 2 +- extra/otelxpg/go.mod | 2 +- go.mod | 2 +- 10 files changed, 10 insertions(+), 10 deletions(-) diff --git a/README.md b/README.md index 9dd9622..6ed45c4 100644 --- a/README.md +++ b/README.md @@ -4,7 +4,7 @@ **Lightweight PostgreSQL wrapper for Go, built on top of [pgx](https://github.com/jackc/pgx).** -![Go Version](https://img.shields.io/badge/go-1.26%2B-blue) +![Go Version](https://img.shields.io/badge/go-1.27%2B-blue) [![License: MIT](https://img.shields.io/badge/license-MIT-green.svg)](LICENSE) diff --git a/Taskfile.yml b/Taskfile.yml index 12b46ca..2f18413 100755 --- a/Taskfile.yml +++ b/Taskfile.yml @@ -1,7 +1,7 @@ version: "3" vars: - LINTER_VER: "v2.12.2" + LINTER_VER: "v2.13.1" GOMODCACHE: sh: go env GOMODCACHE diff --git a/examples/advisory/go.mod b/examples/advisory/go.mod index 1907213..20fe69d 100644 --- a/examples/advisory/go.mod +++ b/examples/advisory/go.mod @@ -1,6 +1,6 @@ module advisory -go 1.26 +go 1.27 require ( github.com/jackc/pgx/v5 v5.10.0 diff --git a/examples/basic/go.mod b/examples/basic/go.mod index c1f1325..4a1fff5 100644 --- a/examples/basic/go.mod +++ b/examples/basic/go.mod @@ -1,6 +1,6 @@ module basic -go 1.26 +go 1.27 require ( github.com/jackc/pgx/v5 v5.10.0 diff --git a/examples/cluster/go.mod b/examples/cluster/go.mod index da620f5..ba9090d 100644 --- a/examples/cluster/go.mod +++ b/examples/cluster/go.mod @@ -1,6 +1,6 @@ module cluster -go 1.26 +go 1.27 require ( github.com/jackc/pgx/v5 v5.10.0 diff --git a/examples/otel/go.mod b/examples/otel/go.mod index 77210a8..bd781d7 100644 --- a/examples/otel/go.mod +++ b/examples/otel/go.mod @@ -1,6 +1,6 @@ module observability -go 1.26 +go 1.27 require ( github.com/jackc/pgx/v5 v5.10.0 diff --git a/examples/shard/go.mod b/examples/shard/go.mod index 2215122..e9cc1be 100644 --- a/examples/shard/go.mod +++ b/examples/shard/go.mod @@ -1,6 +1,6 @@ module github.com/mkbeh/xpg/examples/shard -go 1.26 +go 1.27 require ( github.com/jackc/pgx/v5 v5.10.0 diff --git a/examples/transactions/go.mod b/examples/transactions/go.mod index 81a5c4c..0c3b7cc 100644 --- a/examples/transactions/go.mod +++ b/examples/transactions/go.mod @@ -1,6 +1,6 @@ module transactions -go 1.26 +go 1.27 require ( github.com/mkbeh/xpg v0.2.0 diff --git a/extra/otelxpg/go.mod b/extra/otelxpg/go.mod index 702e6e3..859994a 100644 --- a/extra/otelxpg/go.mod +++ b/extra/otelxpg/go.mod @@ -1,6 +1,6 @@ module github.com/mkbeh/xpg/extra/otelxpg -go 1.26 +go 1.27 require ( github.com/mkbeh/xpg v0.2.0 diff --git a/go.mod b/go.mod index dd70cea..e7fcb52 100644 --- a/go.mod +++ b/go.mod @@ -1,6 +1,6 @@ module github.com/mkbeh/xpg -go 1.26 +go 1.27 require github.com/jackc/pgx/v5 v5.10.0 From 1091c06dbd4368ccf718b1a85478183eda2a6501 Mon Sep 17 00:00:00 2001 From: mkbeh Date: Sun, 23 Aug 2026 14:10:22 +0300 Subject: [PATCH 18/41] chore: update golangci-lint configuration --- .golangci.yml | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/.golangci.yml b/.golangci.yml index dbfafd5..9c32ca8 100755 --- a/.golangci.yml +++ b/.golangci.yml @@ -1,7 +1,7 @@ version: "2" run: - go: "1.26" + go: "1.27" timeout: 5m allow-parallel-runners: true @@ -168,7 +168,8 @@ formatters: - goimports settings: gofumpt: - extra-rules: true + extra: + group-params: true exclusions: generated: lax paths: From b3f7c71cf46f44acd4d97b872c87d5a8e7ce0d57 Mon Sep 17 00:00:00 2001 From: mkbeh Date: Mon, 24 Aug 2026 17:04:29 +0300 Subject: [PATCH 19/41] feat: add logging and tracing support --- .github/workflows/lint.yml | 1 + examples/README.md | 18 +- examples/observability/README.md | 214 ++++++++++++++++++++ examples/{otel => observability}/go.mod | 21 +- examples/{otel => observability}/main.go | 102 ++++++++-- examples/{otel => observability}/metrics.go | 25 +-- examples/observability/otel.go | 68 +++++++ examples/otel/README.md | 140 ------------- extra/slogxpg/go.mod | 13 ++ extra/slogxpg/go.sum | 26 +++ extra/slogxpg/logger.go | 75 +++++++ extra/slogxpg/logger_test.go | 169 ++++++++++++++++ options.go | 63 +++++- pool.go | 4 + 14 files changed, 740 insertions(+), 199 deletions(-) create mode 100644 examples/observability/README.md rename examples/{otel => observability}/go.mod (65%) rename examples/{otel => observability}/main.go (59%) rename examples/{otel => observability}/metrics.go (68%) create mode 100644 examples/observability/otel.go delete mode 100644 examples/otel/README.md create mode 100644 extra/slogxpg/go.mod create mode 100644 extra/slogxpg/go.sum create mode 100644 extra/slogxpg/logger.go create mode 100644 extra/slogxpg/logger_test.go diff --git a/.github/workflows/lint.yml b/.github/workflows/lint.yml index e02f372..4f7bdc8 100644 --- a/.github/workflows/lint.yml +++ b/.github/workflows/lint.yml @@ -31,6 +31,7 @@ jobs: module: - . - extra/otelxpg + - extra/slogxpg - examples/advisory - examples/basic - examples/cluster diff --git a/examples/README.md b/examples/README.md index 32c21d1..29cd2e8 100644 --- a/examples/README.md +++ b/examples/README.md @@ -2,14 +2,14 @@ This directory contains runnable examples demonstrating the main features and usage patterns of `xpg`. -| Example | Demonstrates | -|:-------------------------------|:------------------------------------------------------------------------------------------| -| [`basic`](basic) | Pool lifecycle and common query methods | -| [`transactions`](transactions) | Committing an outer transaction after an optional operation is rolled back to a savepoint | -| [`advisory`](advisory) | Coordinating concurrent transactions with PostgreSQL advisory locks | -| [`otel`](otel) | Exporting pool metrics through OpenTelemetry and Prometheus | -| [`cluster`](cluster) | Routing reads and transactions across primary and replica pools | -| [`sharding`](shard) | Typed routing across immutable standalone and cluster shard targets | +| Example | Demonstrates | +|:---------------------------------|:------------------------------------------------------------------------------------------| +| [`basic`](basic) | Pool lifecycle and common query methods | +| [`transactions`](transactions) | Committing an outer transaction after an optional operation is rolled back to a savepoint | +| [`advisory`](advisory) | Coordinating concurrent transactions with PostgreSQL advisory locks | +| [`observability`](observability) | Logging with `slog`, OpenTelemetry tracing, and Prometheus pool metrics | +| [`cluster`](cluster) | Routing reads and transactions across primary and replica pools | +| [`sharding`](shard) | Typed routing across immutable standalone and cluster shard targets | ## Running the examples @@ -30,4 +30,4 @@ go run . > [!NOTE] > Some examples may require different services or configuration. Refer to the README in the corresponding example -> directory for the exact startup command, connection settings, and expected output. +> directory for the exact startup command, connection settings, and expected output. \ No newline at end of file diff --git a/examples/observability/README.md b/examples/observability/README.md new file mode 100644 index 0000000..f5f8509 --- /dev/null +++ b/examples/observability/README.md @@ -0,0 +1,214 @@ +# Observability Example + +This example shows how to combine logging, distributed tracing, and connection pool metrics with `xpg`. + +**This example demonstrates:** + +* Adapting the standard library `log/slog` logger to `pgx/tracelog` +* Attaching an OpenTelemetry PostgreSQL tracer through `xpg.WithTracer` +* Combining the logger and tracer automatically through the `xpg` tracing pipeline +* Exporting `xpg` pool metrics through OpenTelemetry and Prometheus +* Propagating an application span through concurrent PostgreSQL operations + +The signals remain independent at the application boundary: + +```text +slog otelpgx otelxpg + | | | + +---- pgx tracing -----+ OTel metrics + | | + xpg Prometheus + | + /metrics +``` + +`otelpgx.RecordStats` is intentionally not used here. PostgreSQL tracing is handled by `otelpgx`, while pool metrics are +owned by `xpg` and exported through `extra/otelxpg`. + +## Configuration + +The example uses the following connection string by default: + +```text +postgres://postgres:postgres@localhost:5432/postgres?sslmode=disable +``` + +Set `XPG_DATABASE_URL` to use another PostgreSQL instance: + +```shell +export XPG_DATABASE_URL='postgres://user:password@localhost:5432/database?sslmode=disable' +``` + +The HTTP server listens on `localhost:9464` by default. Override it with: + +```shell +export HTTP_ADDR='localhost:9464' +``` + +## Local setup + +Start PostgreSQL and Adminer from the repository root: + +```shell +docker compose -f examples/docker-compose.yml --profile tools up -d +``` + +Or from this example directory: + +```shell +docker compose -f ../docker-compose.yml --profile tools up -d +``` + +Services are available at: + +```text +PostgreSQL: localhost:5432 +Adminer: http://localhost:8080 +``` + +Sign in to Adminer with: + +```text +System: PostgreSQL +Server: postgres +Username: postgres +Password: postgres +Database: postgres +``` + +## Run + +From this directory: + +```shell +go run . +``` + +Or from the repository root: + +```shell +go run ./examples/observability +``` + +The example starts an HTTP server at: + +```text +http://localhost:9464 +``` + +## Logging + +The example adapts `log/slog` to `tracelog.Logger` and passes it to the pool: + +```go +xpg.WithLogger( + newPGXLogger(logger), + tracelog.LogLevelInfo, +) +``` + +The adapter also copies the active OpenTelemetry `trace_id` and `span_id` into pgx log records when a span is present, +which makes logs and database spans directly correlatable. + +`pgx/tracelog` can include SQL text and query arguments in log records. Production applications should choose logging +levels and redaction policies appropriate for the data they process. + +## Tracing + +`otelpgx` is attached as a normal pgx tracer: + +```go +xpg.WithTracer( + otelpgx.NewTracer( + otelpgx.WithTracerProvider(tracing.TracerProvider()), + ), +) +``` + +The example writes completed spans to stdout as formatted JSON. The stdout exporter and synchronous span processor are +used only to make the example self-contained and immediately observable. Production applications should normally export +traces through OTLP and use a batch span processor. + +Generate a traced workload: + +```shell +curl -X POST 'http://localhost:9464/load' +``` + +The handler creates one application span and passes its context to six concurrent PostgreSQL operations. Database spans +created by `otelpgx` therefore appear as children of that application span. + +## Metrics + +Pool metrics use an explicit OpenTelemetry `MeterProvider` backed by the Prometheus exporter: + +```go +xpg.WithMetrics( + otelxpg.NewMetrics( + otelxpg.WithMeterProvider(metrics.MeterProvider()), + ), +) +``` + +Open the Prometheus endpoint: + +```shell +curl 'http://localhost:9464/metrics' +``` + +Show only database connection pool and `xpg` metrics: + +```shell +curl -s 'http://localhost:9464/metrics' \ + | grep -E '^(db_client_connection|xpg_pool_connection_)' +``` + +The example exports these metric families: + +```text +db_client_connection_count +db_client_connection_max +xpg_pool_connection_constructing +xpg_pool_connection_acquire_count_total +xpg_pool_connection_acquire_time_seconds_total +xpg_pool_connection_acquire_canceled_count_total +xpg_pool_connection_acquire_empty_count_total +xpg_pool_connection_acquire_empty_wait_time_seconds_total +xpg_pool_connection_create_count_total +xpg_pool_connection_destroy_count_total +``` + +The Prometheus exporter converts OpenTelemetry dotted instrument names to Prometheus-compatible names and adds unit and +counter suffixes where required. + +## Generate pool contention + +The pool is intentionally limited to two connections. Run: + +```shell +curl -X POST 'http://localhost:9464/load' +``` + +While it is running, inspect metrics from another terminal: + +```shell +curl -s 'http://localhost:9464/metrics' \ + | grep -E 'db_client_connection_count|xpg_pool_connection_acquire_' +``` + +Six concurrent one-second queries make connection usage and acquire wait metrics visible while also producing related +logs and spans. + +## Stop services + +From the repository root: + +```shell +docker compose -f examples/docker-compose.yml --profile tools down --remove-orphans -v +``` + +Or from this example directory: + +```shell +docker compose -f ../docker-compose.yml --profile tools down --remove-orphans -v +``` diff --git a/examples/otel/go.mod b/examples/observability/go.mod similarity index 65% rename from examples/otel/go.mod rename to examples/observability/go.mod index bd781d7..9a67ba3 100644 --- a/examples/otel/go.mod +++ b/examples/observability/go.mod @@ -3,20 +3,24 @@ module observability go 1.27 require ( + github.com/exaring/otelpgx v0.11.1 github.com/jackc/pgx/v5 v5.10.0 github.com/mkbeh/xpg v0.2.0 github.com/mkbeh/xpg/extra/otelxpg v0.1.0 + github.com/mkbeh/xpg/extra/slogxpg v0.1.0 github.com/prometheus/client_golang v1.24.1 - go.opentelemetry.io/otel v1.44.0 + go.opentelemetry.io/otel v1.45.0 go.opentelemetry.io/otel/exporters/prometheus v0.67.0 + go.opentelemetry.io/otel/exporters/stdout/stdouttrace v1.45.0 go.opentelemetry.io/otel/sdk v1.45.0 go.opentelemetry.io/otel/sdk/metric v1.45.0 + go.opentelemetry.io/otel/trace v1.45.0 ) require ( github.com/beorn7/perks v1.0.1 // indirect github.com/cespare/xxhash/v2 v2.3.0 // indirect - github.com/go-logr/logr v1.4.3 // indirect + github.com/go-logr/logr v1.4.4 // indirect github.com/go-logr/stdr v1.2.2 // indirect github.com/google/uuid v1.6.0 // indirect github.com/jackc/pgpassfile v1.0.0 // indirect @@ -24,15 +28,14 @@ require ( github.com/jackc/puddle/v2 v2.2.2 // indirect github.com/munnerz/goautoneg v0.0.0-20191010083416-a7dc8b61c822 // indirect github.com/prometheus/client_model v0.6.2 // indirect - github.com/prometheus/common v0.69.0 // indirect + github.com/prometheus/common v0.70.1 // indirect github.com/prometheus/otlptranslator v1.0.0 // indirect - github.com/prometheus/procfs v0.20.1 // indirect + github.com/prometheus/procfs v0.21.1 // indirect go.opentelemetry.io/auto/sdk v1.2.1 // indirect - go.opentelemetry.io/otel/metric v1.44.0 // indirect - go.opentelemetry.io/otel/metric/x v0.66.0 // indirect - go.opentelemetry.io/otel/trace v1.44.0 // indirect + go.opentelemetry.io/otel/metric v1.45.0 // indirect + go.opentelemetry.io/otel/metric/x v0.67.0 // indirect golang.org/x/sync v0.22.0 // indirect - golang.org/x/sys v0.46.0 // indirect - golang.org/x/text v0.40.0 // indirect + golang.org/x/sys v0.47.0 // indirect + golang.org/x/text v0.41.0 // indirect google.golang.org/protobuf v1.36.11 // indirect ) diff --git a/examples/otel/main.go b/examples/observability/main.go similarity index 59% rename from examples/otel/main.go rename to examples/observability/main.go index 7660f94..fab9f55 100644 --- a/examples/otel/main.go +++ b/examples/observability/main.go @@ -4,16 +4,21 @@ import ( "context" "errors" "fmt" - "log" + "log/slog" "net/http" "os" "os/signal" "syscall" "time" + "github.com/exaring/otelpgx" "github.com/jackc/pgx/v5/pgxpool" + "github.com/jackc/pgx/v5/tracelog" "github.com/mkbeh/xpg" "github.com/mkbeh/xpg/extra/otelxpg" + "github.com/mkbeh/xpg/extra/slogxpg" + "go.opentelemetry.io/otel/codes" + "go.opentelemetry.io/otel/trace" ) const ( @@ -22,16 +27,35 @@ const ( ) func main() { - ctx, stop := signal.NotifyContext(context.Background(), os.Interrupt, syscall.SIGTERM) + logger := slog.New( + slog.NewTextHandler(os.Stdout, &slog.HandlerOptions{ + Level: slog.LevelDebug, + }), + ) + + ctx, stop := signal.NotifyContext( + context.Background(), + os.Interrupt, + syscall.SIGTERM, + ) defer stop() - if err := run(ctx); err != nil { - log.Fatal(err) + if err := run(ctx, logger); err != nil { + logger.Error( + "observability example failed", + slog.Any("error", err), + ) + os.Exit(1) } } -func run(ctx context.Context) (runErr error) { - metrics, err := newMetricsRuntime(ctx) +func run(ctx context.Context, logger *slog.Logger) (runErr error) { + res, err := newOTelResource(ctx) + if err != nil { + return err + } + + metrics, err := newMetricsRuntime(res) if err != nil { return fmt.Errorf("initialize metrics: %w", err) } @@ -45,6 +69,20 @@ func run(ctx context.Context) (runErr error) { ) }() + tracing, err := newTracingRuntime(res) + if err != nil { + return fmt.Errorf("initialize tracing: %w", err) + } + defer func() { + shutdownCtx, cancel := context.WithTimeout(context.Background(), 5*time.Second) + defer cancel() + + runErr = errors.Join( + runErr, + tracing.Shutdown(shutdownCtx), + ) + }() + config, err := pgxpool.ParseConfig(databaseURL()) if err != nil { return fmt.Errorf("parse PostgreSQL config: %w", err) @@ -53,13 +91,30 @@ func run(ctx context.Context) (runErr error) { // A small pool makes contention visible when POST /load runs. config.MaxConns = 2 + pgxLogger := logger.With(slog.String("component", "pgx")) + pool, err := xpg.New( ctx, config, - xpg.WithName("otel-example"), + xpg.WithName("observability-example"), xpg.WithLabel("xpg.pool.role", "primary"), + xpg.WithLogger( + slogxpg.New(pgxLogger), + tracelog.LogLevelInfo, + ), + xpg.WithTracer( + otelpgx.NewTracer( + otelpgx.WithTracerProvider( + tracing.TracerProvider(), + ), + ), + ), xpg.WithMetrics( - otelxpg.NewMetrics(), + otelxpg.NewMetrics( + otelxpg.WithMeterProvider( + metrics.MeterProvider(), + ), + ), ), ) if err != nil { @@ -72,9 +127,8 @@ func run(ctx context.Context) (runErr error) { } mux := http.NewServeMux() - mux.Handle("GET /metrics", metrics.Handler()) - mux.HandleFunc("POST /load", loadHandler(pool)) + mux.HandleFunc("POST /load", loadHandler(pool, tracing.Tracer())) server := &http.Server{ Addr: httpAddress(), @@ -82,7 +136,10 @@ func run(ctx context.Context) (runErr error) { ReadHeaderTimeout: 5 * time.Second, } - log.Printf("OpenTelemetry example listening on http://%s", server.Addr) + logger.Info( + "observability example listening", + slog.String("address", "http://"+server.Addr), + ) if err := serveHTTP(ctx, server); err != nil { return fmt.Errorf("serve HTTP: %w", err) @@ -105,7 +162,6 @@ func serveHTTP(ctx context.Context, server *http.Server) error { } return err - case <-ctx.Done(): } @@ -123,11 +179,17 @@ func serveHTTP(ctx context.Context, server *http.Server) error { return nil } -func loadHandler(pool *xpg.Pool) http.HandlerFunc { +func loadHandler(pool *xpg.Pool, tracer trace.Tracer) http.HandlerFunc { return func(w http.ResponseWriter, request *http.Request) { + ctx, span := tracer.Start(request.Context(), "run-load") + defer span.End() + startedAt := time.Now() - if err := runWorkload(request.Context(), pool); err != nil { + if err := runWorkload(ctx, pool); err != nil { + span.RecordError(err) + span.SetStatus(codes.Error, "workload failed") + http.Error( w, fmt.Sprintf("run workload: %v", err), @@ -137,11 +199,7 @@ func loadHandler(pool *xpg.Pool) http.HandlerFunc { return } - _, _ = fmt.Fprintf( - w, - "workload completed in %s\n", - time.Since(startedAt), - ) + _, _ = fmt.Fprintf(w, "workload completed in %s\n", time.Since(startedAt)) } } @@ -154,7 +212,11 @@ func runWorkload(ctx context.Context, pool *xpg.Pool) error { for range workerCount { go func() { <-start - _, err := pool.Exec(ctx, "SELECT pg_sleep(1)") + + _, err := pool.Exec( + ctx, + "SELECT pg_sleep(1)", + ) results <- err }() } diff --git a/examples/otel/metrics.go b/examples/observability/metrics.go similarity index 68% rename from examples/otel/metrics.go rename to examples/observability/metrics.go index 29f9cbc..3d548a0 100644 --- a/examples/otel/metrics.go +++ b/examples/observability/metrics.go @@ -7,11 +7,9 @@ import ( promclient "github.com/prometheus/client_golang/prometheus" "github.com/prometheus/client_golang/prometheus/promhttp" - "go.opentelemetry.io/otel" otelprom "go.opentelemetry.io/otel/exporters/prometheus" sdkmetric "go.opentelemetry.io/otel/sdk/metric" "go.opentelemetry.io/otel/sdk/resource" - semconv "go.opentelemetry.io/otel/semconv/v1.37.0" ) type metricsRuntime struct { @@ -19,7 +17,7 @@ type metricsRuntime struct { meterProvider *sdkmetric.MeterProvider } -func newMetricsRuntime(ctx context.Context) (*metricsRuntime, error) { +func newMetricsRuntime(res *resource.Resource) (*metricsRuntime, error) { registry := promclient.NewRegistry() exporter, err := otelprom.New( @@ -30,28 +28,11 @@ func newMetricsRuntime(ctx context.Context) (*metricsRuntime, error) { return nil, fmt.Errorf("create Prometheus exporter: %w", err) } - res, err := resource.New( - ctx, - resource.WithFromEnv(), - resource.WithTelemetrySDK(), - resource.WithAttributes( - semconv.ServiceName( - "xpg-observability-example", - ), - semconv.ServiceVersion("dev"), - ), - ) - if err != nil { - return nil, fmt.Errorf("create OpenTelemetry resource: %w", err) - } - meterProvider := sdkmetric.NewMeterProvider( sdkmetric.WithResource(res), sdkmetric.WithReader(exporter), ) - otel.SetMeterProvider(meterProvider) - return &metricsRuntime{ handler: promhttp.HandlerFor( registry, @@ -65,6 +46,10 @@ func (m *metricsRuntime) Handler() http.Handler { return m.handler } +func (m *metricsRuntime) MeterProvider() *sdkmetric.MeterProvider { + return m.meterProvider +} + func (m *metricsRuntime) Shutdown(ctx context.Context) error { return m.meterProvider.Shutdown(ctx) } diff --git a/examples/observability/otel.go b/examples/observability/otel.go new file mode 100644 index 0000000..6a76c81 --- /dev/null +++ b/examples/observability/otel.go @@ -0,0 +1,68 @@ +package main + +import ( + "context" + "fmt" + + "go.opentelemetry.io/otel/exporters/stdout/stdouttrace" + "go.opentelemetry.io/otel/sdk/resource" + sdktrace "go.opentelemetry.io/otel/sdk/trace" + semconv "go.opentelemetry.io/otel/semconv/v1.43.0" + "go.opentelemetry.io/otel/trace" +) + +const tracingInstrumentationName = "github.com/mkbeh/xpg/examples/observability" + +type tracingRuntime struct { + tracerProvider *sdktrace.TracerProvider +} + +func newOTelResource(ctx context.Context) (*resource.Resource, error) { + res, err := resource.New( + ctx, + resource.WithFromEnv(), + resource.WithTelemetrySDK(), + resource.WithAttributes( + semconv.ServiceName("xpg-observability-example"), + semconv.ServiceVersion("dev"), + ), + ) + if err != nil { + return nil, fmt.Errorf("create OpenTelemetry resource: %w", err) + } + + return res, nil +} + +func newTracingRuntime(res *resource.Resource) (*tracingRuntime, error) { + exporter, err := stdouttrace.New( + stdouttrace.WithPrettyPrint(), + ) + if err != nil { + return nil, fmt.Errorf("create stdout trace exporter: %w", err) + } + + tracerProvider := sdktrace.NewTracerProvider( + sdktrace.WithResource(res), + // A synchronous processor keeps this runnable example easy to inspect. + // Production applications should normally prefer WithBatcher with an + // OTLP exporter. + sdktrace.WithSyncer(exporter), + ) + + return &tracingRuntime{ + tracerProvider: tracerProvider, + }, nil +} + +func (t *tracingRuntime) TracerProvider() *sdktrace.TracerProvider { + return t.tracerProvider +} + +func (t *tracingRuntime) Tracer() trace.Tracer { + return t.tracerProvider.Tracer(tracingInstrumentationName) +} + +func (t *tracingRuntime) Shutdown(ctx context.Context) error { + return t.tracerProvider.Shutdown(ctx) +} diff --git a/examples/otel/README.md b/examples/otel/README.md deleted file mode 100644 index 4e94ae5..0000000 --- a/examples/otel/README.md +++ /dev/null @@ -1,140 +0,0 @@ -# OpenTelemetry Metrics Example - -This example shows how to export `xpg` connection pool metrics through the OpenTelemetry Prometheus exporter. - -**This example demonstrates:** - -* Exporting `xpg` pool metrics with OpenTelemetry and Prometheus -* Registering pool metrics through the global `MeterProvider` -* Generating pool contention through the `/load` endpoint -* Shutting down the pool and metrics provider in the correct order - -## Configuration - -The example uses the following connection string by default: - -```text -postgres://postgres:postgres@localhost:5432/postgres?sslmode=disable -``` - -Set `XPG_DATABASE_URL` to use another PostgreSQL instance: - -```shell -export XPG_DATABASE_URL='postgres://user:password@localhost:5432/database?sslmode=disable' -``` - -## Local setup - -Start PostgreSQL and Adminer from the repository root: - -```shell -docker compose -f examples/docker-compose.yml --profile tools up -d -``` - -Or from this example directory: - -```shell -docker compose -f ../docker-compose.yml --profile tools up -d -``` - -Services are available at: - -```text -PostgreSQL: localhost:5432 -Adminer: http://localhost:8080 -``` - -Sign in to Adminer with: - -```text -System: PostgreSQL -Server: postgres -Username: postgres -Password: postgres -Database: postgres -``` - -## Run - -From this directory: - -```shell -go run . -``` - -Or from the repository root: - -```shell -go run ./examples/otel -``` - -The HTTP server starts on: - -```text -http://localhost:9464 -``` - -## View metrics - -Open the Prometheus endpoint: - -```shell -curl 'http://localhost:9464/metrics' -``` - -Show only database connection pool and `xpg` metrics: - -```shell -curl -s 'http://localhost:9464/metrics' \ - | grep -E '^(db_client_connection|xpg_pool_connection_)' -``` - -The example exports these metric families: - -```text -db_client_connection_count -db_client_connection_max -xpg_pool_connection_constructing -xpg_pool_connection_acquire_count_total -xpg_pool_connection_acquire_time_seconds_total -xpg_pool_connection_acquire_canceled_count_total -xpg_pool_connection_acquire_empty_count_total -xpg_pool_connection_acquire_empty_wait_time_seconds_total -xpg_pool_connection_create_count_total -xpg_pool_connection_destroy_count_total -``` - -The Prometheus exporter converts OpenTelemetry dotted instrument names to Prometheus-compatible names and adds unit and -counter suffixes where required. - -## Generate load - -Run the debug workload: - -```shell -curl -X POST 'http://localhost:9464/load' -``` - -While it is running, inspect the pool metrics from another terminal: - -```shell -curl -s 'http://localhost:9464/metrics' \ - | grep -E 'db_client_connection_count|xpg_pool_connection_acquire_' -``` - -The workload runs six concurrent queries against a pool limited to two connections, making connection usage and wait -metrics visible. - -## Stop services - -From the repository root: - -```shell -docker compose -f examples/docker-compose.yml --profile tools down --remove-orphans -v -``` - -Or from this example directory: - -```shell -docker compose -f ../docker-compose.yml --profile tools down --remove-orphans -v -``` \ No newline at end of file diff --git a/extra/slogxpg/go.mod b/extra/slogxpg/go.mod new file mode 100644 index 0000000..3fe3a33 --- /dev/null +++ b/extra/slogxpg/go.mod @@ -0,0 +1,13 @@ +module github.com/mkbeh/xpg/extra/slogxpg + +go 1.27 + +require github.com/jackc/pgx/v5 v5.10.0 + +require ( + github.com/jackc/pgpassfile v1.0.0 // indirect + github.com/jackc/pgservicefile v0.0.0-20240606120523-5a60cdf6a761 // indirect + github.com/jackc/puddle/v2 v2.2.2 // indirect + golang.org/x/sync v0.17.0 // indirect + golang.org/x/text v0.29.0 // indirect +) diff --git a/extra/slogxpg/go.sum b/extra/slogxpg/go.sum new file mode 100644 index 0000000..c0e505b --- /dev/null +++ b/extra/slogxpg/go.sum @@ -0,0 +1,26 @@ +github.com/davecgh/go-spew v1.1.0/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38= +github.com/davecgh/go-spew v1.1.1 h1:vj9j/u1bqnvCEfJOwUhtlOARqs3+rkHYY13jYWTU97c= +github.com/davecgh/go-spew v1.1.1/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38= +github.com/jackc/pgpassfile v1.0.0 h1:/6Hmqy13Ss2zCq62VdNG8tM1wchn8zjSGOBJ6icpsIM= +github.com/jackc/pgpassfile v1.0.0/go.mod h1:CEx0iS5ambNFdcRtxPj5JhEz+xB6uRky5eyVu/W2HEg= +github.com/jackc/pgservicefile v0.0.0-20240606120523-5a60cdf6a761 h1:iCEnooe7UlwOQYpKFhBabPMi4aNAfoODPEFNiAnClxo= +github.com/jackc/pgservicefile v0.0.0-20240606120523-5a60cdf6a761/go.mod h1:5TJZWKEWniPve33vlWYSoGYefn3gLQRzjfDlhSJ9ZKM= +github.com/jackc/pgx/v5 v5.10.0 h1:VhSvgU2jSli8o3AqIEOTJr7rZwAEUVo4E4XhR94Zfr0= +github.com/jackc/pgx/v5 v5.10.0/go.mod h1:mal1tBGAFfLHvZzaYh77YS/eC6IX9OWbRV1QIIM0Jn4= +github.com/jackc/puddle/v2 v2.2.2 h1:PR8nw+E/1w0GLuRFSmiioY6UooMp6KJv0/61nB7icHo= +github.com/jackc/puddle/v2 v2.2.2/go.mod h1:vriiEXHvEE654aYKXXjOvZM39qJ0q+azkZFrfEOc3H4= +github.com/pmezard/go-difflib v1.0.0 h1:4DBwDE0NGyQoBHbLQYPwSUPoCMWR5BEzIk/f1lZbAQM= +github.com/pmezard/go-difflib v1.0.0/go.mod h1:iKH77koFhYxTK1pcRnkKkqfTogsbg7gZNVY4sRDYZ/4= +github.com/stretchr/objx v0.1.0/go.mod h1:HFkY916IF+rwdDfMAkV7OtwuqBVzrE8GR6GFx+wExME= +github.com/stretchr/testify v1.3.0/go.mod h1:M5WIy9Dh21IEIfnGCwXGc5bZfKNJtfHm1UVUgZn+9EI= +github.com/stretchr/testify v1.7.0/go.mod h1:6Fq8oRcR53rry900zMqJjRRixrwX3KX962/h/Wwjteg= +github.com/stretchr/testify v1.11.1 h1:7s2iGBzp5EwR7/aIZr8ao5+dra3wiQyKjjFuvgVKu7U= +github.com/stretchr/testify v1.11.1/go.mod h1:wZwfW3scLgRK+23gO65QZefKpKQRnfz6sD981Nm4B6U= +golang.org/x/sync v0.17.0 h1:l60nONMj9l5drqw6jlhIELNv9I0A4OFgRsG9k2oT9Ug= +golang.org/x/sync v0.17.0/go.mod h1:9KTHXmSnoGruLpwFjVSX0lNNA75CykiMECbovNTZqGI= +golang.org/x/text v0.29.0 h1:1neNs90w9YzJ9BocxfsQNHKuAT4pkghyXc4nhZ6sJvk= +golang.org/x/text v0.29.0/go.mod h1:7MhJOA9CD2qZyOKYazxdYMF85OwPdEr9jTtBpO7ydH4= +gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0= +gopkg.in/yaml.v3 v3.0.0-20200313102051-9f266ea9e77c/go.mod h1:K4uyk7z7BCEPqu6E+C64Yfv1cQ7kz7rIZviUmN+EgEM= +gopkg.in/yaml.v3 v3.0.1 h1:fxVm/GzAzEWqLHuvctI91KS9hhNmmWOoWu0XTYJS7CA= +gopkg.in/yaml.v3 v3.0.1/go.mod h1:K4uyk7z7BCEPqu6E+C64Yfv1cQ7kz7rIZviUmN+EgEM= diff --git a/extra/slogxpg/logger.go b/extra/slogxpg/logger.go new file mode 100644 index 0000000..7d12482 --- /dev/null +++ b/extra/slogxpg/logger.go @@ -0,0 +1,75 @@ +// Package slogxpg adapts log/slog loggers to pgx tracelog.Logger. +package slogxpg + +import ( + "context" + "log/slog" + + "github.com/jackc/pgx/v5/tracelog" +) + +const pgxLogLevelKey = "pgx_log_level" + +type adapter struct { + logger *slog.Logger +} + +var _ tracelog.Logger = (*adapter)(nil) + +// New adapts logger to pgx tracelog.Logger. +// +// New returns nil when logger is nil. +func New(logger *slog.Logger) tracelog.Logger { + if logger == nil { + return nil + } + + return &adapter{ + logger: logger, + } +} + +func (l *adapter) Log( + ctx context.Context, + level tracelog.LogLevel, + msg string, + data map[string]any, +) { + attrs := make([]slog.Attr, 0, len(data)+1) + + for key, value := range data { + attrs = append(attrs, slog.Any(key, value)) + } + + var logLevel slog.Level + + switch level { + case tracelog.LogLevelTrace: + logLevel = slog.LevelDebug - 1 + attrs = append( + attrs, + slog.String(pgxLogLevelKey, level.String()), + ) + case tracelog.LogLevelDebug: + logLevel = slog.LevelDebug + case tracelog.LogLevelInfo: + logLevel = slog.LevelInfo + case tracelog.LogLevelWarn: + logLevel = slog.LevelWarn + case tracelog.LogLevelError: + logLevel = slog.LevelError + default: + logLevel = slog.LevelError + attrs = append( + attrs, + slog.String(pgxLogLevelKey, level.String()), + ) + } + + l.logger.LogAttrs( + ctx, + logLevel, + msg, + attrs..., + ) +} diff --git a/extra/slogxpg/logger_test.go b/extra/slogxpg/logger_test.go new file mode 100644 index 0000000..ca8aecc --- /dev/null +++ b/extra/slogxpg/logger_test.go @@ -0,0 +1,169 @@ +package slogxpg + +import ( + "context" + "log/slog" + "testing" + + "github.com/jackc/pgx/v5/tracelog" +) + +func TestNewNil(t *testing.T) { + t.Parallel() + + if logger := New(nil); logger != nil { + t.Fatal("expected nil logger") + } +} + +func TestLoggerLevels(t *testing.T) { + t.Parallel() + + tests := []struct { + name string + pgx tracelog.LogLevel + want slog.Level + wantPGX string + }{ + { + name: "trace", + pgx: tracelog.LogLevelTrace, + want: slog.LevelDebug - 1, + wantPGX: tracelog.LogLevelTrace.String(), + }, + { + name: "debug", + pgx: tracelog.LogLevelDebug, + want: slog.LevelDebug, + }, + { + name: "info", + pgx: tracelog.LogLevelInfo, + want: slog.LevelInfo, + }, + { + name: "warn", + pgx: tracelog.LogLevelWarn, + want: slog.LevelWarn, + }, + { + name: "error", + pgx: tracelog.LogLevelError, + want: slog.LevelError, + }, + { + name: "unknown", + pgx: tracelog.LogLevel(255), + want: slog.LevelError, + wantPGX: tracelog.LogLevel(255).String(), + }, + } + + for _, test := range tests { + t.Run(test.name, func(t *testing.T) { + t.Parallel() + + handler := &captureHandler{} + logger := New(slog.New(handler)) + + logger.Log( + context.Background(), + test.pgx, + "message", + nil, + ) + + record := handler.record + + if record.Level != test.want { + t.Fatalf("level = %v, want %v", record.Level, test.want) + } + + attrs := recordAttrs(record) + + if test.wantPGX == "" { + if _, ok := attrs[pgxLogLevelKey]; ok { + t.Fatalf("unexpected %s attribute", pgxLogLevelKey) + } + + return + } + + if attrs[pgxLogLevelKey] != test.wantPGX { + t.Fatalf( + "%s = %v, want %q", + pgxLogLevelKey, + attrs[pgxLogLevelKey], + test.wantPGX, + ) + } + }) + } +} + +func TestLoggerData(t *testing.T) { + t.Parallel() + + handler := &captureHandler{} + logger := New(slog.New(handler)) + + logger.Log( + context.Background(), + tracelog.LogLevelInfo, + "query", + map[string]any{ + "sql": "select 1", + "args": 1, + }, + ) + + record := handler.record + + if record.Message != "query" { + t.Fatalf("message = %q, want %q", record.Message, "query") + } + + attrs := recordAttrs(record) + + if attrs["sql"] != "select 1" { + t.Fatalf("sql = %v, want %q", attrs["sql"], "select 1") + } + + if attrs["args"] != int64(1) { + t.Fatalf("args = %v, want %v", attrs["args"], int64(1)) + } +} + +type captureHandler struct { + record slog.Record +} + +func (h *captureHandler) Enabled(context.Context, slog.Level) bool { + return true +} + +func (h *captureHandler) Handle(_ context.Context, record slog.Record) error { + h.record = record.Clone() + + return nil +} + +func (h *captureHandler) WithAttrs([]slog.Attr) slog.Handler { + return h +} + +func (h *captureHandler) WithGroup(string) slog.Handler { + return h +} + +func recordAttrs(record slog.Record) map[string]any { + attrs := make(map[string]any, record.NumAttrs()) + + record.Attrs(func(attr slog.Attr) bool { + attrs[attr.Key] = attr.Value.Any() + + return true + }) + + return attrs +} diff --git a/options.go b/options.go index 4bea3ff..64c96b9 100644 --- a/options.go +++ b/options.go @@ -7,6 +7,10 @@ import ( "net" "strconv" "strings" + + "github.com/jackc/pgx/v5" + "github.com/jackc/pgx/v5/multitracer" + "github.com/jackc/pgx/v5/tracelog" ) // Option configures a Pool. @@ -26,9 +30,10 @@ type settings struct { name string labels map[string]string metrics Metrics + tracers []pgx.QueryTracer } -func (s settings) poolName(host string, port uint16, database string) string { +func (s *settings) poolName(host string, port uint16, database string) string { if s.name != "" { return s.name } @@ -44,6 +49,14 @@ func (s settings) poolName(host string, port uint16, database string) string { return address + "/" + database } +func (s *settings) buildTracer() pgx.QueryTracer { + if len(s.tracers) == 0 { + return nil + } + + return multitracer.New(s.tracers...) +} + func defaultSettings() *settings { return &settings{ labels: make(map[string]string), @@ -134,6 +147,54 @@ func WithMetrics(metrics Metrics) Option { }) } +// WithLogger attaches a pgx-compatible logger to the pool. +// +// Logging is implemented through pgx tracelog and participates in the same +// tracing pipeline as custom tracers. If other tracers are configured, xpg +// combines them automatically with pgx multitracer. pgx tracelog may include +// SQL text and query arguments in log records; applications are responsible for +// choosing an appropriate level and handling sensitive values. +func WithLogger(logger tracelog.Logger, level tracelog.LogLevel) Option { + return optionFunc(func(settings *settings) error { + if logger == nil { + return errors.New("pool logger is nil") + } + + settings.tracers = append(settings.tracers, &tracelog.TraceLog{ + Logger: logger, + LogLevel: level, + }) + + return nil + }) +} + +// WithTracer attaches a pgx query tracer to the pool. +// +// The option may be specified multiple times. Tracers are invoked in the order +// they are configured, after any tracer already present in +// config.ConnConfig.Tracer. When more than one tracer is present, xpg combines +// them with pgx multitracer. Additional pgx tracing capabilities implemented by +// a tracer, such as batch, copy, prepare, connect, acquire, and release tracing, +// are preserved by multitracer. +func WithTracer(tracer pgx.QueryTracer) Option { + return WithTracers(tracer) +} + +func WithTracers(tracers ...pgx.QueryTracer) Option { + return optionFunc(func(settings *settings) error { + for _, tracer := range tracers { + if tracer == nil { + return errors.New("pool tracer is nil") + } + } + + settings.tracers = append(settings.tracers, tracers...) + + return nil + }) +} + func cloneLabels(labels map[string]string) map[string]string { if len(labels) == 0 { return nil diff --git a/pool.go b/pool.go index 759dfeb..84317ef 100644 --- a/pool.go +++ b/pool.go @@ -54,6 +54,10 @@ func New(ctx context.Context, config *pgxpool.Config, options ...Option) (*Pool, poolConfig := config.Copy() connConfig := poolConfig.ConnConfig + if tracer := settings.buildTracer(); tracer != nil { + connConfig.Tracer = tracer + } + pgxPool, err := pgxpool.NewWithConfig(ctx, poolConfig) if err != nil { return nil, fmt.Errorf("xpg: create pool: %w", err) From 46e9846a8cae72b427b10fa1986fb7b661b835cb Mon Sep 17 00:00:00 2001 From: mkbeh Date: Mon, 24 Aug 2026 21:01:57 +0300 Subject: [PATCH 20/41] refactor: simplify slog adapter --- extra/slogxpg/logger.go | 43 ++++++---------- extra/slogxpg/logger_test.go | 97 +++++++++++++++++------------------- 2 files changed, 63 insertions(+), 77 deletions(-) diff --git a/extra/slogxpg/logger.go b/extra/slogxpg/logger.go index 7d12482..de5ec2c 100644 --- a/extra/slogxpg/logger.go +++ b/extra/slogxpg/logger.go @@ -1,4 +1,3 @@ -// Package slogxpg adapts log/slog loggers to pgx tracelog.Logger. package slogxpg import ( @@ -8,8 +7,6 @@ import ( "github.com/jackc/pgx/v5/tracelog" ) -const pgxLogLevelKey = "pgx_log_level" - type adapter struct { logger *slog.Logger } @@ -29,47 +26,39 @@ func New(logger *slog.Logger) tracelog.Logger { } } -func (l *adapter) Log( +func (a *adapter) Log( ctx context.Context, level tracelog.LogLevel, msg string, data map[string]any, ) { - attrs := make([]slog.Attr, 0, len(data)+1) + attrs := make([]slog.Attr, 0, len(data)) for key, value := range data { attrs = append(attrs, slog.Any(key, value)) } - var logLevel slog.Level + a.logger.LogAttrs( + ctx, + slogLevel(level), + msg, //nolint:sloglint // pgx provides the log message dynamically. + attrs..., + ) +} +func slogLevel(level tracelog.LogLevel) slog.Level { switch level { case tracelog.LogLevelTrace: - logLevel = slog.LevelDebug - 1 - attrs = append( - attrs, - slog.String(pgxLogLevelKey, level.String()), - ) + return slog.LevelDebug - 1 case tracelog.LogLevelDebug: - logLevel = slog.LevelDebug + return slog.LevelDebug case tracelog.LogLevelInfo: - logLevel = slog.LevelInfo + return slog.LevelInfo case tracelog.LogLevelWarn: - logLevel = slog.LevelWarn + return slog.LevelWarn case tracelog.LogLevelError: - logLevel = slog.LevelError + return slog.LevelError default: - logLevel = slog.LevelError - attrs = append( - attrs, - slog.String(pgxLogLevelKey, level.String()), - ) + return slog.LevelError } - - l.logger.LogAttrs( - ctx, - logLevel, - msg, - attrs..., - ) } diff --git a/extra/slogxpg/logger_test.go b/extra/slogxpg/logger_test.go index ca8aecc..cd8cbab 100644 --- a/extra/slogxpg/logger_test.go +++ b/extra/slogxpg/logger_test.go @@ -16,20 +16,18 @@ func TestNewNil(t *testing.T) { } } -func TestLoggerLevels(t *testing.T) { +func TestSlogLevel(t *testing.T) { t.Parallel() tests := []struct { - name string - pgx tracelog.LogLevel - want slog.Level - wantPGX string + name string + pgx tracelog.LogLevel + want slog.Level }{ { - name: "trace", - pgx: tracelog.LogLevelTrace, - want: slog.LevelDebug - 1, - wantPGX: tracelog.LogLevelTrace.String(), + name: "trace", + pgx: tracelog.LogLevelTrace, + want: slog.LevelDebug - 1, }, { name: "debug", @@ -52,10 +50,9 @@ func TestLoggerLevels(t *testing.T) { want: slog.LevelError, }, { - name: "unknown", - pgx: tracelog.LogLevel(255), - want: slog.LevelError, - wantPGX: tracelog.LogLevel(255).String(), + name: "unknown", + pgx: tracelog.LogLevel(255), + want: slog.LevelError, }, } @@ -63,45 +60,14 @@ func TestLoggerLevels(t *testing.T) { t.Run(test.name, func(t *testing.T) { t.Parallel() - handler := &captureHandler{} - logger := New(slog.New(handler)) - - logger.Log( - context.Background(), - test.pgx, - "message", - nil, - ) - - record := handler.record - - if record.Level != test.want { - t.Fatalf("level = %v, want %v", record.Level, test.want) - } - - attrs := recordAttrs(record) - - if test.wantPGX == "" { - if _, ok := attrs[pgxLogLevelKey]; ok { - t.Fatalf("unexpected %s attribute", pgxLogLevelKey) - } - - return - } - - if attrs[pgxLogLevelKey] != test.wantPGX { - t.Fatalf( - "%s = %v, want %q", - pgxLogLevelKey, - attrs[pgxLogLevelKey], - test.wantPGX, - ) + if got := slogLevel(test.pgx); got != test.want { + t.Fatalf("slogLevel(%v) = %v, want %v", test.pgx, got, test.want) } }) } } -func TestLoggerData(t *testing.T) { +func TestLoggerLog(t *testing.T) { t.Parallel() handler := &captureHandler{} @@ -123,7 +89,11 @@ func TestLoggerData(t *testing.T) { t.Fatalf("message = %q, want %q", record.Message, "query") } - attrs := recordAttrs(record) + if record.Level != slog.LevelInfo { + t.Fatalf("level = %v, want %v", record.Level, slog.LevelInfo) + } + + attrs := recordAttrs(&record) if attrs["sql"] != "select 1" { t.Fatalf("sql = %v, want %q", attrs["sql"], "select 1") @@ -134,6 +104,30 @@ func TestLoggerData(t *testing.T) { } } +func TestLoggerUnknownLevel(t *testing.T) { + t.Parallel() + + handler := &captureHandler{} + logger := New(slog.New(handler)) + + logger.Log( + context.Background(), + tracelog.LogLevel(255), + "unknown", + nil, + ) + + record := handler.record + + if record.Message != "unknown" { + t.Fatalf("message = %q, want %q", record.Message, "unknown") + } + + if record.Level != slog.LevelError { + t.Fatalf("level = %v, want %v", record.Level, slog.LevelError) + } +} + type captureHandler struct { record slog.Record } @@ -142,7 +136,10 @@ func (h *captureHandler) Enabled(context.Context, slog.Level) bool { return true } -func (h *captureHandler) Handle(_ context.Context, record slog.Record) error { +func (h *captureHandler) Handle( + _ context.Context, + record slog.Record, //nolint:gocritic // slog.Handler requires slog.Record by value. +) error { h.record = record.Clone() return nil @@ -156,7 +153,7 @@ func (h *captureHandler) WithGroup(string) slog.Handler { return h } -func recordAttrs(record slog.Record) map[string]any { +func recordAttrs(record *slog.Record) map[string]any { attrs := make(map[string]any, record.NumAttrs()) record.Attrs(func(attr slog.Attr) bool { From 00eee4a80ebbbc593264fafd225ddae09eb97060 Mon Sep 17 00:00:00 2001 From: mkbeh Date: Mon, 24 Aug 2026 21:02:11 +0300 Subject: [PATCH 21/41] refactor: simplify cluster routing --- cluster/cluster.go | 40 ++++++++++++++++------------------------ cluster/resolver.go | 38 ++++++++++++++++---------------------- cluster/selector.go | 12 ++++-------- options.go | 12 +++--------- 4 files changed, 39 insertions(+), 63 deletions(-) diff --git a/cluster/cluster.go b/cluster/cluster.go index f7b9911..050c14d 100644 --- a/cluster/cluster.go +++ b/cluster/cluster.go @@ -54,7 +54,7 @@ type Cluster struct { // At least one pool is required. When Selector is nil, replicas are selected // using round-robin. func New(config Config) (*Cluster, error) { - if config.Primary != nil && invalidPool(config.Primary) { + if config.Primary != nil && config.Primary.Raw() == nil { return nil, errors.New("xpg/cluster: primary pool is invalid") } @@ -70,13 +70,13 @@ func New(config Config) (*Cluster, error) { metadata := make(replicaMetadata, len(replicas)) for index, replica := range replicas { - if invalidPool(replica) { - return nil, fmt.Errorf("xpg/cluster: replica %d is nil", index) + if replica == nil || replica.Raw() == nil { + return nil, fmt.Errorf("pg/cluster: replica %d is invalid", index) } metadata[index] = ReplicaInfo{ name: replica.Name(), - labels: cloneLabels(replica.Labels()), + labels: replica.Labels(), } } @@ -95,24 +95,6 @@ func New(config Config) (*Cluster, error) { }, nil } -func invalidPool(pool *xpg.Pool) bool { - return pool == nil || pool.Raw() == nil -} - -func validateLabels(labels map[string]string) error { - for key, value := range labels { - if key == "" { - return errors.New("label key must not be empty") - } - - if value == "" { - return fmt.Errorf("label %q value must not be empty", key) - } - } - - return nil -} - // ID returns the stable logical cluster ID. func (c *Cluster) ID() ID { if c == nil { @@ -182,8 +164,8 @@ func (c *Cluster) Close() { } c.closeOnce.Do(func() { - for index := len(c.replicas) - 1; index >= 0; index-- { - c.replicas[index].Close() + for _, v := range slices.Backward(c.replicas) { + v.Close() } if c.primary != nil { @@ -191,3 +173,13 @@ func (c *Cluster) Close() { } }) } + +func validateLabels(labels map[string]string) error { + for key := range labels { + if key == "" { + return errors.New("label key must not be empty") + } + } + + return nil +} diff --git a/cluster/resolver.go b/cluster/resolver.go index 024b192..3294ba5 100644 --- a/cluster/resolver.go +++ b/cluster/resolver.go @@ -40,10 +40,7 @@ func ParsePolicy(value string) (ReadPolicy, error) { case readPolicyReplicaRequired: return ReadReplicaRequired, nil default: - return 0, fmt.Errorf( - "xpg/cluster: unknown read policy %q", - value, - ) + return 0, fmt.Errorf("xpg/cluster: unknown read policy %q", value) } } @@ -73,14 +70,10 @@ func (c *Cluster) ReadPool(ctx context.Context, policy ReadPolicy) (*xpg.Pool, e switch policy { case ReadPrimary: - if c.primary == nil { - return nil, ErrNoPrimary - } - - return c.primary, nil + return c.resolvePrimary() case ReadReplicaPreferred: - replica, err := c.selectReplica(ctx) + replica, err := c.resolveReplica(ctx) if err == nil { return replica, nil } @@ -89,24 +82,25 @@ func (c *Cluster) ReadPool(ctx context.Context, policy ReadPolicy) (*xpg.Pool, e return nil, err } - if c.primary == nil { - return nil, ErrNoPrimary - } - - return c.primary, nil + return c.resolvePrimary() case ReadReplicaRequired: - return c.selectReplica(ctx) + return c.resolveReplica(ctx) default: - return nil, fmt.Errorf( - "xpg/cluster: unsupported read policy %d", - policy, - ) + return nil, fmt.Errorf("xpg/cluster: unsupported read policy %d", policy) } } -func (c *Cluster) selectReplica(ctx context.Context) (*xpg.Pool, error) { +func (c *Cluster) resolvePrimary() (*xpg.Pool, error) { + if c.primary == nil { + return nil, ErrNoPrimary + } + + return c.primary, nil +} + +func (c *Cluster) resolveReplica(ctx context.Context) (*xpg.Pool, error) { if len(c.replicas) == 0 { return nil, ErrNoReplica } @@ -118,7 +112,7 @@ func (c *Cluster) selectReplica(ctx context.Context) (*xpg.Pool, error) { if index < 0 || index >= len(c.replicas) { return nil, fmt.Errorf( - "xpg/cluster: replica selector returned index %d for %d replicas", + "xpg/cluster: replica selector returned invalid index %d for %d replicas", index, len(c.replicas), ) diff --git a/cluster/selector.go b/cluster/selector.go index 5f0fb85..930abf9 100644 --- a/cluster/selector.go +++ b/cluster/selector.go @@ -77,11 +77,10 @@ func RoundRobinSelector() ReplicaSelector { func (selector *roundRobinSelector) Select(_ context.Context, replicas ReplicaSet) (int, error) { length := replicas.Len() - if length == 0 { + switch length { + case 0: return -1, ErrNoReplica - } - - if length == 1 { + case 1: return 0, nil } @@ -95,8 +94,5 @@ func cloneLabels(labels map[string]string) map[string]string { return nil } - cloned := make(map[string]string, len(labels)) - maps.Copy(cloned, labels) - - return cloned + return maps.Clone(labels) } diff --git a/options.go b/options.go index 64c96b9..6f44ed1 100644 --- a/options.go +++ b/options.go @@ -104,9 +104,8 @@ func WithLabels(labels map[string]string) Option { return optionFunc(func(settings *settings) error { for key, value := range labels { - key = strings.TrimSpace(key) if key == "" { - return errors.New("label key must not be blank") + return errors.New("label key must not be empty") } settings.labels[key] = value @@ -118,11 +117,9 @@ func WithLabels(labels map[string]string) Option { // WithLabel adds or replaces one pool label. func WithLabel(key, value string) Option { - key = strings.TrimSpace(key) - return optionFunc(func(settings *settings) error { if key == "" { - return errors.New("label key must not be blank") + return errors.New("label key must not be empty") } settings.labels[key] = value @@ -200,8 +197,5 @@ func cloneLabels(labels map[string]string) map[string]string { return nil } - cloned := make(map[string]string, len(labels)) - maps.Copy(cloned, labels) - - return cloned + return maps.Clone(labels) } From b943baa5ef428d9ba06d24d3f9813ad49580e089 Mon Sep 17 00:00:00 2001 From: mkbeh Date: Mon, 24 Aug 2026 21:18:29 +0300 Subject: [PATCH 22/41] test: add cluster coverage --- cluster/cluster.go | 2 +- cluster/cluster_test.go | 320 ++++++++++++++++++++++++++++++++ cluster/helpers_test.go | 51 ++++++ cluster/resolver_test.go | 385 +++++++++++++++++++++++++++++++++++++++ cluster/selector_test.go | 202 ++++++++++++++++++++ cluster/tx_test.go | 152 ++++++++++++++++ 6 files changed, 1111 insertions(+), 1 deletion(-) create mode 100644 cluster/cluster_test.go create mode 100644 cluster/helpers_test.go create mode 100644 cluster/resolver_test.go create mode 100644 cluster/selector_test.go create mode 100644 cluster/tx_test.go diff --git a/cluster/cluster.go b/cluster/cluster.go index 050c14d..1aa9648 100644 --- a/cluster/cluster.go +++ b/cluster/cluster.go @@ -71,7 +71,7 @@ func New(config Config) (*Cluster, error) { for index, replica := range replicas { if replica == nil || replica.Raw() == nil { - return nil, fmt.Errorf("pg/cluster: replica %d is invalid", index) + return nil, fmt.Errorf("xpg/cluster: replica %d is invalid", index) } metadata[index] = ReplicaInfo{ diff --git a/cluster/cluster_test.go b/cluster/cluster_test.go new file mode 100644 index 0000000..491a45e --- /dev/null +++ b/cluster/cluster_test.go @@ -0,0 +1,320 @@ +package cluster + +import ( + "context" + "testing" + + "github.com/mkbeh/xpg" +) + +func TestNewRequiresPool(t *testing.T) { + t.Parallel() + + cluster, err := New(Config{}) + if err == nil { + cluster.Close() + t.Fatal("expected error") + } + + if got, want := err.Error(), "xpg/cluster: at least one pool is required"; got != want { + t.Fatalf("error = %q, want %q", got, want) + } +} + +func TestNewRejectsInvalidPrimary(t *testing.T) { + t.Parallel() + + cluster, err := New(Config{ + Primary: &xpg.Pool{}, + }) + if err == nil { + cluster.Close() + t.Fatal("expected error") + } + + if got, want := err.Error(), "xpg/cluster: primary pool is invalid"; got != want { + t.Fatalf("error = %q, want %q", got, want) + } +} + +func TestNewRejectsInvalidReplica(t *testing.T) { + t.Parallel() + + tests := []struct { + name string + replica *xpg.Pool + }{ + { + name: "nil", + replica: nil, + }, + { + name: "zero value", + replica: &xpg.Pool{}, + }, + } + + for _, test := range tests { + t.Run(test.name, func(t *testing.T) { + t.Parallel() + + cluster, err := New(Config{ + Replicas: []*xpg.Pool{test.replica}, + }) + if err == nil { + cluster.Close() + t.Fatal("expected error") + } + + if got, want := err.Error(), "xpg/cluster: replica 0 is invalid"; got != want { + t.Fatalf("error = %q, want %q", got, want) + } + }) + } +} + +func TestNewLabels(t *testing.T) { + t.Parallel() + + primary := newTestPool(t, "primary", nil) + labels := map[string]string{ + "environment": "", + " ": "whitespace-key", + } + + cluster, err := New(Config{ + ID: "orders", + Labels: labels, + Primary: primary, + }) + if err != nil { + t.Fatalf("New() error = %v", err) + } + t.Cleanup(cluster.Close) + + labels["environment"] = "changed" + labels["new"] = "value" + + if got, want := cluster.ID(), ID("orders"); got != want { + t.Fatalf("ID() = %q, want %q", got, want) + } + + if got, ok := cluster.Label("environment"); !ok || got != "" { + t.Fatalf("Label(environment) = %q, %v; want empty value, true", got, ok) + } + + if got, ok := cluster.Label(" "); !ok || got != "whitespace-key" { + t.Fatalf("Label(whitespace) = %q, %v; want %q, true", got, ok, "whitespace-key") + } + + if _, ok := cluster.Label("new"); ok { + t.Fatal("cluster labels changed after input map mutation") + } + + cloned := cluster.Labels() + cloned["environment"] = "mutated" + delete(cloned, " ") + + if got, _ := cluster.Label("environment"); got != "" { + t.Fatalf("Label(environment) after Labels mutation = %q, want empty value", got) + } + + if got, ok := cluster.Label(" "); !ok || got != "whitespace-key" { + t.Fatalf("Label(whitespace) after Labels mutation = %q, %v; want %q, true", got, ok, "whitespace-key") + } +} + +func TestNewRejectsEmptyLabelKey(t *testing.T) { + t.Parallel() + + primary := newTestPool(t, "primary", nil) + + cluster, err := New(Config{ + Labels: map[string]string{ + "": "value", + }, + Primary: primary, + }) + if err == nil { + cluster.Close() + t.Fatal("expected error") + } + + if got, want := err.Error(), "xpg/cluster: label key must not be empty"; got != want { + t.Fatalf("error = %q, want %q", got, want) + } +} + +func TestNewClonesReplicaSlice(t *testing.T) { + t.Parallel() + + replicaA := newTestPool(t, "replica-a", nil) + replicaB := newTestPool(t, "replica-b", nil) + replicas := []*xpg.Pool{replicaA} + + cluster, err := New(Config{ + Replicas: replicas, + }) + if err != nil { + t.Fatalf("New() error = %v", err) + } + t.Cleanup(cluster.Close) + + replicas[0] = replicaB + + if got := cluster.ReplicaAt(0); got != replicaA { + t.Fatalf("ReplicaAt(0) = %p, want original replica %p", got, replicaA) + } +} + +func TestNewAllowsDuplicatePools(t *testing.T) { + t.Parallel() + + pool := newTestPool(t, "shared", nil) + + cluster, err := New(Config{ + Primary: pool, + Replicas: []*xpg.Pool{ + pool, + pool, + }, + }) + if err != nil { + t.Fatalf("New() error = %v", err) + } + t.Cleanup(cluster.Close) + + if cluster.Primary() != pool { + t.Fatal("Primary() did not return the configured pool") + } + + if got, want := cluster.ReplicaCount(), 2; got != want { + t.Fatalf("ReplicaCount() = %d, want %d", got, want) + } + + if cluster.ReplicaAt(0) != pool || cluster.ReplicaAt(1) != pool { + t.Fatal("duplicate topology entries were not preserved") + } +} + +func TestNewCapturesReplicaMetadata(t *testing.T) { + t.Parallel() + + replica := newTestPool(t, "replica-a", map[string]string{ + "region": "eu", + "role": "read", + }) + + selectorCalls := 0 + selector := ReplicaSelectorFunc(func(_ context.Context, replicas ReplicaSet) (int, error) { + selectorCalls++ + + if got, want := replicas.Len(), 1; got != want { + t.Fatalf("replicas.Len() = %d, want %d", got, want) + } + + info := replicas.At(0) + + if got, want := info.Name(), "replica-a"; got != want { + t.Fatalf("replica name = %q, want %q", got, want) + } + + if got, ok := info.Label("region"); !ok || got != "eu" { + t.Fatalf("region label = %q, %v; want %q, true", got, ok, "eu") + } + + labels := info.Labels() + labels["region"] = "mutated" + + if got, _ := info.Label("region"); got != "eu" { + t.Fatalf("replica metadata mutated through Labels(): got %q", got) + } + + return 0, nil + }) + + cluster, err := New(Config{ + Replicas: []*xpg.Pool{replica}, + Selector: selector, + }) + if err != nil { + t.Fatalf("New() error = %v", err) + } + t.Cleanup(cluster.Close) + + resolved, err := cluster.ReadPool(context.Background(), ReadReplicaRequired) + if err != nil { + t.Fatalf("ReadPool() error = %v", err) + } + + if resolved != replica { + t.Fatalf("ReadPool() = %p, want %p", resolved, replica) + } + + if selectorCalls != 1 { + t.Fatalf("selector calls = %d, want 1", selectorCalls) + } +} + +func TestClusterNilReceiverMetadata(t *testing.T) { + t.Parallel() + + var cluster *Cluster + + if got := cluster.ID(); got != "" { + t.Fatalf("ID() = %q, want empty", got) + } + + if value, ok := cluster.Label("key"); ok || value != "" { + t.Fatalf("Label() = %q, %v; want empty, false", value, ok) + } + + if labels := cluster.Labels(); labels != nil { + t.Fatalf("Labels() = %v, want nil", labels) + } + + if primary := cluster.Primary(); primary != nil { + t.Fatalf("Primary() = %p, want nil", primary) + } + + if got := cluster.ReplicaCount(); got != 0 { + t.Fatalf("ReplicaCount() = %d, want 0", got) + } + + cluster.Close() +} + +func TestReplicaAtPanicsOutOfRange(t *testing.T) { + t.Parallel() + + replica := newTestPool(t, "replica", nil) + cluster := newTestCluster(t, Config{ + Replicas: []*xpg.Pool{replica}, + }) + + defer func() { + if recover() == nil { + t.Fatal("ReplicaAt() did not panic") + } + }() + + _ = cluster.ReplicaAt(1) +} + +func TestCloseIsIdempotent(t *testing.T) { + t.Parallel() + + primary := newTestPool(t, "primary", nil) + replica := newTestPool(t, "replica", nil) + + cluster, err := New(Config{ + Primary: primary, + Replicas: []*xpg.Pool{replica}, + }) + if err != nil { + t.Fatalf("New() error = %v", err) + } + + cluster.Close() + cluster.Close() +} diff --git a/cluster/helpers_test.go b/cluster/helpers_test.go new file mode 100644 index 0000000..77f5411 --- /dev/null +++ b/cluster/helpers_test.go @@ -0,0 +1,51 @@ +package cluster + +import ( + "context" + "testing" + + "github.com/jackc/pgx/v5/pgxpool" + "github.com/mkbeh/xpg" +) + +const testDatabaseURL = "postgres://postgres:postgres@127.0.0.1:1/postgres?sslmode=disable" //nolint:gosec // Test-only DSN with non-production credentials. + +func newTestPool(t *testing.T, name string, labels map[string]string) *xpg.Pool { + t.Helper() + + config, err := pgxpool.ParseConfig(testDatabaseURL) + if err != nil { + t.Fatalf("pgxpool.ParseConfig() error = %v", err) + } + + config.MinConns = 0 + config.MaxConns = 1 + + options := []xpg.Option{ + xpg.WithName(name), + } + + if len(labels) != 0 { + options = append(options, xpg.WithLabels(labels)) + } + + pool, err := xpg.New(context.Background(), config, options...) + if err != nil { + t.Fatalf("xpg.New() error = %v", err) + } + t.Cleanup(pool.Close) + + return pool +} + +func newTestCluster(t *testing.T, config Config) *Cluster { + t.Helper() + + cluster, err := New(config) + if err != nil { + t.Fatalf("New() error = %v", err) + } + t.Cleanup(cluster.Close) + + return cluster +} diff --git a/cluster/resolver_test.go b/cluster/resolver_test.go new file mode 100644 index 0000000..7d2fc6c --- /dev/null +++ b/cluster/resolver_test.go @@ -0,0 +1,385 @@ +package cluster + +import ( + "context" + "errors" + "strings" + "testing" + + "github.com/mkbeh/xpg" +) + +func TestParsePolicy(t *testing.T) { + t.Parallel() + + tests := []struct { + value string + want ReadPolicy + }{ + {value: "primary", want: ReadPrimary}, + {value: "replica_preferred", want: ReadReplicaPreferred}, + {value: "replica_required", want: ReadReplicaRequired}, + } + + for _, test := range tests { + t.Run(test.value, func(t *testing.T) { + t.Parallel() + + got, err := ParsePolicy(test.value) + if err != nil { + t.Fatalf("ParsePolicy() error = %v", err) + } + + if got != test.want { + t.Fatalf("ParsePolicy(%q) = %v, want %v", test.value, got, test.want) + } + }) + } +} + +func TestParsePolicyRejectsUnknown(t *testing.T) { + t.Parallel() + + _, err := ParsePolicy("nearest") + if err == nil { + t.Fatal("expected error") + } + + if got, want := err.Error(), `xpg/cluster: unknown read policy "nearest"`; got != want { + t.Fatalf("error = %q, want %q", got, want) + } +} + +func TestReadPolicyString(t *testing.T) { + t.Parallel() + + tests := []struct { + policy ReadPolicy + want string + }{ + {policy: ReadPrimary, want: "primary"}, + {policy: ReadReplicaPreferred, want: "replica_preferred"}, + {policy: ReadReplicaRequired, want: "replica_required"}, + {policy: ReadPolicy(255), want: "unknown"}, + } + + for _, test := range tests { + if got := test.policy.String(); got != test.want { + t.Fatalf("ReadPolicy(%d).String() = %q, want %q", test.policy, got, test.want) + } + } +} + +func TestReadPoolNilCluster(t *testing.T) { + t.Parallel() + + var cluster *Cluster + + pool, err := cluster.ReadPool(context.Background(), ReadPrimary) + if pool != nil { + t.Fatalf("ReadPool() pool = %p, want nil", pool) + } + + if err == nil { + t.Fatal("expected error") + } + + if got, want := err.Error(), "xpg/cluster: cluster is nil"; got != want { + t.Fatalf("error = %q, want %q", got, want) + } +} + +func TestReadPoolPrimary(t *testing.T) { + t.Parallel() + + primary := newTestPool(t, "primary", nil) + cluster := newTestCluster(t, Config{Primary: primary}) + + got, err := cluster.ReadPool(context.Background(), ReadPrimary) + if err != nil { + t.Fatalf("ReadPool() error = %v", err) + } + + if got != primary { + t.Fatalf("ReadPool() = %p, want %p", got, primary) + } +} + +func TestReadPoolPrimaryWithoutPrimary(t *testing.T) { + t.Parallel() + + replica := newTestPool(t, "replica", nil) + cluster := newTestCluster(t, Config{ + Replicas: []*xpg.Pool{replica}, + }) + + pool, err := cluster.ReadPool(context.Background(), ReadPrimary) + if pool != nil { + t.Fatalf("ReadPool() pool = %p, want nil", pool) + } + + if !errors.Is(err, ErrNoPrimary) { + t.Fatalf("ReadPool() error = %v, want ErrNoPrimary", err) + } +} + +func TestReadPoolReplicaPreferredUsesReplica(t *testing.T) { + t.Parallel() + + primary := newTestPool(t, "primary", nil) + replica := newTestPool(t, "replica", nil) + cluster := newTestCluster(t, Config{ + Primary: primary, + Replicas: []*xpg.Pool{replica}, + }) + + got, err := cluster.ReadPool(context.Background(), ReadReplicaPreferred) + if err != nil { + t.Fatalf("ReadPool() error = %v", err) + } + + if got != replica { + t.Fatalf("ReadPool() = %p, want replica %p", got, replica) + } +} + +func TestReadPoolReplicaPreferredFallsBackWithoutReplicas(t *testing.T) { + t.Parallel() + + primary := newTestPool(t, "primary", nil) + cluster := newTestCluster(t, Config{Primary: primary}) + + got, err := cluster.ReadPool(context.Background(), ReadReplicaPreferred) + if err != nil { + t.Fatalf("ReadPool() error = %v", err) + } + + if got != primary { + t.Fatalf("ReadPool() = %p, want primary %p", got, primary) + } +} + +func TestReadPoolReplicaPreferredFallsBackOnErrNoReplica(t *testing.T) { + t.Parallel() + + primary := newTestPool(t, "primary", nil) + replica := newTestPool(t, "replica", nil) + selector := ReplicaSelectorFunc(func(context.Context, ReplicaSet) (int, error) { + return -1, errors.Join(errors.New("temporarily unavailable"), ErrNoReplica) + }) + + cluster := newTestCluster(t, Config{ + Primary: primary, + Replicas: []*xpg.Pool{replica}, + Selector: selector, + }) + + got, err := cluster.ReadPool(context.Background(), ReadReplicaPreferred) + if err != nil { + t.Fatalf("ReadPool() error = %v", err) + } + + if got != primary { + t.Fatalf("ReadPool() = %p, want primary %p", got, primary) + } +} + +func TestReadPoolReplicaPreferredDoesNotFallbackOnSelectorError(t *testing.T) { + t.Parallel() + + errSelector := errors.New("selector failed") + primary := newTestPool(t, "primary", nil) + replica := newTestPool(t, "replica", nil) + selector := ReplicaSelectorFunc(func(context.Context, ReplicaSet) (int, error) { + return -1, errSelector + }) + + cluster := newTestCluster(t, Config{ + Primary: primary, + Replicas: []*xpg.Pool{replica}, + Selector: selector, + }) + + pool, err := cluster.ReadPool(context.Background(), ReadReplicaPreferred) + if pool != nil { + t.Fatalf("ReadPool() pool = %p, want nil", pool) + } + + if !errors.Is(err, errSelector) { + t.Fatalf("ReadPool() error = %v, want selector error", err) + } +} + +func TestReadPoolReplicaPreferredWithoutPrimary(t *testing.T) { + t.Parallel() + + replica := newTestPool(t, "replica", nil) + selector := ReplicaSelectorFunc(func(context.Context, ReplicaSet) (int, error) { + return -1, ErrNoReplica + }) + + cluster := newTestCluster(t, Config{ + Replicas: []*xpg.Pool{replica}, + Selector: selector, + }) + + pool, err := cluster.ReadPool(context.Background(), ReadReplicaPreferred) + if pool != nil { + t.Fatalf("ReadPool() pool = %p, want nil", pool) + } + + if !errors.Is(err, ErrNoPrimary) { + t.Fatalf("ReadPool() error = %v, want ErrNoPrimary", err) + } +} + +func TestReadPoolReplicaRequired(t *testing.T) { + t.Parallel() + + replica := newTestPool(t, "replica", nil) + cluster := newTestCluster(t, Config{ + Replicas: []*xpg.Pool{replica}, + }) + + got, err := cluster.ReadPool(context.Background(), ReadReplicaRequired) + if err != nil { + t.Fatalf("ReadPool() error = %v", err) + } + + if got != replica { + t.Fatalf("ReadPool() = %p, want %p", got, replica) + } +} + +func TestReadPoolReplicaRequiredWithoutReplicas(t *testing.T) { + t.Parallel() + + primary := newTestPool(t, "primary", nil) + cluster := newTestCluster(t, Config{Primary: primary}) + + pool, err := cluster.ReadPool(context.Background(), ReadReplicaRequired) + if pool != nil { + t.Fatalf("ReadPool() pool = %p, want nil", pool) + } + + if !errors.Is(err, ErrNoReplica) { + t.Fatalf("ReadPool() error = %v, want ErrNoReplica", err) + } +} + +func TestReadPoolDefaultSelectorRoundRobin(t *testing.T) { + t.Parallel() + + replicaA := newTestPool(t, "replica-a", nil) + replicaB := newTestPool(t, "replica-b", nil) + cluster := newTestCluster(t, Config{ + Replicas: []*xpg.Pool{replicaA, replicaB}, + }) + + want := []*xpg.Pool{ + replicaA, + replicaB, + replicaA, + replicaB, + } + + for call, wantPool := range want { + got, err := cluster.ReadPool(context.Background(), ReadReplicaRequired) + if err != nil { + t.Fatalf("ReadPool() call %d error = %v", call, err) + } + + if got != wantPool { + t.Fatalf("ReadPool() call %d = %p, want %p", call, got, wantPool) + } + } +} + +func TestReadPoolRejectsSelectorIndex(t *testing.T) { + t.Parallel() + + tests := []struct { + name string + index int + }{ + {name: "negative", index: -1}, + {name: "past end", index: 1}, + } + + for _, test := range tests { + t.Run(test.name, func(t *testing.T) { + t.Parallel() + + replica := newTestPool(t, "replica", nil) + selector := ReplicaSelectorFunc(func(context.Context, ReplicaSet) (int, error) { + return test.index, nil + }) + + cluster := newTestCluster(t, Config{ + Replicas: []*xpg.Pool{replica}, + Selector: selector, + }) + + pool, err := cluster.ReadPool(context.Background(), ReadReplicaRequired) + if pool != nil { + t.Fatalf("ReadPool() pool = %p, want nil", pool) + } + + if err == nil { + t.Fatal("expected error") + } + + if !strings.Contains(err.Error(), "replica selector returned invalid index") { + t.Fatalf("error = %q, want invalid replica index error", err) + } + }) + } +} + +func TestReadPoolPreservesSelectorError(t *testing.T) { + t.Parallel() + + errSelector := errors.New("boom") + replica := newTestPool(t, "replica", nil) + selector := ReplicaSelectorFunc(func(context.Context, ReplicaSet) (int, error) { + return -1, errSelector + }) + + cluster := newTestCluster(t, Config{ + Replicas: []*xpg.Pool{replica}, + Selector: selector, + }) + + pool, err := cluster.ReadPool(context.Background(), ReadReplicaRequired) + if pool != nil { + t.Fatalf("ReadPool() pool = %p, want nil", pool) + } + + if !errors.Is(err, errSelector) { + t.Fatalf("ReadPool() error = %v, want wrapped selector error", err) + } + + if !strings.Contains(err.Error(), "xpg/cluster: select replica") { + t.Fatalf("error = %q, want selector context", err) + } +} + +func TestReadPoolRejectsUnsupportedPolicy(t *testing.T) { + t.Parallel() + + primary := newTestPool(t, "primary", nil) + cluster := newTestCluster(t, Config{Primary: primary}) + + pool, err := cluster.ReadPool(context.Background(), ReadPolicy(255)) + if pool != nil { + t.Fatalf("ReadPool() pool = %p, want nil", pool) + } + + if err == nil { + t.Fatal("expected error") + } + + if got, want := err.Error(), "xpg/cluster: unsupported read policy 255"; got != want { + t.Fatalf("error = %q, want %q", got, want) + } +} diff --git a/cluster/selector_test.go b/cluster/selector_test.go new file mode 100644 index 0000000..de5005b --- /dev/null +++ b/cluster/selector_test.go @@ -0,0 +1,202 @@ +package cluster + +import ( + "context" + "errors" + "sync" + "testing" +) + +func TestReplicaInfoLabelsAreDefensive(t *testing.T) { + t.Parallel() + + info := ReplicaInfo{ + name: "replica", + labels: map[string]string{ + "region": "eu", + }, + } + + labels := info.Labels() + labels["region"] = "us" + labels["new"] = "value" + + if got, want := info.Name(), "replica"; got != want { + t.Fatalf("Name() = %q, want %q", got, want) + } + + if got, ok := info.Label("region"); !ok || got != "eu" { + t.Fatalf("Label(region) = %q, %v; want %q, true", got, ok, "eu") + } + + if _, ok := info.Label("new"); ok { + t.Fatal("Labels() exposed internal map") + } +} + +func TestReplicaMetadataAtPanicsOutOfRange(t *testing.T) { + t.Parallel() + + replicas := replicaMetadata{{name: "replica"}} + + defer func() { + if recover() == nil { + t.Fatal("At() did not panic") + } + }() + + _ = replicas.At(1) +} + +func TestReplicaSelectorFunc(t *testing.T) { + t.Parallel() + + type contextKey struct{} + ctx := context.WithValue(context.Background(), contextKey{}, "value") + replicas := replicaMetadata{{name: "replica"}} + + selector := ReplicaSelectorFunc(func(gotCtx context.Context, gotReplicas ReplicaSet) (int, error) { + if got := gotCtx.Value(contextKey{}); got != "value" { + t.Fatalf("context value = %v, want %q", got, "value") + } + + if got, want := gotReplicas.Len(), 1; got != want { + t.Fatalf("replicas.Len() = %d, want %d", got, want) + } + + return 0, nil + }) + + index, err := selector.Select(ctx, replicas) + if err != nil { + t.Fatalf("Select() error = %v", err) + } + + if index != 0 { + t.Fatalf("Select() index = %d, want 0", index) + } +} + +func TestReplicaSelectorFuncNil(t *testing.T) { + t.Parallel() + + var selector ReplicaSelectorFunc + + index, err := selector.Select(context.Background(), nil) + if index != -1 { + t.Fatalf("Select() index = %d, want -1", index) + } + + if err == nil { + t.Fatal("expected error") + } + + if got, want := err.Error(), "xpg/cluster: replica selector function is nil"; got != want { + t.Fatalf("error = %q, want %q", got, want) + } +} + +func TestRoundRobinSelectorEmpty(t *testing.T) { + t.Parallel() + + selector := RoundRobinSelector() + + index, err := selector.Select(context.Background(), replicaMetadata(nil)) + if index != -1 { + t.Fatalf("Select() index = %d, want -1", index) + } + + if !errors.Is(err, ErrNoReplica) { + t.Fatalf("Select() error = %v, want ErrNoReplica", err) + } +} + +func TestRoundRobinSelectorSingleReplica(t *testing.T) { + t.Parallel() + + selector := RoundRobinSelector() + replicas := replicaMetadata{{name: "replica"}} + + for range 10 { + index, err := selector.Select(context.Background(), replicas) + if err != nil { + t.Fatalf("Select() error = %v", err) + } + + if index != 0 { + t.Fatalf("Select() index = %d, want 0", index) + } + } +} + +func TestRoundRobinSelectorSequence(t *testing.T) { + t.Parallel() + + selector := RoundRobinSelector() + replicas := replicaMetadata{ + {name: "replica-a"}, + {name: "replica-b"}, + {name: "replica-c"}, + } + want := []int{0, 1, 2, 0, 1, 2, 0} + + for call, wantIndex := range want { + index, err := selector.Select(context.Background(), replicas) + if err != nil { + t.Fatalf("Select() call %d error = %v", call, err) + } + + if index != wantIndex { + t.Fatalf("Select() call %d index = %d, want %d", call, index, wantIndex) + } + } +} + +func TestRoundRobinSelectorConcurrent(t *testing.T) { + t.Parallel() + + const ( + replicaCount = 3 + callCount = 600 + ) + + selector := RoundRobinSelector() + replicas := make(replicaMetadata, replicaCount) + results := make(chan int, callCount) + + var waitGroup sync.WaitGroup + waitGroup.Add(callCount) + + for range callCount { + go func() { + defer waitGroup.Done() + + index, err := selector.Select(context.Background(), replicas) + if err != nil { + results <- -1 + return + } + + results <- index + }() + } + + waitGroup.Wait() + close(results) + + counts := make([]int, replicaCount) + + for index := range results { + if index < 0 || index >= replicaCount { + t.Fatalf("Select() returned invalid index %d", index) + } + + counts[index]++ + } + + for index, count := range counts { + if count != callCount/replicaCount { + t.Fatalf("replica %d selections = %d, want %d", index, count, callCount/replicaCount) + } + } +} diff --git a/cluster/tx_test.go b/cluster/tx_test.go new file mode 100644 index 0000000..e839f27 --- /dev/null +++ b/cluster/tx_test.go @@ -0,0 +1,152 @@ +package cluster + +import ( + "context" + "errors" + "testing" + + "github.com/jackc/pgx/v5" + "github.com/mkbeh/xpg" +) + +func TestInPrimaryTxNilCluster(t *testing.T) { + t.Parallel() + + var cluster *Cluster + called := false + + err := cluster.InPrimaryTx( + context.Background(), + pgx.TxOptions{}, + func(context.Context, pgx.Tx) error { + called = true + return nil + }, + ) + if err == nil { + t.Fatal("expected error") + } + + if got, want := err.Error(), "xpg/cluster: cluster is nil"; got != want { + t.Fatalf("error = %q, want %q", got, want) + } + + if called { + t.Fatal("transaction callback was called") + } +} + +func TestInPrimaryTxWithoutPrimary(t *testing.T) { + t.Parallel() + + replica := newTestPool(t, "replica", nil) + cluster := newTestCluster(t, Config{ + Replicas: []*xpg.Pool{replica}, + }) + + called := false + err := cluster.InPrimaryTx( + context.Background(), + pgx.TxOptions{}, + func(context.Context, pgx.Tx) error { + called = true + return nil + }, + ) + + if !errors.Is(err, ErrNoPrimary) { + t.Fatalf("InPrimaryTx() error = %v, want ErrNoPrimary", err) + } + + if called { + t.Fatal("transaction callback was called") + } +} + +func TestInPrimaryTxDelegatesToPool(t *testing.T) { + t.Parallel() + + primary := newTestPool(t, "primary", nil) + cluster := newTestCluster(t, Config{Primary: primary}) + + ctx, cancel := context.WithCancel(context.Background()) + cancel() + + called := false + err := cluster.InPrimaryTx( + ctx, + pgx.TxOptions{}, + func(context.Context, pgx.Tx) error { + called = true + return nil + }, + ) + + if err == nil { + t.Fatal("expected error") + } + + if called { + t.Fatal("transaction callback was called") + } +} + +func TestInReadTxRoutingError(t *testing.T) { + t.Parallel() + + primary := newTestPool(t, "primary", nil) + cluster := newTestCluster(t, Config{Primary: primary}) + + called := false + err := cluster.InReadTx( + context.Background(), + ReadReplicaRequired, + ReadTxOptions{}, + func(context.Context, pgx.Tx) error { + called = true + return nil + }, + ) + + if !errors.Is(err, ErrNoReplica) { + t.Fatalf("InReadTx() error = %v, want ErrNoReplica", err) + } + + if called { + t.Fatal("transaction callback was called") + } +} + +func TestInReadTxDelegatesToResolvedPool(t *testing.T) { + t.Parallel() + + replica := newTestPool(t, "replica", nil) + cluster := newTestCluster(t, Config{ + Replicas: []*xpg.Pool{replica}, + }) + + ctx, cancel := context.WithCancel(context.Background()) + cancel() + + called := false + err := cluster.InReadTx( + ctx, + ReadReplicaRequired, + ReadTxOptions{ + IsoLevel: pgx.Serializable, + DeferrableMode: pgx.Deferrable, + }, + func(context.Context, pgx.Tx) error { + called = true + return nil + }, + ) + + if err == nil { + t.Fatal("expected error") + } + + if called { + t.Fatal("transaction callback was called") + } +} From cae2772d5dd11a7ff920f05de72329506836e234 Mon Sep 17 00:00:00 2001 From: mkbeh Date: Tue, 25 Aug 2026 00:58:46 +0300 Subject: [PATCH 23/41] refactor: simplify shard routing --- shard/doc.go | 4 +-- shard/foreach.go | 32 ++++++++++++------------ shard/group.go | 29 +++++++++++----------- shard/resolver.go | 1 + shard/resolver/custom.go | 17 +++++-------- shard/resolver/doc.go | 4 +-- shard/resolver/encoder.go | 7 +++--- shard/resolver/hash.go | 47 +++++++++++++++--------------------- shard/resolver/range.go | 18 +++++--------- shard/resolver/time_range.go | 23 ++++++------------ shard/topology.go | 40 ++++++++++-------------------- 11 files changed, 90 insertions(+), 132 deletions(-) diff --git a/shard/doc.go b/shard/doc.go index 55fad02..9fa129f 100644 --- a/shard/doc.go +++ b/shard/doc.go @@ -7,8 +7,8 @@ // ranges; applications may also provide custom routing logic. Resolvers borrow // their topology and do not own its clusters. // -// The package also provides shard grouping, colocation checks, bounded fan-out, -// and connection-budget diagnostics. +// The package also provides shard grouping, colocation checks, and bounded +// fan-out. // // The package does not inspect SQL, hide shard keys in contexts, move data, // replicate reference tables, or provide distributed transactions. diff --git a/shard/foreach.go b/shard/foreach.go index 41cc265..fc6d5c3 100644 --- a/shard/foreach.go +++ b/shard/foreach.go @@ -18,7 +18,7 @@ type ForEachShardResults []ForEachShardResult // Err returns all shard failures joined in registration order. func (results ForEachShardResults) Err() error { - errs := make([]error, 0, len(results)) + var errs []error for _, result := range results { if result.Err == nil { @@ -50,7 +50,7 @@ func (t *Topology) ForEachShard( fn func(context.Context, Shard) error, ) (ForEachShardResults, error) { if t == nil || len(t.shards) == 0 { - return nil, errors.New("xpg/shard: topology is nil or empty") + return nil, errors.New("xpg/shard: topology is nil") } if concurrency <= 0 { @@ -62,7 +62,6 @@ func (t *Topology) ForEachShard( } results := make(ForEachShardResults, len(t.shards)) - for index, shard := range t.shards { results[index].ShardID = shard.ID() } @@ -71,10 +70,15 @@ func (t *Topology) ForEachShard( jobs := make(chan int) var workers sync.WaitGroup + workers.Add(workerCount) for range workerCount { - workers.Go(func() { + go func() { + defer workers.Done() + for index := range jobs { + // An index may have been scheduled immediately before context + // cancellation. Skip callbacks that have not started yet. if err := ctx.Err(); err != nil { results[index].Err = err continue @@ -82,34 +86,28 @@ func (t *Topology) ForEachShard( results[index].Err = fn(ctx, t.shards[index]) } - }) + }() } nextIndex := 0 -schedule: - for nextIndex < len(t.shards) { - // Check cancellation before entering select so that a ready worker does - // not repeatedly win against an already canceled context. - if ctx.Err() != nil { - break - } - + for nextIndex < len(t.shards) && ctx.Err() == nil { + // The explicit context check above prevents scheduling new work after + // cancellation has already been observed. The select still handles + // cancellation that happens while waiting for a worker. select { case jobs <- nextIndex: nextIndex++ - case <-ctx.Done(): - break schedule } } close(jobs) workers.Wait() + // Workers own results for scheduled indexes [0, nextIndex). After all + // workers finish, remaining indexes can be marked canceled without races. if err := ctx.Err(); err != nil { - // Every index before nextIndex was handed to exactly one worker. - // Everything from nextIndex onward was never scheduled. for index := nextIndex; index < len(results); index++ { results[index].Err = err } diff --git a/shard/group.go b/shard/group.go index 870ad17..59116ec 100644 --- a/shard/group.go +++ b/shard/group.go @@ -24,18 +24,20 @@ func SameShard[K any](resolver Resolver[K], keys ...K) (Shard, error) { expectedID := expected.ID() for index := 1; index < len(keys); index++ { - actual, resolveErr := resolver.Resolve(keys[index]) - if resolveErr != nil { - return Shard{}, fmt.Errorf("xpg/shard: resolve key %d: %w", index, resolveErr) + actual, err := resolver.Resolve(keys[index]) + if err != nil { + return Shard{}, fmt.Errorf("xpg/shard: resolve key %d: %w", index, err) } actualID := actual.ID() - if actualID != expectedID { - return Shard{}, &MismatchError{ - Expected: expectedID, - Actual: actualID, - Index: index, - } + if actualID == expectedID { + continue + } + + return Shard{}, &MismatchError{ + Expected: expectedID, + Actual: actualID, + Index: index, } } @@ -72,12 +74,9 @@ func GroupByShard[K any](resolver Resolver[K], keys []K) ([]Group[K], error) { groupIndex = len(groups) indexByID[id] = groupIndex - groups = append( - groups, - Group[K]{ - Shard: resolved, - }, - ) + groups = append(groups, Group[K]{ + Shard: resolved, + }) } groups[groupIndex].Keys = append(groups[groupIndex].Keys, key) diff --git a/shard/resolver.go b/shard/resolver.go index 67f287a..e7e6927 100644 --- a/shard/resolver.go +++ b/shard/resolver.go @@ -2,6 +2,7 @@ package shard // Resolver maps a typed application key to one shard. // +// Resolve should return ErrNoShard when the key cannot be mapped to a shard. // Implementations shared by concurrent callers must be concurrency-safe. type Resolver[K any] interface { Resolve(key K) (Shard, error) diff --git a/shard/resolver/custom.go b/shard/resolver/custom.go index 7c07c15..b5e04f4 100644 --- a/shard/resolver/custom.go +++ b/shard/resolver/custom.go @@ -8,8 +8,9 @@ import ( // ResolveFunc maps an application key to a shard ID within topology. // -// Implementations shared by concurrent callers must be deterministic and -// concurrency-safe. They should not perform hidden I/O or modify topology. +// Resolve functions should return shard.ErrNoShard when a key cannot be mapped +// to a shard. Implementations shared by concurrent callers must be deterministic +// and concurrency-safe. They should not perform hidden I/O or modify topology. type ResolveFunc[K any] func(key K, topology *shard.Topology) (shard.ID, error) // CustomResolver adapts ResolveFunc to shard.Resolver. @@ -36,12 +37,8 @@ func NewCustom[K any](topology *shard.Topology, resolve ResolveFunc[K]) (*Custom // Resolve maps key to a shard and rejects IDs absent from the bound topology. func (resolver *CustomResolver[K]) Resolve(key K) (shard.Shard, error) { - if resolver == nil || - resolver.topology == nil || - resolver.resolve == nil { - return shard.Shard{}, errors.New( - "xpg/shard/resolver: custom resolver is not initialized", - ) + if resolver == nil || resolver.topology == nil || resolver.resolve == nil { + return shard.Shard{}, errors.New("xpg/shard/resolver: custom resolver is not initialized") } id, err := resolver.resolve(key, resolver.topology) @@ -51,9 +48,7 @@ func (resolver *CustomResolver[K]) Resolve(key K) (shard.Shard, error) { resolved, ok := resolver.topology.Shard(id) if !ok { - return shard.Shard{}, &shard.UnknownShardError{ - ShardID: id, - } + return shard.Shard{}, &shard.UnknownShardError{ShardID: id} } return resolved, nil diff --git a/shard/resolver/doc.go b/shard/resolver/doc.go index 047e8ca..7138e4e 100644 --- a/shard/resolver/doc.go +++ b/shard/resolver/doc.go @@ -1,7 +1,7 @@ // Package resolver provides routing strategies for shard.Topology. // -// Resolvers are bound to an immutable topology and return shard.Shard values. -// The package provides rendezvous hashing, ordered numeric or string ranges, +// Resolvers are bound to an immutable topology and map application keys to +// shard.Shard values. The package provides rendezvous hashing, ordered ranges, // time ranges, and an adapter for custom routing functions. // // Resolvers borrow their topology and must not outlive it. diff --git a/shard/resolver/encoder.go b/shard/resolver/encoder.go index a0a1c3e..1009493 100644 --- a/shard/resolver/encoder.go +++ b/shard/resolver/encoder.go @@ -8,9 +8,10 @@ import ( // KeyEncoder converts a typed key into stable canonical bytes. // -// Implementations used for persistent shard placement must remain -// deterministic across processes and releases. Changing an encoder changes -// hash placement and may require data migration. +// Implementations used for persistent shard placement must remain deterministic +// across processes and releases. Implementations shared by concurrent Resolve +// calls must be concurrency-safe. Changing an encoder changes hash placement +// and may require data migration. type KeyEncoder[K any] interface { Encode(K) ([]byte, error) } diff --git a/shard/resolver/hash.go b/shard/resolver/hash.go index 6bb3e70..08cac0a 100644 --- a/shard/resolver/hash.go +++ b/shard/resolver/hash.go @@ -7,7 +7,6 @@ import ( "errors" "fmt" "math" - "strings" "github.com/mkbeh/xpg/shard" ) @@ -47,14 +46,8 @@ func NewHash[K any]( return nil, errors.New("xpg/shard/resolver: key encoder is nil") } - trimmedNamespace := strings.TrimSpace(namespace) - - if trimmedNamespace == "" { - return nil, errors.New("xpg/shard/resolver: hash namespace must not be blank") - } - - if trimmedNamespace != namespace { - return nil, errors.New("xpg/shard/resolver: hash namespace must not contain surrounding whitespace") + if namespace == "" { + return nil, errors.New("xpg/shard/resolver: hash namespace must not be empty") } if len(namespace) > math.MaxUint32 { @@ -71,23 +64,30 @@ func NewHash[K any]( return nil, errors.New("xpg/shard/resolver: shard ID is too large") } + // Resolve reuses one score-input buffer for all candidates, sized for + // the largest shard ID in the topology. maxIDLength = max(maxIDLength, len(id)) } - // Prefix is invariant for the lifetime of the resolver: + // Prefix is part of the persistent placement format: // // domain || namespace_length || namespace - prefix := make([]byte, len(rendezvousDomain)+rendezvousLengthSize+len(namespace)) + // + // Changing this layout changes shard placement and requires data migration. + prefixSize := len(rendezvousDomain) + rendezvousLengthSize + len(namespace) + prefix := make([]byte, prefixSize) - offset := copy(prefix, rendezvousDomain) + lengthOffset := len(rendezvousDomain) + namespaceOffset := lengthOffset + rendezvousLengthSize + + copy(prefix, rendezvousDomain) binary.BigEndian.PutUint32( - prefix[offset:offset+rendezvousLengthSize], + prefix[lengthOffset:namespaceOffset], uint32(len(namespace)), ) - offset += rendezvousLengthSize - copy(prefix[offset:], namespace) + copy(prefix[namespaceOffset:], namespace) return &HashResolver[K]{ shards: shards, @@ -97,14 +97,8 @@ func NewHash[K any]( }, nil } -// Resolve selects the shard with the lexicographically greatest SHA-256 score. -// -// Resolve performs only in-memory routing. It does not acquire a connection or -// execute a PostgreSQL query. func (resolver *HashResolver[K]) Resolve(key K) (shard.Shard, error) { - if resolver == nil || - len(resolver.shards) == 0 || - resolver.encoder == nil { + if resolver == nil || len(resolver.shards) == 0 || resolver.encoder == nil { return shard.Shard{}, errors.New("xpg/shard/resolver: hash resolver is not initialized") } @@ -117,13 +111,13 @@ func (resolver *HashResolver[K]) Resolve(key K) (shard.Shard, error) { return shard.Shard{}, errors.New("xpg/shard/resolver: encoded key is too large") } - // Build the candidate-independent prefix once. The shard ID suffix is - // overwritten for each candidate below. keyLengthOffset := len(resolver.prefix) keyOffset := keyLengthOffset + rendezvousLengthSize idLengthOffset := keyOffset + len(encoded) idOffset := idLengthOffset + rendezvousLengthSize + // The candidate-independent part is written once. Only the shard ID suffix + // is overwritten while evaluating rendezvous scores. scoreInput := make([]byte, idOffset+resolver.maxIDLength) copy(scoreInput, resolver.prefix) @@ -154,14 +148,11 @@ func (resolver *HashResolver[K]) Resolve(key K) (shard.Shard, error) { copy(scoreInput[idOffset:inputEnd], candidateID) score := sha256.Sum256(scoreInput[:inputEnd]) - comparison := bytes.Compare(score[:], best[:]) // Shard ID is the deterministic tie-breaker, so placement does not // depend on topology registration order when scores are equal. - if !hasBest || - comparison > 0 || - comparison == 0 && candidateID < bestID { + if !hasBest || comparison > 0 || (comparison == 0 && candidateID < bestID) { selected = candidate best = score bestID = candidateID diff --git a/shard/resolver/range.go b/shard/resolver/range.go index 70f0aa2..67595fe 100644 --- a/shard/resolver/range.go +++ b/shard/resolver/range.go @@ -54,9 +54,9 @@ func NewRange[K cmp.Ordered](topology *shard.Topology, ranges []Range[K]) (*Rang return nil, fmt.Errorf("xpg/shard/resolver: range %d: %w", index, err) } - // This rejects empty, reversed, and NaN-bounded ranges. - valid := valueRange.Start < valueRange.End - if !valid { + // Using < intentionally rejects empty and reversed ranges as well as + // ranges with NaN boundaries for floating-point key types. + if !(valueRange.Start < valueRange.End) { //nolint:staticcheck // Negated comparison intentionally rejects NaN boundaries. return nil, fmt.Errorf("xpg/shard/resolver: range %d must satisfy start < end", index) } @@ -79,16 +79,10 @@ func NewRange[K cmp.Ordered](topology *shard.Topology, ranges []Range[K]) (*Rang } } - // sourceIndex provides deterministic ordering for equal starts and keeps - // overlap errors tied to the caller's original slice. - slices.SortFunc( + slices.SortStableFunc( entries, func(left, right rangeEntry[K]) int { - if order := cmp.Compare(left.start, right.start); order != 0 { - return order - } - - return cmp.Compare(left.sourceIndex, right.sourceIndex) + return cmp.Compare(left.start, right.start) }, ) @@ -123,7 +117,7 @@ func NewRange[K cmp.Ordered](topology *shard.Topology, ranges []Range[K]) (*Rang // or execute a PostgreSQL query. func (resolver *RangeResolver[K]) Resolve(key K) (shard.Shard, error) { if resolver == nil || len(resolver.ranges) == 0 { - return shard.Shard{}, shard.ErrNoShard + return shard.Shard{}, errors.New("xpg/shard/resolver: range resolver is not initialized") } // Non-overlap validation guarantees strictly increasing upper boundaries, diff --git a/shard/resolver/time_range.go b/shard/resolver/time_range.go index 79a54a2..4e28a80 100644 --- a/shard/resolver/time_range.go +++ b/shard/resolver/time_range.go @@ -1,7 +1,6 @@ package resolver import ( - "cmp" "errors" "fmt" "slices" @@ -56,8 +55,8 @@ func NewTimeRange(topology *shard.Topology, ranges []TimeRange) (*TimeRangeResol return nil, fmt.Errorf("xpg/shard/resolver: time range %d: %w", index, err) } - start := normalizeTime(valueRange.Start) - end := normalizeTime(valueRange.End) + start := timeToUTC(valueRange.Start) + end := timeToUTC(valueRange.End) if !start.Before(end) { return nil, fmt.Errorf("xpg/shard/resolver: time range %d must satisfy start < end", index) @@ -82,16 +81,10 @@ func NewTimeRange(topology *shard.Topology, ranges []TimeRange) (*TimeRangeResol } } - // sourceIndex keeps overlap diagnostics tied to the caller's original - // slice and provides deterministic ordering for equal starts. - slices.SortFunc( + slices.SortStableFunc( entries, func(left, right timeRangeEntry) int { - if order := left.start.Compare(right.start); order != 0 { - return order - } - - return cmp.Compare(left.sourceIndex, right.sourceIndex) + return left.start.Compare(right.start) }, ) @@ -124,10 +117,10 @@ func NewTimeRange(topology *shard.Topology, ranges []TimeRange) (*TimeRangeResol // or execute a PostgreSQL query. func (resolver *TimeRangeResolver) Resolve(key time.Time) (shard.Shard, error) { if resolver == nil || len(resolver.ranges) == 0 { - return shard.Shard{}, shard.ErrNoShard + return shard.Shard{}, errors.New("xpg/shard/resolver: time range resolver is not initialized") } - key = normalizeTime(key) + key = timeToUTC(key) // Non-overlap validation guarantees strictly increasing upper boundaries, // making this search predicate monotonic. @@ -153,6 +146,6 @@ func (resolver *TimeRangeResolver) Resolve(key time.Time) (shard.Shard, error) { return entry.shard, nil } -func normalizeTime(value time.Time) time.Time { - return value.UTC() +func timeToUTC(t time.Time) time.Time { + return t.UTC() } diff --git a/shard/topology.go b/shard/topology.go index d79ec20..4c1647f 100644 --- a/shard/topology.go +++ b/shard/topology.go @@ -36,23 +36,27 @@ func NewTopology(configs []Config) (*Topology, error) { return nil, errors.New("xpg/shard: topology must contain at least one shard") } - shards := make([]Shard, 0, len(configs)) + shards := make([]Shard, len(configs)) indexByID := make(map[ID]int, len(configs)) for index, config := range configs { - resolved, err := newShard(config) - if err != nil { - return nil, fmt.Errorf("xpg/shard: shard %d: %w", index, err) + if config.Cluster == nil { + return nil, fmt.Errorf("xpg/shard: shard %d: cluster is nil", index) } - id := resolved.ID() + id := config.Cluster.ID() + if id == "" { + return nil, fmt.Errorf("xpg/shard: shard %d: cluster ID must not be empty", index) + } if _, exists := indexByID[id]; exists { return nil, fmt.Errorf("xpg/shard: duplicate shard ID %q", id) } - indexByID[id] = len(shards) - shards = append(shards, resolved) + shards[index] = Shard{ + cluster: config.Cluster, + } + indexByID[id] = index } return &Topology{ @@ -87,10 +91,6 @@ func (t *Topology) Shards() []Shard { // Shard returns one shard by stable ID. func (t *Topology) Shard(id ID) (Shard, bool) { - return t.lookup(id) -} - -func (t *Topology) lookup(id ID) (Shard, bool) { if t == nil { return Shard{}, false } @@ -111,22 +111,8 @@ func (t *Topology) Close() { } t.closeOnce.Do(func() { - for index := len(t.shards) - 1; index >= 0; index-- { - t.shards[index].cluster.Close() + for _, v := range slices.Backward(t.shards) { + v.cluster.Close() } }) } - -func newShard(config Config) (Shard, error) { - if config.Cluster == nil { - return Shard{}, errors.New("cluster is nil") - } - - if config.Cluster.ID() == "" { - return Shard{}, errors.New("cluster ID must not be empty") - } - - return Shard{ - cluster: config.Cluster, - }, nil -} From a2917c675c95922c23cf1be6186208773a1544ec Mon Sep 17 00:00:00 2001 From: mkbeh Date: Tue, 25 Aug 2026 01:14:26 +0300 Subject: [PATCH 24/41] test: cover shard topology and resolvers --- shard/errors_test.go | 38 ++++ shard/foreach_test.go | 347 ++++++++++++++++++++++++++++++ shard/group_test.go | 248 +++++++++++++++++++++ shard/helpers_test.go | 74 +++++++ shard/resolver/custom_test.go | 147 +++++++++++++ shard/resolver/encoder_test.go | 146 +++++++++++++ shard/resolver/hash_test.go | 212 ++++++++++++++++++ shard/resolver/helpers_test.go | 60 ++++++ shard/resolver/range_test.go | 226 +++++++++++++++++++ shard/resolver/time_range_test.go | 216 +++++++++++++++++++ shard/shard_test.go | 98 +++++++++ shard/topology_test.go | 162 ++++++++++++++ 12 files changed, 1974 insertions(+) create mode 100644 shard/errors_test.go create mode 100644 shard/foreach_test.go create mode 100644 shard/group_test.go create mode 100644 shard/helpers_test.go create mode 100644 shard/resolver/custom_test.go create mode 100644 shard/resolver/encoder_test.go create mode 100644 shard/resolver/hash_test.go create mode 100644 shard/resolver/helpers_test.go create mode 100644 shard/resolver/range_test.go create mode 100644 shard/resolver/time_range_test.go create mode 100644 shard/shard_test.go create mode 100644 shard/topology_test.go diff --git a/shard/errors_test.go b/shard/errors_test.go new file mode 100644 index 0000000..5523408 --- /dev/null +++ b/shard/errors_test.go @@ -0,0 +1,38 @@ +package shard + +import ( + "errors" + "testing" +) + +func TestUnknownShardError(t *testing.T) { + t.Parallel() + + err := &UnknownShardError{ShardID: "missing"} + + if got, want := err.Error(), `xpg/shard: unknown shard "missing"`; got != want { + t.Fatalf("Error() = %q, want %q", got, want) + } + + if !errors.Is(err, ErrUnknownShard) { + t.Fatal("errors.Is() = false, want ErrUnknownShard") + } +} + +func TestMismatchError(t *testing.T) { + t.Parallel() + + err := &MismatchError{ + Expected: "shard-a", + Actual: "shard-b", + Index: 2, + } + + if got, want := err.Error(), `xpg/shard: key 2 resolved to shard "shard-b" instead of "shard-a"`; got != want { + t.Fatalf("Error() = %q, want %q", got, want) + } + + if !errors.Is(err, ErrShardMismatch) { + t.Fatal("errors.Is() = false, want ErrShardMismatch") + } +} diff --git a/shard/foreach_test.go b/shard/foreach_test.go new file mode 100644 index 0000000..2c304e7 --- /dev/null +++ b/shard/foreach_test.go @@ -0,0 +1,347 @@ +package shard + +import ( + "context" + "errors" + "strings" + "sync/atomic" + "testing" + "time" +) + +func TestForEachShardValidatesArguments(t *testing.T) { + t.Parallel() + + topology := newTestTopology(t, "shard-a") + + tests := []struct { + name string + topology *Topology + concurrency int + fn func(context.Context, Shard) error + wantError string + }{ + { + name: "nil topology", + topology: nil, + concurrency: 1, + fn: func(context.Context, Shard) error { return nil }, + wantError: "topology is nil", + }, + { + name: "zero concurrency", + topology: topology, + concurrency: 0, + fn: func(context.Context, Shard) error { return nil }, + wantError: "concurrency must be positive", + }, + { + name: "nil callback", + topology: topology, + concurrency: 1, + fn: nil, + wantError: "callback is nil", + }, + } + + for _, test := range tests { + t.Run(test.name, func(t *testing.T) { + _, err := test.topology.ForEachShard( + context.Background(), + test.concurrency, + test.fn, + ) + if err == nil { + t.Fatal("expected error") + } + + if !strings.Contains(err.Error(), test.wantError) { + t.Fatalf("error = %q, want substring %q", err, test.wantError) + } + }) + } +} + +func TestForEachShardPreservesRegistrationOrder(t *testing.T) { + t.Parallel() + + topology := newTestTopology(t, "shard-c", "shard-a", "shard-b") + + results, err := topology.ForEachShard( + context.Background(), + 2, + func(context.Context, Shard) error { return nil }, + ) + if err != nil { + t.Fatalf("ForEachShard() error = %v", err) + } + + want := []ID{"shard-c", "shard-a", "shard-b"} + for index, wantID := range want { + if got := results[index].ShardID; got != wantID { + t.Fatalf("results[%d].ShardID = %q, want %q", index, got, wantID) + } + + if results[index].Err != nil { + t.Fatalf("results[%d].Err = %v", index, results[index].Err) + } + } +} + +func TestForEachShardHonorsConcurrencyLimit(t *testing.T) { + t.Parallel() + + topology := newTestTopology( + t, + "shard-a", + "shard-b", + "shard-c", + "shard-d", + "shard-e", + "shard-f", + ) + + ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second) + defer cancel() + + started := make(chan struct{}, topology.Len()) + release := make(chan struct{}) + + var active atomic.Int32 + var maximum atomic.Int32 + var calls atomic.Int32 + + done := make(chan struct { + results ForEachShardResults + err error + }, 1) + + go func() { + results, err := topology.ForEachShard( + ctx, + 2, + func(ctx context.Context, _ Shard) error { + current := active.Add(1) + defer active.Add(-1) + + calls.Add(1) + + for { + observed := maximum.Load() + if current <= observed || maximum.CompareAndSwap(observed, current) { + break + } + } + + started <- struct{}{} + + select { + case <-release: + return nil + case <-ctx.Done(): + return ctx.Err() + } + }, + ) + + done <- struct { + results ForEachShardResults + err error + }{ + results: results, + err: err, + } + }() + + for range 2 { + select { + case <-started: + case <-ctx.Done(): + close(release) + t.Fatal("two callbacks did not start concurrently") + } + } + + close(release) + + var outcome struct { + results ForEachShardResults + err error + } + + select { + case outcome = <-done: + case <-ctx.Done(): + t.Fatal("ForEachShard() did not finish") + } + + if outcome.err != nil { + t.Fatalf("ForEachShard() error = %v", outcome.err) + } + + if got, want := calls.Load(), int32(topology.Len()); got != want { + t.Fatalf("callback calls = %d, want %d", got, want) + } + + if got, want := maximum.Load(), int32(2); got != want { + t.Fatalf("maximum concurrent callbacks = %d, want %d", got, want) + } + + if err := outcome.results.Err(); err != nil { + t.Fatalf("results.Err() = %v", err) + } +} + +func TestForEachShardCallbackErrorsDoNotStopOtherShards(t *testing.T) { + t.Parallel() + + topology := newTestTopology(t, "shard-a", "shard-b", "shard-c") + sentinel := errors.New("callback failed") + var calls atomic.Int32 + + results, err := topology.ForEachShard( + context.Background(), + 2, + func(_ context.Context, current Shard) error { + calls.Add(1) + if current.ID() == "shard-b" { + return sentinel + } + + return nil + }, + ) + if err != nil { + t.Fatalf("ForEachShard() error = %v", err) + } + + if got, want := calls.Load(), int32(3); got != want { + t.Fatalf("callback calls = %d, want %d", got, want) + } + + if results[0].Err != nil || !errors.Is(results[1].Err, sentinel) || results[2].Err != nil { + t.Fatalf("results = %+v", results) + } + + joined := results.Err() + if !errors.Is(joined, sentinel) { + t.Fatalf("results.Err() = %v, want wrapped sentinel", joined) + } + + if !strings.Contains(joined.Error(), `shard "shard-b" callback`) { + t.Fatalf("results.Err() = %q, want shard context", joined) + } +} + +func TestForEachShardCanceledBeforeScheduling(t *testing.T) { + t.Parallel() + + topology := newTestTopology(t, "shard-a", "shard-b", "shard-c") + ctx, cancel := context.WithCancel(context.Background()) + cancel() + + var calls atomic.Int32 + + results, err := topology.ForEachShard( + ctx, + 2, + func(context.Context, Shard) error { + calls.Add(1) + return nil + }, + ) + if err != nil { + t.Fatalf("ForEachShard() error = %v", err) + } + + if got := calls.Load(); got != 0 { + t.Fatalf("callback calls = %d, want 0", got) + } + + for index, result := range results { + if !errors.Is(result.Err, context.Canceled) { + t.Fatalf("results[%d].Err = %v, want context.Canceled", index, result.Err) + } + } +} + +func TestForEachShardCancellationSkipsCallbacksNotStarted(t *testing.T) { + t.Parallel() + + topology := newTestTopology(t, "shard-a", "shard-b", "shard-c", "shard-d") + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + + started := make(chan struct{}) + var calls atomic.Int32 + + done := make(chan struct { + results ForEachShardResults + err error + }, 1) + + go func() { + results, err := topology.ForEachShard( + ctx, + 1, + func(ctx context.Context, _ Shard) error { + if calls.Add(1) == 1 { + close(started) + } + + <-ctx.Done() + return ctx.Err() + }, + ) + + done <- struct { + results ForEachShardResults + err error + }{results: results, err: err} + }() + + <-started + cancel() + + outcome := <-done + if outcome.err != nil { + t.Fatalf("ForEachShard() error = %v", outcome.err) + } + + if got := calls.Load(); got != 1 { + t.Fatalf("callback calls = %d, want 1", got) + } + + for index, result := range outcome.results { + if !errors.Is(result.Err, context.Canceled) { + t.Fatalf("results[%d].Err = %v, want context.Canceled", index, result.Err) + } + } +} + +func TestForEachShardResultsErr(t *testing.T) { + t.Parallel() + + first := errors.New("first") + second := errors.New("second") + + results := ForEachShardResults{ + {ShardID: "shard-a", Err: first}, + {ShardID: "shard-b"}, + {ShardID: "shard-c", Err: second}, + } + + err := results.Err() + if !errors.Is(err, first) || !errors.Is(err, second) { + t.Fatalf("Err() = %v, want both failures", err) + } + + if got := err.Error(); !strings.Contains(got, `shard "shard-a" callback`) || + !strings.Contains(got, `shard "shard-c" callback`) { + t.Fatalf("Err() = %q, want shard context", got) + } + + if err := (ForEachShardResults{{ShardID: "shard-a"}}).Err(); err != nil { + t.Fatalf("Err() = %v, want nil", err) + } +} diff --git a/shard/group_test.go b/shard/group_test.go new file mode 100644 index 0000000..9a65457 --- /dev/null +++ b/shard/group_test.go @@ -0,0 +1,248 @@ +package shard + +import ( + "errors" + "slices" + "testing" +) + +func TestSameShard(t *testing.T) { + t.Parallel() + + topology := newTestTopology(t, "shard-a", "shard-b") + shardA := topology.At(0) + shardB := topology.At(1) + + resolver := testResolverFunc[int](func(key int) (Shard, error) { + if key < 100 { + return shardA, nil + } + + return shardB, nil + }) + + resolved, err := SameShard(resolver, 1, 2, 3) + if err != nil { + t.Fatalf("SameShard() error = %v", err) + } + + if got, want := resolved.ID(), ID("shard-a"); got != want { + t.Fatalf("SameShard().ID() = %q, want %q", got, want) + } +} + +func TestSameShardRejectsNilResolver(t *testing.T) { + t.Parallel() + + _, err := SameShard[int](nil, 1) + if err == nil { + t.Fatal("expected error") + } + + if got, want := err.Error(), "xpg/shard: resolver is nil"; got != want { + t.Fatalf("error = %q, want %q", got, want) + } +} + +func TestSameShardRequiresKey(t *testing.T) { + t.Parallel() + + resolver := testResolverFunc[int](func(int) (Shard, error) { + t.Fatal("resolver should not be called") + return Shard{}, nil + }) + + _, err := SameShard(resolver) + if !errors.Is(err, ErrNoShard) { + t.Fatalf("error = %v, want ErrNoShard", err) + } +} + +func TestSameShardWrapsFirstResolveError(t *testing.T) { + t.Parallel() + + sentinel := errors.New("resolve failed") + resolver := testResolverFunc[int](func(int) (Shard, error) { + return Shard{}, sentinel + }) + + _, err := SameShard(resolver, 1) + if !errors.Is(err, sentinel) { + t.Fatalf("error = %v, want wrapped sentinel", err) + } + + if got, want := err.Error(), "xpg/shard: resolve key 0: resolve failed"; got != want { + t.Fatalf("error = %q, want %q", got, want) + } +} + +func TestSameShardWrapsResolveErrorWithIndex(t *testing.T) { + t.Parallel() + + sentinel := errors.New("resolve failed") + resolver := testResolverFunc[int](func(key int) (Shard, error) { + if key == 2 { + return Shard{}, sentinel + } + + return Shard{}, nil + }) + + _, err := SameShard(resolver, 1, 2) + if !errors.Is(err, sentinel) { + t.Fatalf("error = %v, want wrapped sentinel", err) + } + + if got, want := err.Error(), "xpg/shard: resolve key 1: resolve failed"; got != want { + t.Fatalf("error = %q, want %q", got, want) + } +} + +func TestSameShardReturnsMismatchDetails(t *testing.T) { + t.Parallel() + + topology := newTestTopology(t, "shard-a", "shard-b") + shardA := topology.At(0) + shardB := topology.At(1) + + resolver := testResolverFunc[int](func(key int) (Shard, error) { + if key == 3 { + return shardB, nil + } + + return shardA, nil + }) + + _, err := SameShard(resolver, 1, 2, 3) + if !errors.Is(err, ErrShardMismatch) { + t.Fatalf("error = %v, want ErrShardMismatch", err) + } + + var mismatch *MismatchError + if !errors.As(err, &mismatch) { + t.Fatalf("error = %T, want *MismatchError", err) + } + + if mismatch.Expected != "shard-a" || mismatch.Actual != "shard-b" || mismatch.Index != 2 { + t.Fatalf("mismatch = %+v", mismatch) + } +} + +func TestGroupByShardPreservesStableOrder(t *testing.T) { + t.Parallel() + + topology := newTestTopology(t, "shard-a", "shard-b") + shardA := topology.At(0) + shardB := topology.At(1) + + resolver := testResolverFunc[int](func(key int) (Shard, error) { + if key < 100 { + return shardA, nil + } + + return shardB, nil + }) + + groups, err := GroupByShard(resolver, []int{142, 42, 143, 43}) + if err != nil { + t.Fatalf("GroupByShard() error = %v", err) + } + + if got, want := len(groups), 2; got != want { + t.Fatalf("len(groups) = %d, want %d", got, want) + } + + if got, want := groups[0].Shard.ID(), ID("shard-b"); got != want { + t.Fatalf("groups[0].Shard.ID() = %q, want %q", got, want) + } + if got, want := groups[0].Keys, []int{142, 143}; !slices.Equal(got, want) { + t.Fatalf("groups[0].Keys = %v, want %v", got, want) + } + + if got, want := groups[1].Shard.ID(), ID("shard-a"); got != want { + t.Fatalf("groups[1].Shard.ID() = %q, want %q", got, want) + } + if got, want := groups[1].Keys, []int{42, 43}; !slices.Equal(got, want) { + t.Fatalf("groups[1].Keys = %v, want %v", got, want) + } +} + +func TestGroupByShardResolvesEachKeyOnce(t *testing.T) { + t.Parallel() + + topology := newTestTopology(t, "shard-a") + resolved := topology.At(0) + calls := 0 + + resolver := testResolverFunc[int](func(int) (Shard, error) { + calls++ + return resolved, nil + }) + + keys := []int{1, 2, 3, 4} + groups, err := GroupByShard(resolver, keys) + if err != nil { + t.Fatalf("GroupByShard() error = %v", err) + } + + if got, want := calls, len(keys); got != want { + t.Fatalf("resolve calls = %d, want %d", got, want) + } + + if len(groups) != 1 { + t.Fatalf("len(groups) = %d, want 1", len(groups)) + } +} + +func TestGroupByShardRejectsNilResolver(t *testing.T) { + t.Parallel() + + _, err := GroupByShard[int](nil, []int{1}) + if err == nil { + t.Fatal("expected error") + } + + if got, want := err.Error(), "xpg/shard: resolver is nil"; got != want { + t.Fatalf("error = %q, want %q", got, want) + } +} + +func TestGroupByShardEmptyKeys(t *testing.T) { + t.Parallel() + + resolver := testResolverFunc[int](func(int) (Shard, error) { + t.Fatal("resolver should not be called") + return Shard{}, nil + }) + + groups, err := GroupByShard(resolver, nil) + if err != nil { + t.Fatalf("GroupByShard() error = %v", err) + } + + if len(groups) != 0 { + t.Fatalf("len(groups) = %d, want 0", len(groups)) + } +} + +func TestGroupByShardWrapsResolveErrorWithIndex(t *testing.T) { + t.Parallel() + + sentinel := errors.New("resolve failed") + resolver := testResolverFunc[int](func(key int) (Shard, error) { + if key == 3 { + return Shard{}, sentinel + } + + return Shard{}, nil + }) + + _, err := GroupByShard(resolver, []int{1, 2, 3}) + if !errors.Is(err, sentinel) { + t.Fatalf("error = %v, want wrapped sentinel", err) + } + + if got, want := err.Error(), "xpg/shard: resolve key 2: resolve failed"; got != want { + t.Fatalf("error = %q, want %q", got, want) + } +} diff --git a/shard/helpers_test.go b/shard/helpers_test.go new file mode 100644 index 0000000..553a0f5 --- /dev/null +++ b/shard/helpers_test.go @@ -0,0 +1,74 @@ +package shard + +import ( + "context" + "testing" + + "github.com/jackc/pgx/v5/pgxpool" + "github.com/mkbeh/xpg" + "github.com/mkbeh/xpg/cluster" +) + +const testDatabaseURL = "postgres://postgres@127.0.0.1:1/postgres?sslmode=disable" + +func newTestCluster(t *testing.T, id ID, labels map[string]string) *cluster.Cluster { + t.Helper() + + config, err := pgxpool.ParseConfig(testDatabaseURL) + if err != nil { + t.Fatalf("pgxpool.ParseConfig() error = %v", err) + } + + config.MinConns = 0 + config.MaxConns = 1 + + pool, err := xpg.New( + context.Background(), + config, + xpg.WithName("shard."+string(id)+".primary"), + ) + if err != nil { + t.Fatalf("xpg.New() error = %v", err) + } + + shardCluster, err := cluster.New(cluster.Config{ + ID: id, + Labels: labels, + Primary: pool, + }) + if err != nil { + pool.Close() + t.Fatalf("cluster.New() error = %v", err) + } + + t.Cleanup(shardCluster.Close) + + return shardCluster +} + +func newTestTopology(t *testing.T, ids ...ID) *Topology { + t.Helper() + + configs := make([]Config, len(ids)) + + for index, id := range ids { + configs[index] = Config{ + Cluster: newTestCluster(t, id, nil), + } + } + + topology, err := NewTopology(configs) + if err != nil { + t.Fatalf("NewTopology() error = %v", err) + } + + t.Cleanup(topology.Close) + + return topology +} + +type testResolverFunc[K any] func(K) (Shard, error) + +func (resolve testResolverFunc[K]) Resolve(key K) (Shard, error) { + return resolve(key) +} diff --git a/shard/resolver/custom_test.go b/shard/resolver/custom_test.go new file mode 100644 index 0000000..a02d4b4 --- /dev/null +++ b/shard/resolver/custom_test.go @@ -0,0 +1,147 @@ +package resolver + +import ( + "errors" + "testing" + + "github.com/mkbeh/xpg/shard" +) + +func TestNewCustomValidatesArguments(t *testing.T) { + t.Parallel() + + topology := newTestTopology(t, "shard-a") + + if resolver, err := NewCustom[int](nil, func(int, *shard.Topology) (shard.ID, error) { + return "shard-a", nil + }); err == nil { + _ = resolver + t.Fatal("expected topology error") + } + + if resolver, err := NewCustom[int](topology, nil); err == nil { + _ = resolver + t.Fatal("expected resolve function error") + } +} + +func TestCustomResolverResolve(t *testing.T) { + t.Parallel() + + topology := newTestTopology(t, "shard-a", "shard-b") + + resolver, err := NewCustom( + topology, + func(key int, gotTopology *shard.Topology) (shard.ID, error) { + if gotTopology != topology { + t.Fatal("resolve function received a different topology") + } + + if key < 100 { + return "shard-a", nil + } + + return "shard-b", nil + }, + ) + if err != nil { + t.Fatalf("NewCustom() error = %v", err) + } + + resolved, err := resolver.Resolve(142) + if err != nil { + t.Fatalf("Resolve() error = %v", err) + } + + if got, want := resolved.ID(), shard.ID("shard-b"); got != want { + t.Fatalf("Resolve().ID() = %q, want %q", got, want) + } +} + +func TestCustomResolverPropagatesResolveError(t *testing.T) { + t.Parallel() + + topology := newTestTopology(t, "shard-a") + sentinel := errors.New("resolve failed") + + resolver, err := NewCustom( + topology, + func(int, *shard.Topology) (shard.ID, error) { + return "", sentinel + }, + ) + if err != nil { + t.Fatalf("NewCustom() error = %v", err) + } + + _, err = resolver.Resolve(1) + if !errors.Is(err, sentinel) { + t.Fatalf("Resolve() error = %v, want sentinel", err) + } +} + +func TestCustomResolverPropagatesErrNoShard(t *testing.T) { + t.Parallel() + + topology := newTestTopology(t, "shard-a") + + resolver, err := NewCustom( + topology, + func(int, *shard.Topology) (shard.ID, error) { + return "", shard.ErrNoShard + }, + ) + if err != nil { + t.Fatalf("NewCustom() error = %v", err) + } + + _, err = resolver.Resolve(1) + if !errors.Is(err, shard.ErrNoShard) { + t.Fatalf("Resolve() error = %v, want ErrNoShard", err) + } +} + +func TestCustomResolverRejectsUnknownShard(t *testing.T) { + t.Parallel() + + topology := newTestTopology(t, "shard-a") + + resolver, err := NewCustom( + topology, + func(int, *shard.Topology) (shard.ID, error) { + return "missing", nil + }, + ) + if err != nil { + t.Fatalf("NewCustom() error = %v", err) + } + + _, err = resolver.Resolve(1) + if !errors.Is(err, shard.ErrUnknownShard) { + t.Fatalf("Resolve() error = %v, want ErrUnknownShard", err) + } + + var unknown *shard.UnknownShardError + if !errors.As(err, &unknown) { + t.Fatalf("Resolve() error = %T, want *shard.UnknownShardError", err) + } + + if got, want := unknown.ShardID, shard.ID("missing"); got != want { + t.Fatalf("ShardID = %q, want %q", got, want) + } +} + +func TestCustomResolverUninitialized(t *testing.T) { + t.Parallel() + + var resolver *CustomResolver[int] + + _, err := resolver.Resolve(1) + if err == nil { + t.Fatal("expected error") + } + + if got, want := err.Error(), "xpg/shard/resolver: custom resolver is not initialized"; got != want { + t.Fatalf("error = %q, want %q", got, want) + } +} diff --git a/shard/resolver/encoder_test.go b/shard/resolver/encoder_test.go new file mode 100644 index 0000000..c2285ed --- /dev/null +++ b/shard/resolver/encoder_test.go @@ -0,0 +1,146 @@ +package resolver + +import ( + "bytes" + "encoding/hex" + "errors" + "testing" +) + +func TestKeyEncoderFuncNil(t *testing.T) { + t.Parallel() + + var encoder KeyEncoderFunc[int] + + _, err := encoder.Encode(1) + if err == nil { + t.Fatal("expected error") + } + + if got, want := err.Error(), "xpg/shard/resolver: key encoder function is nil"; got != want { + t.Fatalf("error = %q, want %q", got, want) + } +} + +func TestKeyEncoderFunc(t *testing.T) { + t.Parallel() + + sentinel := errors.New("encode failed") + encoder := KeyEncoderFunc[int](func(key int) ([]byte, error) { + if key < 0 { + return nil, sentinel + } + + return []byte{byte(key)}, nil + }) + + encoded, err := encoder.Encode(7) + if err != nil { + t.Fatalf("Encode() error = %v", err) + } + if got, want := encoded, []byte{7}; !bytes.Equal(got, want) { + t.Fatalf("Encode() = %v, want %v", got, want) + } + + if _, err := encoder.Encode(-1); !errors.Is(err, sentinel) { + t.Fatalf("Encode() error = %v, want sentinel", err) + } +} + +func TestStringKeyEncoder(t *testing.T) { + t.Parallel() + + encoded, err := StringKeyEncoder().Encode("a\x00b") + if err != nil { + t.Fatalf("Encode() error = %v", err) + } + + if got, want := string(encoded), "a\x00b"; got != want { + t.Fatalf("Encode() = %q, want %q", got, want) + } +} + +func TestBytesKeyEncoderReturnsDefensiveCopy(t *testing.T) { + t.Parallel() + + key := []byte{1, 2, 3} + encoded, err := BytesKeyEncoder().Encode(key) + if err != nil { + t.Fatalf("Encode() error = %v", err) + } + + encoded[0] = 9 + + if got, want := key[0], byte(1); got != want { + t.Fatalf("input key changed to %d, want %d", got, want) + } +} + +func TestIntegerKeyEncoders(t *testing.T) { + t.Parallel() + + tests := []struct { + name string + got func() ([]byte, error) + want string + }{ + { + name: "int64 positive", + got: func() ([]byte, error) { return Int64KeyEncoder().Encode(1) }, + want: "0000000000000001", + }, + { + name: "int64 negative", + got: func() ([]byte, error) { return Int64KeyEncoder().Encode(-1) }, + want: "ffffffffffffffff", + }, + { + name: "uint64", + got: func() ([]byte, error) { return Uint64KeyEncoder().Encode(0x0102030405060708) }, + want: "0102030405060708", + }, + } + + for _, test := range tests { + t.Run(test.name, func(t *testing.T) { + encoded, err := test.got() + if err != nil { + t.Fatalf("Encode() error = %v", err) + } + + if got := hex.EncodeToString(encoded); got != test.want { + t.Fatalf("Encode() = %s, want %s", got, test.want) + } + }) + } +} + +func TestFixedSizeKeyEncoders(t *testing.T) { + t.Parallel() + + var key16 [16]byte + for index := range key16 { + key16[index] = byte(index) + } + + encoded16, err := Bytes16KeyEncoder().Encode(key16) + if err != nil { + t.Fatalf("Bytes16KeyEncoder.Encode() error = %v", err) + } + if got, want := encoded16, key16[:]; !bytes.Equal(got, want) { + t.Fatalf("Bytes16KeyEncoder.Encode() = %v, want %v", got, want) + } + + var key32 [32]byte + for index := range key32 { + key32[index] = byte(31 - index) + } + + encoded32, err := Bytes32KeyEncoder().Encode(key32) + if err != nil { + t.Fatalf("Bytes32KeyEncoder.Encode() error = %v", err) + } + if got, want := encoded32, key32[:]; !bytes.Equal(got, want) { + t.Fatalf("Bytes32KeyEncoder.Encode() = %v, want %v", got, want) + } +} diff --git a/shard/resolver/hash_test.go b/shard/resolver/hash_test.go new file mode 100644 index 0000000..6fdf064 --- /dev/null +++ b/shard/resolver/hash_test.go @@ -0,0 +1,212 @@ +package resolver + +import ( + "errors" + "fmt" + "testing" + + "github.com/mkbeh/xpg/shard" +) + +func TestNewHashValidatesArguments(t *testing.T) { + t.Parallel() + + topology := newTestTopology(t, "shard-a") + + if resolver, err := NewHash[string](nil, "users", StringKeyEncoder()); err == nil { + _ = resolver + t.Fatal("expected topology error") + } + + if resolver, err := NewHash[string](topology, "users", nil); err == nil { + _ = resolver + t.Fatal("expected encoder error") + } + + if resolver, err := NewHash(topology, "", StringKeyEncoder()); err == nil { + _ = resolver + t.Fatal("expected namespace error") + } +} + +func TestNewHashTreatsNamespaceAsOpaqueNonEmptyString(t *testing.T) { + t.Parallel() + + topology := newTestTopology(t, "shard-a") + + for _, namespace := range []string{"users", " users ", " "} { + resolver, err := NewHash(topology, namespace, StringKeyEncoder()) + if err != nil { + t.Fatalf("NewHash(%q) error = %v", namespace, err) + } + + resolved, err := resolver.Resolve("alice") + if err != nil { + t.Fatalf("Resolve() error = %v", err) + } + + if got, want := resolved.ID(), shard.ID("shard-a"); got != want { + t.Fatalf("Resolve().ID() = %q, want %q", got, want) + } + } +} + +func TestHashResolverStablePlacementVectors(t *testing.T) { + t.Parallel() + + topology := newTestTopology(t, "shard-a", "shard-b", "shard-c") + resolver, err := NewHash(topology, "users", StringKeyEncoder()) + if err != nil { + t.Fatalf("NewHash() error = %v", err) + } + + tests := []struct { + key string + want shard.ID + }{ + {key: "alice", want: "shard-a"}, + {key: "bob", want: "shard-b"}, + {key: "carol", want: "shard-b"}, + {key: "dave", want: "shard-b"}, + {key: "eve", want: "shard-c"}, + {key: "0", want: "shard-a"}, + {key: "1", want: "shard-b"}, + {key: "2", want: "shard-c"}, + } + + for _, test := range tests { + t.Run(test.key, func(t *testing.T) { + resolved, err := resolver.Resolve(test.key) + if err != nil { + t.Fatalf("Resolve() error = %v", err) + } + + if got := resolved.ID(); got != test.want { + t.Fatalf("Resolve(%q).ID() = %q, want %q", test.key, got, test.want) + } + }) + } +} + +func TestHashResolverPlacementDoesNotDependOnTopologyOrder(t *testing.T) { + t.Parallel() + + first := newTestTopology(t, "shard-a", "shard-b", "shard-c") + second := newTestTopology(t, "shard-c", "shard-a", "shard-b") + + firstResolver, err := NewHash(first, "users", StringKeyEncoder()) + if err != nil { + t.Fatalf("NewHash(first) error = %v", err) + } + secondResolver, err := NewHash(second, "users", StringKeyEncoder()) + if err != nil { + t.Fatalf("NewHash(second) error = %v", err) + } + + for _, key := range []string{"alice", "bob", "carol", "dave", "eve", "user-123"} { + firstShard, err := firstResolver.Resolve(key) + if err != nil { + t.Fatalf("first Resolve(%q) error = %v", key, err) + } + secondShard, err := secondResolver.Resolve(key) + if err != nil { + t.Fatalf("second Resolve(%q) error = %v", key, err) + } + + if firstShard.ID() != secondShard.ID() { + t.Fatalf( + "Resolve(%q) = %q and %q for different topology orders", + key, + firstShard.ID(), + secondShard.ID(), + ) + } + } +} + +func TestHashResolverAddingShardOnlyMovesKeysToNewShard(t *testing.T) { + t.Parallel() + + before := newTestTopology(t, "shard-a", "shard-b") + after := newTestTopology(t, "shard-a", "shard-b", "shard-c") + + beforeResolver, err := NewHash(before, "users", StringKeyEncoder()) + if err != nil { + t.Fatalf("NewHash(before) error = %v", err) + } + afterResolver, err := NewHash(after, "users", StringKeyEncoder()) + if err != nil { + t.Fatalf("NewHash(after) error = %v", err) + } + + moved := 0 + + for index := range 256 { + key := fmt.Sprintf("user-%d", index) + + previous, err := beforeResolver.Resolve(key) + if err != nil { + t.Fatalf("before Resolve(%q) error = %v", key, err) + } + current, err := afterResolver.Resolve(key) + if err != nil { + t.Fatalf("after Resolve(%q) error = %v", key, err) + } + + if previous.ID() == current.ID() { + continue + } + + moved++ + if current.ID() != "shard-c" { + t.Fatalf( + "Resolve(%q) moved from %q to existing shard %q", + key, + previous.ID(), + current.ID(), + ) + } + } + + if moved == 0 { + t.Fatal("expected at least one key to move to the new shard") + } +} + +func TestHashResolverWrapsEncoderError(t *testing.T) { + t.Parallel() + + topology := newTestTopology(t, "shard-a") + sentinel := errors.New("encode failed") + + resolver, err := NewHash( + topology, + "users", + KeyEncoderFunc[string](func(string) ([]byte, error) { + return nil, sentinel + }), + ) + if err != nil { + t.Fatalf("NewHash() error = %v", err) + } + + _, err = resolver.Resolve("alice") + if !errors.Is(err, sentinel) { + t.Fatalf("Resolve() error = %v, want wrapped sentinel", err) + } +} + +func TestHashResolverUninitialized(t *testing.T) { + t.Parallel() + + var resolver *HashResolver[string] + + _, err := resolver.Resolve("alice") + if err == nil { + t.Fatal("expected error") + } + + if got, want := err.Error(), "xpg/shard/resolver: hash resolver is not initialized"; got != want { + t.Fatalf("error = %q, want %q", got, want) + } +} diff --git a/shard/resolver/helpers_test.go b/shard/resolver/helpers_test.go new file mode 100644 index 0000000..cfeabf3 --- /dev/null +++ b/shard/resolver/helpers_test.go @@ -0,0 +1,60 @@ +package resolver + +import ( + "context" + "testing" + + "github.com/jackc/pgx/v5/pgxpool" + "github.com/mkbeh/xpg" + "github.com/mkbeh/xpg/cluster" + "github.com/mkbeh/xpg/shard" +) + +const testDatabaseURL = "postgres://postgres@127.0.0.1:1/postgres?sslmode=disable" + +func newTestTopology(t *testing.T, ids ...shard.ID) *shard.Topology { + t.Helper() + + configs := make([]shard.Config, len(ids)) + + for index, id := range ids { + poolConfig, err := pgxpool.ParseConfig(testDatabaseURL) + if err != nil { + t.Fatalf("pgxpool.ParseConfig() error = %v", err) + } + + poolConfig.MinConns = 0 + poolConfig.MaxConns = 1 + + pool, err := xpg.New( + context.Background(), + poolConfig, + xpg.WithName("shard."+string(id)+".primary"), + ) + if err != nil { + t.Fatalf("xpg.New() error = %v", err) + } + + shardCluster, err := cluster.New(cluster.Config{ + ID: id, + Primary: pool, + }) + if err != nil { + pool.Close() + t.Fatalf("cluster.New() error = %v", err) + } + + t.Cleanup(shardCluster.Close) + + configs[index] = shard.Config{Cluster: shardCluster} + } + + topology, err := shard.NewTopology(configs) + if err != nil { + t.Fatalf("shard.NewTopology() error = %v", err) + } + + t.Cleanup(topology.Close) + + return topology +} diff --git a/shard/resolver/range_test.go b/shard/resolver/range_test.go new file mode 100644 index 0000000..c4b4ae2 --- /dev/null +++ b/shard/resolver/range_test.go @@ -0,0 +1,226 @@ +package resolver + +import ( + "errors" + "math" + "slices" + "strings" + "testing" + + "github.com/mkbeh/xpg/shard" +) + +func TestNewRangeValidatesArguments(t *testing.T) { + t.Parallel() + + topology := newTestTopology(t, "shard-a") + + if resolver, err := NewRange[int](nil, []Range[int]{{Start: 0, End: 10, ShardID: "shard-a"}}); err == nil { + _ = resolver + t.Fatal("expected topology error") + } + + if resolver, err := NewRange[int](topology, nil); err == nil { + _ = resolver + t.Fatal("expected empty ranges error") + } + + tests := []struct { + name string + valueRange Range[int] + want string + }{ + { + name: "empty shard ID", + valueRange: Range[int]{Start: 0, End: 10}, + want: "shard ID must not be empty", + }, + { + name: "empty interval", + valueRange: Range[int]{Start: 10, End: 10, ShardID: "shard-a"}, + want: "must satisfy start < end", + }, + { + name: "reversed interval", + valueRange: Range[int]{Start: 20, End: 10, ShardID: "shard-a"}, + want: "must satisfy start < end", + }, + { + name: "unknown shard", + valueRange: Range[int]{Start: 0, End: 10, ShardID: "missing"}, + want: `unknown shard "missing"`, + }, + } + + for _, test := range tests { + t.Run(test.name, func(t *testing.T) { + resolver, err := NewRange(topology, []Range[int]{test.valueRange}) + if err == nil { + _ = resolver + t.Fatal("expected error") + } + + if !strings.Contains(err.Error(), test.want) { + t.Fatalf("error = %q, want substring %q", err, test.want) + } + }) + } +} + +func TestNewRangeRejectsNaNBoundaries(t *testing.T) { + t.Parallel() + + topology := newTestTopology(t, "shard-a") + + tests := []Range[float64]{ + {Start: math.NaN(), End: 10, ShardID: "shard-a"}, + {Start: 0, End: math.NaN(), ShardID: "shard-a"}, + } + + for index, valueRange := range tests { + resolver, err := NewRange(topology, []Range[float64]{valueRange}) + if err == nil { + _ = resolver + t.Fatalf("range %d: expected error", index) + } + } +} + +func TestNewRangeRejectsOverlapUsingSourceIndexes(t *testing.T) { + t.Parallel() + + topology := newTestTopology(t, "shard-a", "shard-b") + + resolver, err := NewRange(topology, []Range[int]{ + {Start: 100, End: 200, ShardID: "shard-b"}, + {Start: 50, End: 150, ShardID: "shard-a"}, + }) + if err == nil { + _ = resolver + t.Fatal("expected overlap error") + } + + if got, want := err.Error(), "xpg/shard/resolver: ranges 1 and 0 overlap"; got != want { + t.Fatalf("error = %q, want %q", got, want) + } +} + +func TestNewRangeDoesNotModifyInputOrder(t *testing.T) { + t.Parallel() + + topology := newTestTopology(t, "shard-a", "shard-b") + ranges := []Range[int]{ + {Start: 100, End: 200, ShardID: "shard-b"}, + {Start: 0, End: 100, ShardID: "shard-a"}, + } + want := append([]Range[int](nil), ranges...) + + if _, err := NewRange(topology, ranges); err != nil { + t.Fatalf("NewRange() error = %v", err) + } + + if !slices.Equal(ranges, want) { + t.Fatalf("ranges = %+v, want unchanged %+v", ranges, want) + } +} + +func TestRangeResolverHalfOpenBoundariesAndGaps(t *testing.T) { + t.Parallel() + + topology := newTestTopology(t, "shard-a", "shard-b") + resolver, err := NewRange(topology, []Range[int]{ + {Start: 20, End: 30, ShardID: "shard-b"}, + {Start: 0, End: 10, ShardID: "shard-a"}, + }) + if err != nil { + t.Fatalf("NewRange() error = %v", err) + } + + tests := []struct { + key int + wantID shard.ID + wantErr error + }{ + {key: 0, wantID: "shard-a"}, + {key: 9, wantID: "shard-a"}, + {key: 10, wantErr: shard.ErrNoShard}, + {key: 19, wantErr: shard.ErrNoShard}, + {key: 20, wantID: "shard-b"}, + {key: 29, wantID: "shard-b"}, + {key: 30, wantErr: shard.ErrNoShard}, + } + + for _, test := range tests { + resolved, err := resolver.Resolve(test.key) + if test.wantErr != nil { + if !errors.Is(err, test.wantErr) { + t.Fatalf("Resolve(%d) error = %v, want %v", test.key, err, test.wantErr) + } + continue + } + + if err != nil { + t.Fatalf("Resolve(%d) error = %v", test.key, err) + } + if got := resolved.ID(); got != test.wantID { + t.Fatalf("Resolve(%d).ID() = %q, want %q", test.key, got, test.wantID) + } + } +} + +func TestRangeResolverAllowsAdjacentRanges(t *testing.T) { + t.Parallel() + + topology := newTestTopology(t, "shard-a", "shard-b") + resolver, err := NewRange(topology, []Range[int]{ + {Start: 0, End: 100, ShardID: "shard-a"}, + {Start: 100, End: 200, ShardID: "shard-b"}, + }) + if err != nil { + t.Fatalf("NewRange() error = %v", err) + } + + resolved, err := resolver.Resolve(100) + if err != nil { + t.Fatalf("Resolve() error = %v", err) + } + if got, want := resolved.ID(), shard.ID("shard-b"); got != want { + t.Fatalf("Resolve().ID() = %q, want %q", got, want) + } +} + +func TestRangeResolverSupportsStrings(t *testing.T) { + t.Parallel() + + topology := newTestTopology(t, "shard-a", "shard-b") + resolver, err := NewRange(topology, []Range[string]{ + {Start: "a", End: "m", ShardID: "shard-a"}, + {Start: "m", End: "z", ShardID: "shard-b"}, + }) + if err != nil { + t.Fatalf("NewRange() error = %v", err) + } + + resolved, err := resolver.Resolve("m") + if err != nil { + t.Fatalf("Resolve() error = %v", err) + } + if got, want := resolved.ID(), shard.ID("shard-b"); got != want { + t.Fatalf("Resolve().ID() = %q, want %q", got, want) + } +} + +func TestRangeResolverUninitialized(t *testing.T) { + t.Parallel() + + var resolver *RangeResolver[int] + + _, err := resolver.Resolve(1) + if err == nil { + t.Fatal("expected error") + } + + if got, want := err.Error(), "xpg/shard/resolver: range resolver is not initialized"; got != want { + t.Fatalf("error = %q, want %q", got, want) + } +} diff --git a/shard/resolver/time_range_test.go b/shard/resolver/time_range_test.go new file mode 100644 index 0000000..6611611 --- /dev/null +++ b/shard/resolver/time_range_test.go @@ -0,0 +1,216 @@ +package resolver + +import ( + "errors" + "slices" + "strings" + "testing" + "time" + + "github.com/mkbeh/xpg/shard" +) + +func TestNewTimeRangeValidatesArguments(t *testing.T) { + t.Parallel() + + topology := newTestTopology(t, "shard-a") + start := time.Date(2026, 1, 1, 0, 0, 0, 0, time.UTC) + end := start.Add(time.Hour) + + if resolver, err := NewTimeRange(nil, []TimeRange{{Start: start, End: end, ShardID: "shard-a"}}); err == nil { + _ = resolver + t.Fatal("expected topology error") + } + + if resolver, err := NewTimeRange(topology, nil); err == nil { + _ = resolver + t.Fatal("expected empty ranges error") + } + + tests := []struct { + name string + valueRange TimeRange + want string + }{ + { + name: "empty shard ID", + valueRange: TimeRange{Start: start, End: end}, + want: "shard ID must not be empty", + }, + { + name: "empty interval", + valueRange: TimeRange{Start: start, End: start, ShardID: "shard-a"}, + want: "must satisfy start < end", + }, + { + name: "reversed interval", + valueRange: TimeRange{Start: end, End: start, ShardID: "shard-a"}, + want: "must satisfy start < end", + }, + { + name: "unknown shard", + valueRange: TimeRange{Start: start, End: end, ShardID: "missing"}, + want: `unknown shard "missing"`, + }, + } + + for _, test := range tests { + t.Run(test.name, func(t *testing.T) { + resolver, err := NewTimeRange(topology, []TimeRange{test.valueRange}) + if err == nil { + _ = resolver + t.Fatal("expected error") + } + + if !strings.Contains(err.Error(), test.want) { + t.Fatalf("error = %q, want substring %q", err, test.want) + } + }) + } +} + +func TestNewTimeRangeRejectsOverlapUsingSourceIndexes(t *testing.T) { + t.Parallel() + + topology := newTestTopology(t, "shard-a", "shard-b") + base := time.Date(2026, 1, 1, 0, 0, 0, 0, time.UTC) + + resolver, err := NewTimeRange(topology, []TimeRange{ + {Start: base.Add(2 * time.Hour), End: base.Add(4 * time.Hour), ShardID: "shard-b"}, + {Start: base.Add(time.Hour), End: base.Add(3 * time.Hour), ShardID: "shard-a"}, + }) + if err == nil { + _ = resolver + t.Fatal("expected overlap error") + } + + if got, want := err.Error(), "xpg/shard/resolver: time ranges 1 and 0 overlap"; got != want { + t.Fatalf("error = %q, want %q", got, want) + } +} + +func TestNewTimeRangeDoesNotModifyInput(t *testing.T) { + t.Parallel() + + topology := newTestTopology(t, "shard-a", "shard-b") + location := time.FixedZone("UTC+3", 3*60*60) + base := time.Date(2026, 1, 1, 0, 0, 0, 0, location) + ranges := []TimeRange{ + {Start: base.Add(time.Hour), End: base.Add(2 * time.Hour), ShardID: "shard-b"}, + {Start: base, End: base.Add(time.Hour), ShardID: "shard-a"}, + } + want := append([]TimeRange(nil), ranges...) + + if _, err := NewTimeRange(topology, ranges); err != nil { + t.Fatalf("NewTimeRange() error = %v", err) + } + + if !slices.Equal(ranges, want) { + t.Fatalf("ranges = %+v, want unchanged %+v", ranges, want) + } +} + +func TestTimeRangeResolverHalfOpenBoundariesAndGaps(t *testing.T) { + t.Parallel() + + topology := newTestTopology(t, "shard-a", "shard-b") + base := time.Date(2026, 1, 1, 0, 0, 0, 0, time.UTC) + resolver, err := NewTimeRange(topology, []TimeRange{ + {Start: base.Add(2 * time.Hour), End: base.Add(3 * time.Hour), ShardID: "shard-b"}, + {Start: base, End: base.Add(time.Hour), ShardID: "shard-a"}, + }) + if err != nil { + t.Fatalf("NewTimeRange() error = %v", err) + } + + tests := []struct { + key time.Time + wantID shard.ID + wantErr error + }{ + {key: base, wantID: "shard-a"}, + {key: base.Add(time.Hour - time.Nanosecond), wantID: "shard-a"}, + {key: base.Add(time.Hour), wantErr: shard.ErrNoShard}, + {key: base.Add(2 * time.Hour), wantID: "shard-b"}, + {key: base.Add(3 * time.Hour), wantErr: shard.ErrNoShard}, + } + + for _, test := range tests { + resolved, err := resolver.Resolve(test.key) + if test.wantErr != nil { + if !errors.Is(err, test.wantErr) { + t.Fatalf("Resolve(%v) error = %v, want %v", test.key, err, test.wantErr) + } + continue + } + + if err != nil { + t.Fatalf("Resolve(%v) error = %v", test.key, err) + } + if got := resolved.ID(); got != test.wantID { + t.Fatalf("Resolve(%v).ID() = %q, want %q", test.key, got, test.wantID) + } + } +} + +func TestTimeRangeResolverNormalizesToUTC(t *testing.T) { + t.Parallel() + + topology := newTestTopology(t, "shard-a") + start := time.Date(2026, 1, 1, 0, 0, 0, 0, time.UTC) + end := start.Add(time.Hour) + resolver, err := NewTimeRange(topology, []TimeRange{ + {Start: start, End: end, ShardID: "shard-a"}, + }) + if err != nil { + t.Fatalf("NewTimeRange() error = %v", err) + } + + location := time.FixedZone("UTC+3", 3*60*60) + key := start.Add(30 * time.Minute).In(location) + + resolved, err := resolver.Resolve(key) + if err != nil { + t.Fatalf("Resolve() error = %v", err) + } + if got, want := resolved.ID(), shard.ID("shard-a"); got != want { + t.Fatalf("Resolve().ID() = %q, want %q", got, want) + } +} + +func TestTimeRangeResolverAllowsAdjacentRanges(t *testing.T) { + t.Parallel() + + topology := newTestTopology(t, "shard-a", "shard-b") + base := time.Date(2026, 1, 1, 0, 0, 0, 0, time.UTC) + resolver, err := NewTimeRange(topology, []TimeRange{ + {Start: base, End: base.Add(time.Hour), ShardID: "shard-a"}, + {Start: base.Add(time.Hour), End: base.Add(2 * time.Hour), ShardID: "shard-b"}, + }) + if err != nil { + t.Fatalf("NewTimeRange() error = %v", err) + } + + resolved, err := resolver.Resolve(base.Add(time.Hour)) + if err != nil { + t.Fatalf("Resolve() error = %v", err) + } + if got, want := resolved.ID(), shard.ID("shard-b"); got != want { + t.Fatalf("Resolve().ID() = %q, want %q", got, want) + } +} + +func TestTimeRangeResolverUninitialized(t *testing.T) { + t.Parallel() + + var resolver *TimeRangeResolver + + _, err := resolver.Resolve(time.Now()) + if err == nil { + t.Fatal("expected error") + } + + if got, want := err.Error(), "xpg/shard/resolver: time range resolver is not initialized"; got != want { + t.Fatalf("error = %q, want %q", got, want) + } +} diff --git a/shard/shard_test.go b/shard/shard_test.go new file mode 100644 index 0000000..5355873 --- /dev/null +++ b/shard/shard_test.go @@ -0,0 +1,98 @@ +package shard + +import ( + "context" + "errors" + "testing" + + "github.com/jackc/pgx/v5" + "github.com/mkbeh/xpg/cluster" +) + +func TestShardZeroValue(t *testing.T) { + t.Parallel() + + var shard Shard + + if got := shard.ID(); got != "" { + t.Fatalf("ID() = %q, want empty", got) + } + + if value, ok := shard.Label("region"); ok || value != "" { + t.Fatalf("Label() = %q, %v; want empty, false", value, ok) + } + + if labels := shard.Labels(); labels != nil { + t.Fatalf("Labels() = %#v, want nil", labels) + } + + if primary := shard.Primary(); primary != nil { + t.Fatalf("Primary() = %p, want nil", primary) + } + + if _, err := shard.ReadPool(context.Background(), cluster.ReadPrimary); !errors.Is(err, ErrNoShard) { + t.Fatalf("ReadPool() error = %v, want ErrNoShard", err) + } + + if err := shard.InPrimaryTx(context.Background(), pgx.TxOptions{}, nil); !errors.Is(err, ErrNoShard) { + t.Fatalf("InPrimaryTx() error = %v, want ErrNoShard", err) + } + + if err := shard.InReadTx( + context.Background(), + cluster.ReadPrimary, + cluster.ReadTxOptions{}, + nil, + ); !errors.Is(err, ErrNoShard) { + t.Fatalf("InReadTx() error = %v, want ErrNoShard", err) + } +} + +func TestShardDelegatesClusterMetadataAndRouting(t *testing.T) { + t.Parallel() + + shardCluster := newTestCluster(t, "shard-a", map[string]string{ + "region": "eu-west", + "role": "", + }) + + topology, err := NewTopology([]Config{{Cluster: shardCluster}}) + if err != nil { + t.Fatalf("NewTopology() error = %v", err) + } + t.Cleanup(topology.Close) + + resolved := topology.At(0) + + if got, want := resolved.ID(), ID("shard-a"); got != want { + t.Fatalf("ID() = %q, want %q", got, want) + } + + if got, ok := resolved.Label("region"); !ok || got != "eu-west" { + t.Fatalf("Label(region) = %q, %v", got, ok) + } + + if got, ok := resolved.Label("role"); !ok || got != "" { + t.Fatalf("Label(role) = %q, %v", got, ok) + } + + labels := resolved.Labels() + labels["region"] = "changed" + + if got, _ := resolved.Label("region"); got != "eu-west" { + t.Fatalf("Label(region) after mutation = %q, want eu-west", got) + } + + if resolved.Primary() != shardCluster.Primary() { + t.Fatal("Primary() did not return cluster primary") + } + + pool, err := resolved.ReadPool(context.Background(), cluster.ReadPrimary) + if err != nil { + t.Fatalf("ReadPool() error = %v", err) + } + + if pool != shardCluster.Primary() { + t.Fatal("ReadPool() did not return cluster primary") + } +} diff --git a/shard/topology_test.go b/shard/topology_test.go new file mode 100644 index 0000000..4d50262 --- /dev/null +++ b/shard/topology_test.go @@ -0,0 +1,162 @@ +package shard + +import ( + "strings" + "testing" +) + +func TestNewTopologyRequiresShard(t *testing.T) { + t.Parallel() + + topology, err := NewTopology(nil) + if err == nil { + topology.Close() + t.Fatal("expected error") + } + + if got, want := err.Error(), "xpg/shard: topology must contain at least one shard"; got != want { + t.Fatalf("error = %q, want %q", got, want) + } +} + +func TestNewTopologyRejectsNilCluster(t *testing.T) { + t.Parallel() + + topology, err := NewTopology([]Config{{}}) + if err == nil { + topology.Close() + t.Fatal("expected error") + } + + if !strings.Contains(err.Error(), "cluster is nil") { + t.Fatalf("error = %q, want cluster validation error", err) + } +} + +func TestNewTopologyRejectsEmptyClusterID(t *testing.T) { + t.Parallel() + + shardCluster := newTestCluster(t, "", nil) + topology, err := NewTopology([]Config{{Cluster: shardCluster}}) + if err == nil { + topology.Close() + t.Fatal("expected error") + } + + if !strings.Contains(err.Error(), "cluster ID must not be empty") { + t.Fatalf("error = %q, want cluster ID validation error", err) + } +} + +func TestNewTopologyRejectsDuplicateIDs(t *testing.T) { + t.Parallel() + + first := newTestCluster(t, "shard-a", nil) + second := newTestCluster(t, "shard-a", nil) + topology, err := NewTopology([]Config{ + {Cluster: first}, + {Cluster: second}, + }) + if err == nil { + topology.Close() + t.Fatal("expected error") + } + + if !strings.Contains(err.Error(), `duplicate shard ID "shard-a"`) { + t.Fatalf("error = %q, want duplicate shard ID error", err) + } +} + +func TestTopologyPreservesRegistrationOrder(t *testing.T) { + t.Parallel() + + topology := newTestTopology(t, "shard-b", "shard-a", "shard-c") + + if got, want := topology.Len(), 3; got != want { + t.Fatalf("Len() = %d, want %d", got, want) + } + + want := []ID{"shard-b", "shard-a", "shard-c"} + + for index, wantID := range want { + if got := topology.At(index).ID(); got != wantID { + t.Fatalf("At(%d).ID() = %q, want %q", index, got, wantID) + } + + resolved, ok := topology.Shard(wantID) + if !ok { + t.Fatalf("Shard(%q) not found", wantID) + } + + if got := resolved.ID(); got != wantID { + t.Fatalf("Shard(%q).ID() = %q", wantID, got) + } + } +} + +func TestTopologyShardsReturnsDefensiveCopy(t *testing.T) { + t.Parallel() + + topology := newTestTopology(t, "shard-a", "shard-b") + + shards := topology.Shards() + shards[0] = Shard{} + + if got, want := topology.At(0).ID(), ID("shard-a"); got != want { + t.Fatalf("At(0).ID() = %q, want %q", got, want) + } +} + +func TestTopologyShardUnknownID(t *testing.T) { + t.Parallel() + + topology := newTestTopology(t, "shard-a") + + resolved, ok := topology.Shard("missing") + if ok { + t.Fatalf("Shard() = %+v, true; want false", resolved) + } +} + +func TestTopologyAtPanicsOutOfRange(t *testing.T) { + t.Parallel() + + topology := newTestTopology(t, "shard-a") + + defer func() { + if recover() == nil { + t.Fatal("expected panic") + } + }() + + _ = topology.At(1) +} + +func TestTopologyNilReceiver(t *testing.T) { + t.Parallel() + + var topology *Topology + + if got := topology.Len(); got != 0 { + t.Fatalf("Len() = %d, want 0", got) + } + + if shards := topology.Shards(); shards != nil { + t.Fatalf("Shards() = %#v, want nil", shards) + } + + if resolved, ok := topology.Shard("shard-a"); ok || resolved.ID() != "" { + t.Fatalf("Shard() = %+v, %v; want zero, false", resolved, ok) + } + + topology.Close() +} + +func TestTopologyCloseIsIdempotent(t *testing.T) { + t.Parallel() + + topology := newTestTopology(t, "shard-a", "shard-b") + + topology.Close() + topology.Close() +} From e2af674e91c6584a7bf93b18508d5b4d66eba761 Mon Sep 17 00:00:00 2001 From: mkbeh Date: Tue, 25 Aug 2026 14:07:20 +0300 Subject: [PATCH 25/41] refactor: simplify basic example --- examples/basic/README.md | 78 ++++++++++++++++--------------- examples/basic/docker-compose.yml | 30 ++++++++++++ examples/basic/go.mod | 5 +- examples/basic/main.go | 68 +++++++++------------------ examples/basic/setup.sql | 14 ------ examples/basic/sql/schema.sql | 8 ++++ 6 files changed, 102 insertions(+), 101 deletions(-) create mode 100644 examples/basic/docker-compose.yml delete mode 100644 examples/basic/setup.sql create mode 100644 examples/basic/sql/schema.sql diff --git a/examples/basic/README.md b/examples/basic/README.md index adef86e..c971969 100644 --- a/examples/basic/README.md +++ b/examples/basic/README.md @@ -1,53 +1,34 @@ # Basic pool usage -This example opens an `xpg.Pool`, checks PostgreSQL connectivity, writes two rows, and reads them back with the common -query methods exposed by the pool. +This example shows how to use `xpg.Pool` for common PostgreSQL operations: -**This example demonstrates:** - -* Creating and closing an `xpg.Pool` -* Explicitly checking connectivity with `Pool.Ping` -* Executing a statement with `Pool.Exec` -* Reading one row with `Pool.QueryRow` -* Iterating over rows returned by `Pool.Query` -* Using `Pool.Name` as logical pool metadata - -## Configuration - -The example uses the following connection string by default: - -```text -postgres://postgres:postgres@localhost:5432/postgres?sslmode=disable -``` - -Set `XPG_DATABASE_URL` to use another PostgreSQL instance: - -```shell -export XPG_DATABASE_URL='postgres://user:password@localhost:5432/database?sslmode=disable' -``` +* Create a named PostgreSQL connection pool and verify connectivity +* Execute a write and read a single record back +* Query and iterate over multiple rows ## Local setup -Start PostgreSQL and Adminer from the repository root: +From this directory, start PostgreSQL and Adminer: ```shell -docker compose -f examples/docker-compose.yml --profile tools up -d +docker compose up -d ``` -Or from this example directory: +Apply the example schema: ```shell -docker compose -f ../docker-compose.yml --profile tools up -d +psql 'postgres://postgres:postgres@localhost:5432/postgres?sslmode=disable' \ + < sql/schema.sql ``` -Services are available at: +The services are available at: ```text PostgreSQL: localhost:5432 Adminer: http://localhost:8080 ``` -Sign in to Adminer with: +To inspect the example data in Adminer, sign in with: ```text System: PostgreSQL @@ -57,6 +38,20 @@ Password: postgres Database: postgres ``` +## Configuration + +By default, the example connects to: + +```text +postgres://postgres:postgres@localhost:5432/postgres?sslmode=disable +``` + +To use another PostgreSQL instance, set `XPG_DATABASE_URL`: + +```shell +export XPG_DATABASE_URL='postgres://user:password@localhost:5432/database?sslmode=disable' +``` + ## Run From this directory: @@ -75,20 +70,29 @@ go run ./examples/basic ```text pool: basic-example -inserted users: 2 +upserted users: 2 selected user: 1 Alice active=true active users: - 1 Alice -- 2 Bob ``` -The embedded `setup.sql` file recreates the `xpg_basic_example` schema before each run. The resulting data remains in -PostgreSQL so it can be inspected in Adminer. +## Cleanup + +To remove the example schema and data: + +```shell +psql 'postgres://postgres:postgres@localhost:5432/postgres?sslmode=disable' \ + -c 'DROP SCHEMA IF EXISTS xpg_basic_example CASCADE;' +``` -## Stop services +Stop the local services: ```shell -docker compose -f examples/docker-compose.yml down +docker compose down ``` -Add `-v` to remove the PostgreSQL volume as well. +To also remove the PostgreSQL data volume: + +```shell +docker compose down -v +``` diff --git a/examples/basic/docker-compose.yml b/examples/basic/docker-compose.yml new file mode 100644 index 0000000..046626a --- /dev/null +++ b/examples/basic/docker-compose.yml @@ -0,0 +1,30 @@ +services: + postgres: + image: postgres:18-alpine + environment: + POSTGRES_DB: postgres + POSTGRES_USER: postgres + POSTGRES_PASSWORD: postgres + ports: + - "5432:5432" + volumes: + - postgres-data:/var/lib/postgresql + healthcheck: + test: "pg_isready -U $${POSTGRES_USER} -d $${POSTGRES_DB}" + interval: 2s + timeout: 5s + retries: 10 + start_period: 5s + + adminer: + image: adminer:standalone + environment: + ADMINER_DEFAULT_SERVER: postgres + ports: + - "8080:8080" + depends_on: + postgres: + condition: service_healthy + +volumes: + postgres-data: \ No newline at end of file diff --git a/examples/basic/go.mod b/examples/basic/go.mod index 4a1fff5..7e5031b 100644 --- a/examples/basic/go.mod +++ b/examples/basic/go.mod @@ -2,7 +2,4 @@ module basic go 1.27 -require ( - github.com/jackc/pgx/v5 v5.10.0 - github.com/mkbeh/xpg v0.2.0 -) +require github.com/mkbeh/xpg v0.2.0 diff --git a/examples/basic/main.go b/examples/basic/main.go index fb7005f..b3c1244 100644 --- a/examples/basic/main.go +++ b/examples/basic/main.go @@ -2,20 +2,15 @@ package main import ( "context" - _ "embed" "fmt" "log" "os" - "github.com/jackc/pgx/v5" "github.com/mkbeh/xpg" ) const defaultDatabaseURL = "postgres://postgres:postgres@localhost:5432/postgres?sslmode=disable" -//go:embed setup.sql -var setupSQL string - type user struct { ID int64 Name string @@ -36,7 +31,7 @@ func run(ctx context.Context) error { xpg.WithName("basic-example"), ) if err != nil { - return fmt.Errorf("create pool: %w", err) + return fmt.Errorf("open pool: %w", err) } defer pool.Close() @@ -44,13 +39,9 @@ func run(ctx context.Context) error { return fmt.Errorf("ping PostgreSQL: %w", err) } - if err := prepareExample(ctx, pool); err != nil { - return fmt.Errorf("prepare example: %w", err) - } - - inserted, err := insertUsers(ctx, pool) + inserted, err := upsertUsers(ctx, pool) if err != nil { - return fmt.Errorf("insert users: %w", err) + return fmt.Errorf("upsert users: %w", err) } selected, err := loadUser(ctx, pool, 1) @@ -58,13 +49,13 @@ func run(ctx context.Context) error { return fmt.Errorf("load user: %w", err) } - users, err := listUsers(ctx, pool) + active, err := listActiveUsers(ctx, pool) if err != nil { - return fmt.Errorf("list users: %w", err) + return fmt.Errorf("list active users: %w", err) } fmt.Printf("pool: %s\n", pool.Name()) - fmt.Printf("inserted users: %d\n", inserted) + fmt.Printf("upserted users: %d\n", inserted) fmt.Printf( "selected user: %d %s <%s> active=%t\n", selected.ID, @@ -72,35 +63,16 @@ func run(ctx context.Context) error { selected.Email, selected.Active, ) - fmt.Println("active users:") - for _, current := range users { + fmt.Println("active users:") + for _, current := range active { fmt.Printf("- %d %s <%s>\n", current.ID, current.Name, current.Email) } return nil } -func prepareExample( - ctx context.Context, - pool *xpg.Pool, -) error { - _, err := pool.Exec( - ctx, - setupSQL, - pgx.QueryExecModeSimpleProtocol, - ) - if err != nil { - return fmt.Errorf("execute setup SQL: %w", err) - } - - return nil -} - -func insertUsers( - ctx context.Context, - pool *xpg.Pool, -) (int64, error) { +func upsertUsers(ctx context.Context, pool *xpg.Pool) (int64, error) { tag, err := pool.Exec( ctx, `INSERT INTO xpg_basic_example.users ( @@ -111,7 +83,11 @@ func insertUsers( ) VALUES ($1, $2, $3, $4), - ($5, $6, $7, $8)`, + ($5, $6, $7, $8) + ON CONFLICT (id) DO UPDATE SET + name = EXCLUDED.name, + email = EXCLUDED.email, + active = EXCLUDED.active`, int64(1), "Alice", "alice@example.com", @@ -119,7 +95,7 @@ func insertUsers( int64(2), "Bob", "bob@example.com", - true, + false, ) if err != nil { return 0, err @@ -142,8 +118,8 @@ func loadUser( name, email, active - FROM xpg_basic_example.users - WHERE id = $1`, + FROM xpg_basic_example.users + WHERE id = $1`, userID, ).Scan( &selected.ID, @@ -158,7 +134,7 @@ func loadUser( return selected, nil } -func listUsers( +func listActiveUsers( ctx context.Context, pool *xpg.Pool, ) ([]user, error) { @@ -169,16 +145,16 @@ func listUsers( name, email, active - FROM xpg_basic_example.users - WHERE active - ORDER BY id`, + FROM xpg_basic_example.users + WHERE active + ORDER BY id`, ) if err != nil { return nil, err } defer rows.Close() - users := make([]user, 0, 2) + var users []user for rows.Next() { var current user diff --git a/examples/basic/setup.sql b/examples/basic/setup.sql deleted file mode 100644 index 5be37b5..0000000 --- a/examples/basic/setup.sql +++ /dev/null @@ -1,14 +0,0 @@ -BEGIN; - -DROP SCHEMA IF EXISTS xpg_basic_example CASCADE; - -CREATE SCHEMA xpg_basic_example; - -CREATE TABLE xpg_basic_example.users ( - id bigint PRIMARY KEY, - name text NOT NULL, - email text NOT NULL UNIQUE, - active boolean NOT NULL -); - -COMMIT; diff --git a/examples/basic/sql/schema.sql b/examples/basic/sql/schema.sql new file mode 100644 index 0000000..846e4b1 --- /dev/null +++ b/examples/basic/sql/schema.sql @@ -0,0 +1,8 @@ +CREATE SCHEMA IF NOT EXISTS xpg_basic_example; + +CREATE TABLE IF NOT EXISTS xpg_basic_example.users ( + id bigint PRIMARY KEY, + name text NOT NULL, + email text NOT NULL UNIQUE, + active boolean NOT NULL +); From c2f36503c9af772c0b888e8b3f85a55dfdeac02a Mon Sep 17 00:00:00 2001 From: mkbeh Date: Tue, 25 Aug 2026 14:17:34 +0300 Subject: [PATCH 26/41] refactor: simplify transactions example --- examples/transactions/README.md | 104 ++++++++---------- .../{ => transactions}/docker-compose.yml | 4 +- examples/transactions/go.mod | 3 +- examples/transactions/main.go | 75 +++---------- examples/transactions/setup.sql | 24 ---- examples/transactions/sql/schema.sql | 16 +++ 6 files changed, 81 insertions(+), 145 deletions(-) rename examples/{ => transactions}/docker-compose.yml (93%) delete mode 100644 examples/transactions/setup.sql create mode 100644 examples/transactions/sql/schema.sql diff --git a/examples/transactions/README.md b/examples/transactions/README.md index 8b666c8..d5ad7d8 100644 --- a/examples/transactions/README.md +++ b/examples/transactions/README.md @@ -1,77 +1,66 @@ # Transactions and savepoints -This example creates an order in a transaction and applies an optional promo code inside a PostgreSQL savepoint. -The promo code is already used, so only the savepoint is rolled back while the outer transaction commits the order. +This example shows how to keep an outer transaction commit-able when an optional operation fails: -**This example demonstrates:** +* Execute the main operation inside a transaction +* Isolate optional work with a savepoint +* Detect an expected PostgreSQL constraint violation +* Roll back only the savepoint while allowing the outer transaction to commit -* Creating an `xpg.Pool` and explicitly checking PostgreSQL connectivity -* Executing a callback with an explicit `pgx.Tx` -* Isolating an optional operation with `xpg.InSavepoint` -* Inspecting a wrapped `pgconn.PgError` -* Continuing and committing the outer transaction after a savepoint rollback +The example attempts to redeem `PROMO2026`, which is already in use. The promo write fails and is rolled back to the +savepoint, while the order is committed successfully. -## Configuration - -The example connects to PostgreSQL using: +## Local setup -```text -postgres://postgres:postgres@localhost:5432/postgres?sslmode=disable -``` - -Set `XPG_DATABASE_URL` to use another connection string: +From this directory, start PostgreSQL and Adminer: ```shell -export XPG_DATABASE_URL='postgres://user:password@localhost:5432/database?sslmode=disable' +docker compose up -d ``` -## Local PostgreSQL setup - -The example can use the local PostgreSQL setup from `examples/docker-compose.yml`. - -From the repository root: +Apply the example schema: ```shell -docker compose -f examples/docker-compose.yml --profile tools up -d +psql 'postgres://postgres:postgres@localhost:5432/postgres?sslmode=disable' \ + < sql/schema.sql ``` -Or from this example directory: +The services are available at: -```shell -docker compose -f ../docker-compose.yml --profile tools up -d +```text +PostgreSQL: localhost:5432 +Adminer: http://localhost:8080 ``` -To start only PostgreSQL, omit `--profile tools`. - -PostgreSQL is available to applications running on the host at: +To inspect the example data in Adminer, sign in with: ```text -localhost:5432 +System: PostgreSQL +Server: postgres +Username: postgres +Password: postgres +Database: postgres ``` -Adminer is available at: +## Configuration + +By default, the example connects to: ```text -http://localhost:8080 +postgres://postgres:postgres@localhost:5432/postgres?sslmode=disable ``` -Sign in to Adminer with: +To use another PostgreSQL instance, set `XPG_DATABASE_URL`: -```text -System: PostgreSQL -Server: postgres -Username: postgres -Password: postgres -Database: postgres +```shell +export XPG_DATABASE_URL='postgres://user:password@localhost:5432/database?sslmode=disable' ``` -> [!IMPORTANT] -> Use `postgres`, not `localhost`, in the **Server** field. Adminer connects to PostgreSQL through the Docker Compose -> network, where the database is discoverable by its service name. +The target database must contain the schema from `sql/schema.sql`. ## Run -From this example directory: +From this directory: ```shell go run . @@ -92,31 +81,26 @@ promo code: PROMO2026 promo applied: false ``` -The order is committed because the unique-key error is confined to the savepoint. `InSavepoint` rolls the failed promo -insert back before the outer transaction decides that this specific error is non-fatal. - -## Inspect the result +The failed promo insert is rolled back to the savepoint. The outer transaction then returns `nil`, so the order is +committed. -The embedded `setup.sql` recreates the `xpg_transactions_example` schema before each run and leaves the resulting data -available for inspection. +## Cleanup -In Adminer, open the `xpg_transactions_example` schema and inspect: +To remove the example schema and data: -```text -orders -promo_redemptions +```shell +psql 'postgres://postgres:postgres@localhost:5432/postgres?sslmode=disable' \ + -c 'DROP SCHEMA IF EXISTS xpg_transactions_example CASCADE;' ``` -## Stop services - -From the repository root: +Stop the local services: ```shell -docker compose -f examples/docker-compose.yml --profile tools down --remove-orphans -v +docker compose down ``` -Or from this example directory: +To also remove the PostgreSQL data volume: ```shell -docker compose -f ../docker-compose.yml --profile tools down --remove-orphans -v -``` \ No newline at end of file +docker compose down -v +``` diff --git a/examples/docker-compose.yml b/examples/transactions/docker-compose.yml similarity index 93% rename from examples/docker-compose.yml rename to examples/transactions/docker-compose.yml index 24523d9..c7ff2bf 100644 --- a/examples/docker-compose.yml +++ b/examples/transactions/docker-compose.yml @@ -18,8 +18,6 @@ services: adminer: image: adminer:standalone - profiles: - - tools environment: ADMINER_DEFAULT_SERVER: postgres ports: @@ -29,4 +27,4 @@ services: condition: service_healthy volumes: - postgres-data: \ No newline at end of file + postgres-data: diff --git a/examples/transactions/go.mod b/examples/transactions/go.mod index 0c3b7cc..ed21479 100644 --- a/examples/transactions/go.mod +++ b/examples/transactions/go.mod @@ -3,5 +3,6 @@ module transactions go 1.27 require ( + github.com/jackc/pgx/v5 v5.10.0 github.com/mkbeh/xpg v0.2.0 -) \ No newline at end of file +) diff --git a/examples/transactions/main.go b/examples/transactions/main.go index 5cc5c1a..782dfd6 100644 --- a/examples/transactions/main.go +++ b/examples/transactions/main.go @@ -2,24 +2,15 @@ package main import ( "context" - _ "embed" - "errors" "fmt" "log" "os" "github.com/jackc/pgx/v5" - "github.com/jackc/pgx/v5/pgconn" "github.com/mkbeh/xpg" ) -const ( - defaultDatabaseURL = "postgres://postgres:postgres@localhost:5432/postgres?sslmode=disable" - uniqueViolationCode = "23505" -) - -//go:embed setup.sql -var setupSQL string +const defaultDatabaseURL = "postgres://postgres:postgres@localhost:5432/postgres?sslmode=disable" func main() { if err := run(context.Background()); err != nil { @@ -34,7 +25,7 @@ func run(ctx context.Context) error { xpg.WithName("transactions-example"), ) if err != nil { - return fmt.Errorf("create pool: %w", err) + return fmt.Errorf("open pool: %w", err) } defer pool.Close() @@ -42,10 +33,6 @@ func run(ctx context.Context) error { return fmt.Errorf("ping PostgreSQL: %w", err) } - if err := prepareExample(ctx, pool); err != nil { - return fmt.Errorf("prepare example: %w", err) - } - const ( orderID = int64(1) promoCode = "PROMO2026" @@ -55,14 +42,9 @@ func run(ctx context.Context) error { return fmt.Errorf("process order: %w", err) } - status, promoApplied, err := loadOrderResult( - ctx, - pool, - orderID, - promoCode, - ) + status, promoApplied, err := loadOrder(ctx, pool, orderID, promoCode) if err != nil { - return fmt.Errorf("load order result: %w", err) + return fmt.Errorf("load order: %w", err) } fmt.Printf("order ID: %d\n", orderID) @@ -73,22 +55,6 @@ func run(ctx context.Context) error { return nil } -func prepareExample( - ctx context.Context, - pool *xpg.Pool, -) error { - _, err := pool.Exec( - ctx, - setupSQL, - pgx.QueryExecModeSimpleProtocol, - ) - if err != nil { - return fmt.Errorf("execute setup SQL: %w", err) - } - - return nil -} - func processOrder( ctx context.Context, pool *xpg.Pool, @@ -99,15 +65,16 @@ func processOrder( ctx, pgx.TxOptions{}, func(ctx context.Context, tx pgx.Tx) error { - // The order must be committed even when the optional promo fails. if _, err := tx.Exec( ctx, `INSERT INTO xpg_transactions_example.orders (id, status) - VALUES ($1, $2)`, + VALUES ($1, $2) + ON CONFLICT (id) DO UPDATE + SET status = EXCLUDED.status`, orderID, "new", ); err != nil { - return fmt.Errorf("create order: %w", err) + return fmt.Errorf("upsert order: %w", err) } err := xpg.InSavepoint( @@ -116,16 +83,11 @@ func processOrder( func(ctx context.Context, savepoint pgx.Tx) error { _, err := savepoint.Exec( ctx, - ` - INSERT INTO xpg_transactions_example.promo_redemptions ( - code, - order_id - ) - VALUES ( - $1, - $2 - ) - `, + `INSERT INTO xpg_transactions_example.promo_redemptions ( + code, + order_id + ) + VALUES ($1, $2)`, promoCode, orderID, ) @@ -137,9 +99,8 @@ func processOrder( return nil } - if pgErr, ok := errors.AsType[*pgconn.PgError](err); ok && - pgErr.Code == uniqueViolationCode { - // InSavepoint has already rolled back the failed promo insert. + if xpg.IsUniqueViolation(err) { + // InSavepoint already rolled back the failed promo insert. return nil } @@ -148,7 +109,7 @@ func processOrder( ) } -func loadOrderResult( +func loadOrder( ctx context.Context, pool *xpg.Pool, orderID int64, @@ -169,8 +130,8 @@ func loadOrderResult( WHERE p.order_id = o.id AND p.code = $2 ) - FROM xpg_transactions_example.orders AS o - WHERE o.id = $1`, + FROM xpg_transactions_example.orders AS o + WHERE o.id = $1`, orderID, promoCode, ).Scan( diff --git a/examples/transactions/setup.sql b/examples/transactions/setup.sql deleted file mode 100644 index a2638ec..0000000 --- a/examples/transactions/setup.sql +++ /dev/null @@ -1,24 +0,0 @@ -BEGIN; - -DROP SCHEMA IF EXISTS xpg_transactions_example CASCADE; - -CREATE SCHEMA xpg_transactions_example; - -CREATE TABLE xpg_transactions_example.orders -( - id bigint PRIMARY KEY, - status text NOT NULL -); - -CREATE TABLE xpg_transactions_example.promo_redemptions -( - code text PRIMARY KEY, - order_id bigint NOT NULL -); - -INSERT INTO xpg_transactions_example.promo_redemptions (code, - order_id) -VALUES ('PROMO2026', - 100); - -COMMIT; \ No newline at end of file diff --git a/examples/transactions/sql/schema.sql b/examples/transactions/sql/schema.sql new file mode 100644 index 0000000..94e4185 --- /dev/null +++ b/examples/transactions/sql/schema.sql @@ -0,0 +1,16 @@ +CREATE SCHEMA IF NOT EXISTS xpg_transactions_example; + +CREATE TABLE IF NOT EXISTS xpg_transactions_example.orders ( + id bigint PRIMARY KEY, + status text NOT NULL +); + +CREATE TABLE IF NOT EXISTS xpg_transactions_example.promo_redemptions ( + code text PRIMARY KEY, + order_id bigint NOT NULL +); + +INSERT INTO xpg_transactions_example.promo_redemptions (code, order_id) +VALUES ('PROMO2026', 100) +ON CONFLICT (code) DO UPDATE +SET order_id = EXCLUDED.order_id; From bf83e5acb4ef6aeb6bd32eabbcd01a5f4b7b936b Mon Sep 17 00:00:00 2001 From: mkbeh Date: Tue, 25 Aug 2026 14:24:41 +0300 Subject: [PATCH 27/41] refactor: simplify advisory example --- examples/advisory/README.md | 112 ++++++++---------- examples/advisory/docker-compose.yml | 30 +++++ examples/advisory/go.mod | 2 +- examples/advisory/main.go | 55 ++------- .../advisory/{setup.sql => sql/schema.sql} | 0 5 files changed, 87 insertions(+), 112 deletions(-) create mode 100644 examples/advisory/docker-compose.yml rename examples/advisory/{setup.sql => sql/schema.sql} (100%) diff --git a/examples/advisory/README.md b/examples/advisory/README.md index dfea9b2..dcfbaf8 100644 --- a/examples/advisory/README.md +++ b/examples/advisory/README.md @@ -1,59 +1,37 @@ # Transaction advisory locks -This example coordinates concurrent workers with PostgreSQL transaction-level advisory locks. The first worker acquires -an advisory lock and holds it until its transaction commits. A second worker uses the non-blocking try variant and -cannot -enter the protected section while the lock is held. After the first transaction commits, a third worker acquires the -same -lock successfully. +This example shows how transaction-level advisory locks protect a shared operation across concurrent workers: -**This example demonstrates:** +* Acquire a lock for the lifetime of a transaction +* Check the same lock without blocking +* Release the lock automatically when the transaction completes -* Acquiring a transaction-level lock with `xpg.AdvisoryXactLock` -* Trying to acquire a lock without waiting with `xpg.TryAdvisoryXactLock` -* Holding a lock for the lifetime of an explicit `pgx.Tx` -* Releasing a transaction-level lock automatically on commit or rollback -* Coordinating concurrent database work without `pg_sleep` +Worker A holds the lock while its transaction is open. Worker B cannot enter the protected section, and worker C +acquires the lock after worker A commits. -## Configuration - -The example uses the following connection string by default: - -```text -postgres://postgres:postgres@localhost:5432/postgres?sslmode=disable -``` +## Local setup -Set `XPG_DATABASE_URL` to use another PostgreSQL instance: +From this directory, start PostgreSQL and Adminer: ```shell -export XPG_DATABASE_URL='postgres://user:password@localhost:5432/database?sslmode=disable' +docker compose up -d ``` -> [!NOTE] -> The example runs two transactions concurrently, so the pool must allow at least two connections. The default pgxpool -> configuration satisfies this requirement. - -## Local PostgreSQL setup - -From the repository root, start PostgreSQL and Adminer: +Apply the example schema: ```shell -docker compose -f examples/docker-compose.yml --profile tools up -d +psql 'postgres://postgres:postgres@localhost:5432/postgres?sslmode=disable' \ + < sql/schema.sql ``` -Or from this directory: - -```shell -docker compose -f ../docker-compose.yml --profile tools up -d -``` - -Adminer is available at: +The services are available at: ```text -http://localhost:8080 +PostgreSQL: localhost:5432 +Adminer: http://localhost:8080 ``` -Sign in to Adminer with: +To inspect the example data in Adminer, sign in with: ```text System: PostgreSQL @@ -63,9 +41,24 @@ Password: postgres Database: postgres ``` -> [!IMPORTANT] -> Use `postgres`, not `localhost`, in the **Server** field. Adminer connects through the Docker Compose network, where -> PostgreSQL is available by its service name. +## Configuration + +By default, the example connects to: + +```text +postgres://postgres:postgres@localhost:5432/postgres?sslmode=disable +``` + +To use another PostgreSQL instance, set `XPG_DATABASE_URL`: + +```shell +export XPG_DATABASE_URL='postgres://user:password@localhost:5432/database?sslmode=disable' +``` + +The target database must contain the schema from `sql/schema.sql`. + +> [!NOTE] +> The example runs two transactions concurrently, so the pool must allow at least two connections. ## Run @@ -81,19 +74,6 @@ Or from the repository root: go run ./examples/advisory ``` -## Flow - -This example is easier to follow as a sequence: - -```text -1. Reset the example state -2. Worker A acquires the advisory lock -3. Worker B tries the same lock without waiting -4. Worker A commits and releases the lock -5. Worker C acquires the released lock -6. Read the committed job runs -``` - ## Expected output ```text @@ -106,20 +86,26 @@ recorded job runs: - worker-c (lock key: 2026) ``` -> [!IMPORTANT] -> Advisory lock keys are application-defined `int64` values. Use a stable key mapping and keep the protected transaction -> short because it holds both the lock and a pool connection until commit or rollback. +Transaction-level advisory locks are released automatically on commit or rollback. Keep the protected transaction short +because it holds both the lock and a pool connection. -## Stop services +## Cleanup -From the repository root: +To remove the example schema and data: ```shell -docker compose -f examples/docker-compose.yml --profile tools down --remove-orphans -v +psql 'postgres://postgres:postgres@localhost:5432/postgres?sslmode=disable' \ + -c 'DROP SCHEMA IF EXISTS xpg_advisory_example CASCADE;' ``` -Or from this directory: +Stop the local services: ```shell -docker compose -f ../docker-compose.yml --profile tools down --remove-orphans -v -``` \ No newline at end of file +docker compose down +``` + +To also remove the PostgreSQL data volume: + +```shell +docker compose down -v +``` diff --git a/examples/advisory/docker-compose.yml b/examples/advisory/docker-compose.yml new file mode 100644 index 0000000..c7ff2bf --- /dev/null +++ b/examples/advisory/docker-compose.yml @@ -0,0 +1,30 @@ +services: + postgres: + image: postgres:18-alpine + environment: + POSTGRES_DB: postgres + POSTGRES_USER: postgres + POSTGRES_PASSWORD: postgres + ports: + - "5432:5432" + volumes: + - postgres-data:/var/lib/postgresql + healthcheck: + test: "pg_isready -U $${POSTGRES_USER} -d $${POSTGRES_DB}" + interval: 2s + timeout: 5s + retries: 10 + start_period: 5s + + adminer: + image: adminer:standalone + environment: + ADMINER_DEFAULT_SERVER: postgres + ports: + - "8080:8080" + depends_on: + postgres: + condition: service_healthy + +volumes: + postgres-data: diff --git a/examples/advisory/go.mod b/examples/advisory/go.mod index 20fe69d..da98d91 100644 --- a/examples/advisory/go.mod +++ b/examples/advisory/go.mod @@ -5,4 +5,4 @@ go 1.27 require ( github.com/jackc/pgx/v5 v5.10.0 github.com/mkbeh/xpg v0.2.0 -) \ No newline at end of file +) diff --git a/examples/advisory/main.go b/examples/advisory/main.go index 8df67db..25fd7c4 100644 --- a/examples/advisory/main.go +++ b/examples/advisory/main.go @@ -2,7 +2,6 @@ package main import ( "context" - _ "embed" "errors" "fmt" "log" @@ -17,9 +16,6 @@ const ( jobLockKey = int64(2026) ) -//go:embed setup.sql -var setupSQL string - type jobRun struct { Worker string LockKey int64 @@ -38,7 +34,7 @@ func run(ctx context.Context) error { xpg.WithName("advisory-example"), ) if err != nil { - return fmt.Errorf("create pool: %w", err) + return fmt.Errorf("open pool: %w", err) } defer pool.Close() @@ -46,10 +42,6 @@ func run(ctx context.Context) error { return fmt.Errorf("ping PostgreSQL: %w", err) } - if err := prepareExample(ctx, pool); err != nil { - return fmt.Errorf("prepare example: %w", err) - } - lockAcquired := make(chan struct{}) releaseLock := make(chan struct{}) holderDone := make(chan error, 1) @@ -78,13 +70,7 @@ func run(ctx context.Context) error { return ctx.Err() } - workerBAcquired, workerBErr := tryRunJob( - ctx, - pool, - "worker-b", - jobLockKey, - ) - fmt.Printf("worker-b acquired the lock: %t\n", workerBAcquired) + workerBAcquired, workerBErr := tryRunJob(ctx, pool, "worker-b", jobLockKey) close(releaseLock) holderErr := <-holderDone @@ -101,14 +87,10 @@ func run(ctx context.Context) error { return errors.New("worker-b acquired a lock that should still be held") } + fmt.Printf("worker-b acquired the lock: %t\n", workerBAcquired) fmt.Println("worker-a committed and released the lock") - workerCAcquired, err := tryRunJob( - ctx, - pool, - "worker-c", - jobLockKey, - ) + workerCAcquired, err := tryRunJob(ctx, pool, "worker-c", jobLockKey) if err != nil { return fmt.Errorf("worker-c: %w", err) } @@ -132,22 +114,6 @@ func run(ctx context.Context) error { return nil } -func prepareExample( - ctx context.Context, - pool *xpg.Pool, -) error { - _, err := pool.Exec( - ctx, - setupSQL, - pgx.QueryExecModeSimpleProtocol, - ) - if err != nil { - return fmt.Errorf("execute setup SQL: %w", err) - } - - return nil -} - func holdJobLock( ctx context.Context, pool *xpg.Pool, @@ -195,14 +161,10 @@ func tryRunJob( var err error acquired, err = xpg.TryAdvisoryXactLock(ctx, tx, lockKey) - if err != nil { + if err != nil || !acquired { return err } - if !acquired { - return nil - } - if err := recordJobRun(ctx, tx, worker, lockKey); err != nil { return fmt.Errorf("record job run: %w", err) } @@ -226,7 +188,7 @@ func recordJobRun( _, err := tx.Exec( ctx, `INSERT INTO xpg_advisory_example.job_runs (worker, lock_key) - VALUES ($1, $2)`, + VALUES ($1, $2)`, worker, lockKey, ) @@ -234,10 +196,7 @@ func recordJobRun( return err } -func loadJobRuns( - ctx context.Context, - pool *xpg.Pool, -) ([]jobRun, error) { +func loadJobRuns(ctx context.Context, pool *xpg.Pool) ([]jobRun, error) { rows, err := pool.Query( ctx, `SELECT worker, lock_key diff --git a/examples/advisory/setup.sql b/examples/advisory/sql/schema.sql similarity index 100% rename from examples/advisory/setup.sql rename to examples/advisory/sql/schema.sql From 1fd7ee62b1b6e1ae138cab087342c7672c6d15bc Mon Sep 17 00:00:00 2001 From: mkbeh Date: Tue, 25 Aug 2026 14:38:55 +0300 Subject: [PATCH 28/41] refactor: simplify observability example --- examples/observability/README.md | 171 ++++------------------ examples/observability/docker-compose.yml | 30 ++++ examples/observability/go.mod | 23 --- examples/observability/main.go | 64 +++++--- examples/observability/metrics.go | 40 ++--- examples/observability/otel.go | 43 ++---- 6 files changed, 132 insertions(+), 239 deletions(-) create mode 100644 examples/observability/docker-compose.yml diff --git a/examples/observability/README.md b/examples/observability/README.md index f5f8509..9f6fff1 100644 --- a/examples/observability/README.md +++ b/examples/observability/README.md @@ -1,45 +1,27 @@ -# Observability Example +# Observability -This example shows how to combine logging, distributed tracing, and connection pool metrics with `xpg`. +This example shows how to add observability to `xpg`: -**This example demonstrates:** - -* Adapting the standard library `log/slog` logger to `pgx/tracelog` -* Attaching an OpenTelemetry PostgreSQL tracer through `xpg.WithTracer` -* Combining the logger and tracer automatically through the `xpg` tracing pipeline -* Exporting `xpg` pool metrics through OpenTelemetry and Prometheus -* Propagating an application span through concurrent PostgreSQL operations - -The signals remain independent at the application boundary: - -```text -slog otelpgx otelxpg - | | | - +---- pgx tracing -----+ OTel metrics - | | - xpg Prometheus - | - /metrics -``` - -`otelpgx.RecordStats` is intentionally not used here. PostgreSQL tracing is handled by `otelpgx`, while pool metrics are -owned by `xpg` and exported through `extra/otelxpg`. +* Log PostgreSQL activity with `slog` +* Trace database operations with OpenTelemetry +* Export connection-pool metrics to Prometheus +* Generate concurrent load to observe pool behavior ## Configuration -The example uses the following connection string by default: +By default, the example connects to: ```text postgres://postgres:postgres@localhost:5432/postgres?sslmode=disable ``` -Set `XPG_DATABASE_URL` to use another PostgreSQL instance: +To use another PostgreSQL instance, set `XPG_DATABASE_URL`: ```shell export XPG_DATABASE_URL='postgres://user:password@localhost:5432/database?sslmode=disable' ``` -The HTTP server listens on `localhost:9464` by default. Override it with: +The HTTP server listens on `localhost:9464`. To use another address, set `HTTP_ADDR`: ```shell export HTTP_ADDR='localhost:9464' @@ -47,35 +29,19 @@ export HTTP_ADDR='localhost:9464' ## Local setup -Start PostgreSQL and Adminer from the repository root: +From this directory, start PostgreSQL and Adminer: ```shell -docker compose -f examples/docker-compose.yml --profile tools up -d +docker compose up -d ``` -Or from this example directory: - -```shell -docker compose -f ../docker-compose.yml --profile tools up -d -``` - -Services are available at: +The services are available at: ```text PostgreSQL: localhost:5432 Adminer: http://localhost:8080 ``` -Sign in to Adminer with: - -```text -System: PostgreSQL -Server: postgres -Username: postgres -Password: postgres -Database: postgres -``` - ## Run From this directory: @@ -90,125 +56,46 @@ Or from the repository root: go run ./examples/observability ``` -The example starts an HTTP server at: - -```text -http://localhost:9464 -``` - -## Logging - -The example adapts `log/slog` to `tracelog.Logger` and passes it to the pool: - -```go -xpg.WithLogger( - newPGXLogger(logger), - tracelog.LogLevelInfo, -) -``` - -The adapter also copies the active OpenTelemetry `trace_id` and `span_id` into pgx log records when a span is present, -which makes logs and database spans directly correlatable. - -`pgx/tracelog` can include SQL text and query arguments in log records. Production applications should choose logging -levels and redaction policies appropriate for the data they process. - -## Tracing - -`otelpgx` is attached as a normal pgx tracer: - -```go -xpg.WithTracer( - otelpgx.NewTracer( - otelpgx.WithTracerProvider(tracing.TracerProvider()), - ), -) -``` - -The example writes completed spans to stdout as formatted JSON. The stdout exporter and synchronous span processor are -used only to make the example self-contained and immediately observable. Production applications should normally export -traces through OTLP and use a batch span processor. - -Generate a traced workload: +The HTTP server starts on: ```shell -curl -X POST 'http://localhost:9464/load' +localhost:9464 ``` -The handler creates one application span and passes its context to six concurrent PostgreSQL operations. Database spans -created by `otelpgx` therefore appear as children of that application span. +## Generate load -## Metrics - -Pool metrics use an explicit OpenTelemetry `MeterProvider` backed by the Prometheus exporter: - -```go -xpg.WithMetrics( - otelxpg.NewMetrics( - otelxpg.WithMeterProvider(metrics.MeterProvider()), - ), -) -``` - -Open the Prometheus endpoint: +Run six concurrent one-second queries against a pool limited to two connections: ```shell -curl 'http://localhost:9464/metrics' +curl -X POST 'http://localhost:9464/load' ``` -Show only database connection pool and `xpg` metrics: +While the request is running, inspect pool contention from another terminal: ```shell curl -s 'http://localhost:9464/metrics' \ - | grep -E '^(db_client_connection|xpg_pool_connection_)' -``` - -The example exports these metric families: - -```text -db_client_connection_count -db_client_connection_max -xpg_pool_connection_constructing -xpg_pool_connection_acquire_count_total -xpg_pool_connection_acquire_time_seconds_total -xpg_pool_connection_acquire_canceled_count_total -xpg_pool_connection_acquire_empty_count_total -xpg_pool_connection_acquire_empty_wait_time_seconds_total -xpg_pool_connection_create_count_total -xpg_pool_connection_destroy_count_total -``` - -The Prometheus exporter converts OpenTelemetry dotted instrument names to Prometheus-compatible names and adds unit and -counter suffixes where required. - -## Generate pool contention - -The pool is intentionally limited to two connections. Run: - -```shell -curl -X POST 'http://localhost:9464/load' + | grep -E 'db_client_connection_count|xpg_pool_connection_acquire_' ``` -While it is running, inspect metrics from another terminal: +To inspect all exported metrics: ```shell -curl -s 'http://localhost:9464/metrics' \ - | grep -E 'db_client_connection_count|xpg_pool_connection_acquire_' +curl 'http://localhost:9464/metrics' ``` -Six concurrent one-second queries make connection usage and acquire wait metrics visible while also producing related -logs and spans. +The request produces pgx logs and OpenTelemetry spans in the application output, while pool metrics remain available +from the `/metrics` endpoint. -## Stop services +## Cleanup -From the repository root: +Stop the local services: ```shell -docker compose -f examples/docker-compose.yml --profile tools down --remove-orphans -v +docker compose down ``` -Or from this example directory: +To also remove the PostgreSQL data volume: ```shell -docker compose -f ../docker-compose.yml --profile tools down --remove-orphans -v -``` +docker compose down -v +``` \ No newline at end of file diff --git a/examples/observability/docker-compose.yml b/examples/observability/docker-compose.yml new file mode 100644 index 0000000..c7ff2bf --- /dev/null +++ b/examples/observability/docker-compose.yml @@ -0,0 +1,30 @@ +services: + postgres: + image: postgres:18-alpine + environment: + POSTGRES_DB: postgres + POSTGRES_USER: postgres + POSTGRES_PASSWORD: postgres + ports: + - "5432:5432" + volumes: + - postgres-data:/var/lib/postgresql + healthcheck: + test: "pg_isready -U $${POSTGRES_USER} -d $${POSTGRES_DB}" + interval: 2s + timeout: 5s + retries: 10 + start_period: 5s + + adminer: + image: adminer:standalone + environment: + ADMINER_DEFAULT_SERVER: postgres + ports: + - "8080:8080" + depends_on: + postgres: + condition: service_healthy + +volumes: + postgres-data: diff --git a/examples/observability/go.mod b/examples/observability/go.mod index 9a67ba3..190c8cb 100644 --- a/examples/observability/go.mod +++ b/examples/observability/go.mod @@ -16,26 +16,3 @@ require ( go.opentelemetry.io/otel/sdk/metric v1.45.0 go.opentelemetry.io/otel/trace v1.45.0 ) - -require ( - github.com/beorn7/perks v1.0.1 // indirect - github.com/cespare/xxhash/v2 v2.3.0 // indirect - github.com/go-logr/logr v1.4.4 // indirect - github.com/go-logr/stdr v1.2.2 // indirect - github.com/google/uuid v1.6.0 // indirect - github.com/jackc/pgpassfile v1.0.0 // indirect - github.com/jackc/pgservicefile v0.0.0-20240606120523-5a60cdf6a761 // indirect - github.com/jackc/puddle/v2 v2.2.2 // indirect - github.com/munnerz/goautoneg v0.0.0-20191010083416-a7dc8b61c822 // indirect - github.com/prometheus/client_model v0.6.2 // indirect - github.com/prometheus/common v0.70.1 // indirect - github.com/prometheus/otlptranslator v1.0.0 // indirect - github.com/prometheus/procfs v0.21.1 // indirect - go.opentelemetry.io/auto/sdk v1.2.1 // indirect - go.opentelemetry.io/otel/metric v1.45.0 // indirect - go.opentelemetry.io/otel/metric/x v0.67.0 // indirect - golang.org/x/sync v0.22.0 // indirect - golang.org/x/sys v0.47.0 // indirect - golang.org/x/text v0.41.0 // indirect - google.golang.org/protobuf v1.36.11 // indirect -) diff --git a/examples/observability/main.go b/examples/observability/main.go index fab9f55..3364567 100644 --- a/examples/observability/main.go +++ b/examples/observability/main.go @@ -50,36 +50,42 @@ func main() { } func run(ctx context.Context, logger *slog.Logger) (runErr error) { - res, err := newOTelResource(ctx) + resource, err := newOTelResource(ctx) if err != nil { return err } - metrics, err := newMetricsRuntime(res) + meterProvider, metricsHandler, err := newMeterProvider(resource) if err != nil { return fmt.Errorf("initialize metrics: %w", err) } defer func() { - shutdownCtx, cancel := context.WithTimeout(context.Background(), 5*time.Second) + shutdownCtx, cancel := context.WithTimeout( + context.Background(), + 5*time.Second, + ) defer cancel() runErr = errors.Join( runErr, - metrics.Shutdown(shutdownCtx), + meterProvider.Shutdown(shutdownCtx), ) }() - tracing, err := newTracingRuntime(res) + tracerProvider, err := newTracerProvider(resource) if err != nil { return fmt.Errorf("initialize tracing: %w", err) } defer func() { - shutdownCtx, cancel := context.WithTimeout(context.Background(), 5*time.Second) + shutdownCtx, cancel := context.WithTimeout( + context.Background(), + 5*time.Second, + ) defer cancel() runErr = errors.Join( runErr, - tracing.Shutdown(shutdownCtx), + tracerProvider.Shutdown(shutdownCtx), ) }() @@ -91,7 +97,9 @@ func run(ctx context.Context, logger *slog.Logger) (runErr error) { // A small pool makes contention visible when POST /load runs. config.MaxConns = 2 - pgxLogger := logger.With(slog.String("component", "pgx")) + pgxLogger := logger.With( + slog.String("component", "pgx"), + ) pool, err := xpg.New( ctx, @@ -104,16 +112,13 @@ func run(ctx context.Context, logger *slog.Logger) (runErr error) { ), xpg.WithTracer( otelpgx.NewTracer( - otelpgx.WithTracerProvider( - tracing.TracerProvider(), - ), + otelpgx.WithTracerProvider(tracerProvider), + otelpgx.WithTrimSQLInSpanName(), ), ), xpg.WithMetrics( otelxpg.NewMetrics( - otelxpg.WithMeterProvider( - metrics.MeterProvider(), - ), + otelxpg.WithMeterProvider(meterProvider), ), ), ) @@ -127,8 +132,14 @@ func run(ctx context.Context, logger *slog.Logger) (runErr error) { } mux := http.NewServeMux() - mux.Handle("GET /metrics", metrics.Handler()) - mux.HandleFunc("POST /load", loadHandler(pool, tracing.Tracer())) + mux.Handle("GET /metrics", metricsHandler) + mux.HandleFunc( + "POST /load", + loadHandler( + pool, + tracerProvider.Tracer(tracingInstrumentationName), + ), + ) server := &http.Server{ Addr: httpAddress(), @@ -165,7 +176,10 @@ func serveHTTP(ctx context.Context, server *http.Server) error { case <-ctx.Done(): } - shutdownCtx, cancel := context.WithTimeout(context.Background(), 5*time.Second) + shutdownCtx, cancel := context.WithTimeout( + context.Background(), + 5*time.Second, + ) defer cancel() if err := server.Shutdown(shutdownCtx); err != nil { @@ -179,9 +193,15 @@ func serveHTTP(ctx context.Context, server *http.Server) error { return nil } -func loadHandler(pool *xpg.Pool, tracer trace.Tracer) http.HandlerFunc { +func loadHandler( + pool *xpg.Pool, + tracer trace.Tracer, +) http.HandlerFunc { return func(w http.ResponseWriter, request *http.Request) { - ctx, span := tracer.Start(request.Context(), "run-load") + ctx, span := tracer.Start( + request.Context(), + "run-load", + ) defer span.End() startedAt := time.Now() @@ -199,7 +219,11 @@ func loadHandler(pool *xpg.Pool, tracer trace.Tracer) http.HandlerFunc { return } - _, _ = fmt.Fprintf(w, "workload completed in %s\n", time.Since(startedAt)) + _, _ = fmt.Fprintf( + w, + "workload completed in %s\n", + time.Since(startedAt), + ) } } diff --git a/examples/observability/metrics.go b/examples/observability/metrics.go index 3d548a0..84fb692 100644 --- a/examples/observability/metrics.go +++ b/examples/observability/metrics.go @@ -1,7 +1,6 @@ package main import ( - "context" "fmt" "net/http" @@ -12,12 +11,9 @@ import ( "go.opentelemetry.io/otel/sdk/resource" ) -type metricsRuntime struct { - handler http.Handler - meterProvider *sdkmetric.MeterProvider -} - -func newMetricsRuntime(res *resource.Resource) (*metricsRuntime, error) { +func newMeterProvider( + resource *resource.Resource, +) (*sdkmetric.MeterProvider, http.Handler, error) { registry := promclient.NewRegistry() exporter, err := otelprom.New( @@ -25,31 +21,21 @@ func newMetricsRuntime(res *resource.Resource) (*metricsRuntime, error) { otelprom.WithoutScopeInfo(), ) if err != nil { - return nil, fmt.Errorf("create Prometheus exporter: %w", err) + return nil, nil, fmt.Errorf( + "create Prometheus exporter: %w", + err, + ) } meterProvider := sdkmetric.NewMeterProvider( - sdkmetric.WithResource(res), + sdkmetric.WithResource(resource), sdkmetric.WithReader(exporter), ) - return &metricsRuntime{ - handler: promhttp.HandlerFor( - registry, - promhttp.HandlerOpts{}, - ), - meterProvider: meterProvider, - }, nil -} - -func (m *metricsRuntime) Handler() http.Handler { - return m.handler -} - -func (m *metricsRuntime) MeterProvider() *sdkmetric.MeterProvider { - return m.meterProvider -} + handler := promhttp.HandlerFor( + registry, + promhttp.HandlerOpts{}, + ) -func (m *metricsRuntime) Shutdown(ctx context.Context) error { - return m.meterProvider.Shutdown(ctx) + return meterProvider, handler, nil } diff --git a/examples/observability/otel.go b/examples/observability/otel.go index 6a76c81..df28add 100644 --- a/examples/observability/otel.go +++ b/examples/observability/otel.go @@ -8,15 +8,10 @@ import ( "go.opentelemetry.io/otel/sdk/resource" sdktrace "go.opentelemetry.io/otel/sdk/trace" semconv "go.opentelemetry.io/otel/semconv/v1.43.0" - "go.opentelemetry.io/otel/trace" ) const tracingInstrumentationName = "github.com/mkbeh/xpg/examples/observability" -type tracingRuntime struct { - tracerProvider *sdktrace.TracerProvider -} - func newOTelResource(ctx context.Context) (*resource.Resource, error) { res, err := resource.New( ctx, @@ -28,41 +23,35 @@ func newOTelResource(ctx context.Context) (*resource.Resource, error) { ), ) if err != nil { - return nil, fmt.Errorf("create OpenTelemetry resource: %w", err) + return nil, fmt.Errorf( + "create OpenTelemetry resource: %w", + err, + ) } return res, nil } -func newTracingRuntime(res *resource.Resource) (*tracingRuntime, error) { +func newTracerProvider( + resource *resource.Resource, +) (*sdktrace.TracerProvider, error) { exporter, err := stdouttrace.New( stdouttrace.WithPrettyPrint(), ) if err != nil { - return nil, fmt.Errorf("create stdout trace exporter: %w", err) + return nil, fmt.Errorf( + "create stdout trace exporter: %w", + err, + ) } tracerProvider := sdktrace.NewTracerProvider( - sdktrace.WithResource(res), - // A synchronous processor keeps this runnable example easy to inspect. - // Production applications should normally prefer WithBatcher with an - // OTLP exporter. + sdktrace.WithResource(resource), + // The synchronous processor keeps the example easy to inspect. + // Production applications should normally use a batch processor with + // an OTLP exporter. sdktrace.WithSyncer(exporter), ) - return &tracingRuntime{ - tracerProvider: tracerProvider, - }, nil -} - -func (t *tracingRuntime) TracerProvider() *sdktrace.TracerProvider { - return t.tracerProvider -} - -func (t *tracingRuntime) Tracer() trace.Tracer { - return t.tracerProvider.Tracer(tracingInstrumentationName) -} - -func (t *tracingRuntime) Shutdown(ctx context.Context) error { - return t.tracerProvider.Shutdown(ctx) + return tracerProvider, nil } From 0b9789ccac8401885a0198655235a0f305a3abb2 Mon Sep 17 00:00:00 2001 From: mkbeh Date: Tue, 25 Aug 2026 14:43:42 +0300 Subject: [PATCH 29/41] refactor: observability minor fixes --- examples/observability/main.go | 5 +---- examples/observability/metrics.go | 9 ++------- examples/observability/otel.go | 14 +++----------- 3 files changed, 6 insertions(+), 22 deletions(-) diff --git a/examples/observability/main.go b/examples/observability/main.go index 3364567..6b8f920 100644 --- a/examples/observability/main.go +++ b/examples/observability/main.go @@ -193,10 +193,7 @@ func serveHTTP(ctx context.Context, server *http.Server) error { return nil } -func loadHandler( - pool *xpg.Pool, - tracer trace.Tracer, -) http.HandlerFunc { +func loadHandler(pool *xpg.Pool, tracer trace.Tracer) http.HandlerFunc { return func(w http.ResponseWriter, request *http.Request) { ctx, span := tracer.Start( request.Context(), diff --git a/examples/observability/metrics.go b/examples/observability/metrics.go index 84fb692..c49115f 100644 --- a/examples/observability/metrics.go +++ b/examples/observability/metrics.go @@ -11,9 +11,7 @@ import ( "go.opentelemetry.io/otel/sdk/resource" ) -func newMeterProvider( - resource *resource.Resource, -) (*sdkmetric.MeterProvider, http.Handler, error) { +func newMeterProvider(resource *resource.Resource) (*sdkmetric.MeterProvider, http.Handler, error) { registry := promclient.NewRegistry() exporter, err := otelprom.New( @@ -21,10 +19,7 @@ func newMeterProvider( otelprom.WithoutScopeInfo(), ) if err != nil { - return nil, nil, fmt.Errorf( - "create Prometheus exporter: %w", - err, - ) + return nil, nil, fmt.Errorf("create Prometheus exporter: %w", err) } meterProvider := sdkmetric.NewMeterProvider( diff --git a/examples/observability/otel.go b/examples/observability/otel.go index df28add..b0d319d 100644 --- a/examples/observability/otel.go +++ b/examples/observability/otel.go @@ -23,26 +23,18 @@ func newOTelResource(ctx context.Context) (*resource.Resource, error) { ), ) if err != nil { - return nil, fmt.Errorf( - "create OpenTelemetry resource: %w", - err, - ) + return nil, fmt.Errorf("create OpenTelemetry resource: %w", err) } return res, nil } -func newTracerProvider( - resource *resource.Resource, -) (*sdktrace.TracerProvider, error) { +func newTracerProvider(resource *resource.Resource) (*sdktrace.TracerProvider, error) { exporter, err := stdouttrace.New( stdouttrace.WithPrettyPrint(), ) if err != nil { - return nil, fmt.Errorf( - "create stdout trace exporter: %w", - err, - ) + return nil, fmt.Errorf("create stdout trace exporter: %w", err) } tracerProvider := sdktrace.NewTracerProvider( From 2a00c49229e99858973c02b907e4751e84a442c7 Mon Sep 17 00:00:00 2001 From: mkbeh Date: Tue, 25 Aug 2026 15:29:23 +0300 Subject: [PATCH 30/41] refactor: simplify cluster example --- examples/cluster/README.md | 107 ++++++++------------- examples/cluster/docker-compose.yml | 5 - examples/cluster/main.go | 134 ++------------------------- examples/cluster/setup.go | 96 +++++++++++++++++++ examples/cluster/sql/primary.sql | 10 +- examples/cluster/sql/replica-one.sql | 4 + examples/cluster/sql/replica-two.sql | 4 + 7 files changed, 158 insertions(+), 202 deletions(-) create mode 100644 examples/cluster/setup.go diff --git a/examples/cluster/README.md b/examples/cluster/README.md index 24c9684..114f36d 100644 --- a/examples/cluster/README.md +++ b/examples/cluster/README.md @@ -1,67 +1,47 @@ # Cluster routing -This example routes reads and transactions across one primary pool and two replica pools with `cluster.Cluster`. +This example shows how `cluster.Cluster` coordinates primary and replica access: -```text - ┌─ primary -application ─ cluster - ├─ replica-one - └─ replica-two -``` - -**This example demonstrates:** - -* Creating a cluster from primary and replica pools -* Routing reads to the primary or replicas -* Distributing replica reads with round-robin -* Running primary and read-only replica transactions +* Route reads explicitly to the primary or replicas +* Distribute replica reads with round-robin selection +* Run write transactions on the primary +* Run read-only transactions on replicas > [!NOTE] -> The local containers are independent PostgreSQL instances used to demonstrate routing. They do not configure streaming -replication. - -## Configuration - -The example uses the following connection strings by default: +> The local services are independent PostgreSQL instances used only to demonstrate routing. They do not configure +> streaming replication. -```text -XPG_PRIMARY_DATABASE_URL=postgres://postgres:postgres@localhost:55432/postgres?sslmode=disable&target_session_attrs=read-write -XPG_REPLICA_ONE_DATABASE_URL=postgres://postgres:postgres@localhost:55433/postgres?sslmode=disable&target_session_attrs=read-only -XPG_REPLICA_TWO_DATABASE_URL=postgres://postgres:postgres@localhost:55434/postgres?sslmode=disable&target_session_attrs=read-only -``` +## Local setup -Set the corresponding environment variables to use different PostgreSQL endpoints: +From this directory, start the three PostgreSQL nodes and Adminer: ```shell -export XPG_PRIMARY_DATABASE_URL='postgres://user:password@primary.example.com:5432/database?sslmode=disable&target_session_attrs=read-write' -export XPG_REPLICA_ONE_DATABASE_URL='postgres://user:password@replica-one.example.com:5432/database?sslmode=disable&target_session_attrs=read-only' -export XPG_REPLICA_TWO_DATABASE_URL='postgres://user:password@replica-two.example.com:5432/database?sslmode=disable&target_session_attrs=read-only' +docker compose up -d ``` -## Local setup - -Start the primary, both replica endpoints, and Adminer from the repository root: +Apply the node-specific setup: ```shell -docker compose -f examples/cluster/docker-compose.yml --profile tools up -d -``` +psql 'postgres://postgres:postgres@localhost:55432/postgres?sslmode=disable' \ + < sql/primary.sql -Or from this example directory: +psql 'postgres://postgres:postgres@localhost:55433/postgres?sslmode=disable' \ + < sql/replica-one.sql -```shell -docker compose --profile tools up -d +psql 'postgres://postgres:postgres@localhost:55434/postgres?sslmode=disable' \ + < sql/replica-two.sql ``` -Services are available at: +The services are available at: ```text -Primary: localhost:55432 -Replica 1: localhost:55433 -Replica 2: localhost:55434 -Adminer: http://localhost:58080 +Primary: localhost:55432 +Replica 1: localhost:55433 +Replica 2: localhost:55434 +Adminer: http://localhost:8080 ``` -Sign in to Adminer with: +To inspect the nodes in Adminer, sign in with: ```text System: PostgreSQL @@ -71,21 +51,20 @@ Password: postgres Database: postgres ``` -Use `postgres-replica-one` or `postgres-replica-two` in the **Server** field to inspect the replica endpoints. +Use `postgres-replica-one` or `postgres-replica-two` in the **Server** field to inspect a replica. -## Run +## Configuration -From this directory: +By default, the example connects to: -```shell -go run . +```text +Primary: postgres://postgres:postgres@localhost:55432/postgres?sslmode=disable&target_session_attrs=read-write +Replica 1: postgres://postgres:postgres@localhost:55433/postgres?sslmode=disable&target_session_attrs=read-only +Replica 2: postgres://postgres:postgres@localhost:55434/postgres?sslmode=disable&target_session_attrs=read-only ``` -Or from the repository root: - -```shell -go run ./examples/basic -``` +To use other PostgreSQL endpoints, set `XPG_PRIMARY_DATABASE_URL`, `XPG_REPLICA_ONE_DATABASE_URL`, and +`XPG_REPLICA_TWO_DATABASE_URL`. ## Run @@ -104,7 +83,7 @@ go run ./examples/cluster ## Expected output ```text -primary: +primary read: - pool=cluster.primary node=primary role=primary replica reads: - pool=cluster.replica-one node=replica-one role=replica @@ -114,25 +93,19 @@ transactions: - replica node=replica-one read_only=on ``` -The example performs one complete round-robin pass across the configured replicas before starting the read-only -transaction. - -Pools remain owned by the caller until `cluster.New` succeeds. After successful cluster creation, `cluster.Cluster` owns -the pools and closes them when `Cluster.Close` is called. +The two replica reads show one complete round-robin cycle. The following read transaction continues from the same +selector and therefore resolves the first replica again. -## Stop services +## Cleanup -From the repository root: +Stop the local services: ```shell -docker compose \ - -f examples/cluster/docker-compose.yml \ - --profile tools \ - down --remove-orphans -v +docker compose down ``` -Or from this example directory: +To also remove all PostgreSQL data volumes: ```shell -docker compose --profile tools down --remove-orphans -v +docker compose down -v ``` diff --git a/examples/cluster/docker-compose.yml b/examples/cluster/docker-compose.yml index 00dc2bc..44e7765 100644 --- a/examples/cluster/docker-compose.yml +++ b/examples/cluster/docker-compose.yml @@ -9,7 +9,6 @@ services: - "55432:5432" volumes: - primary-data:/var/lib/postgresql - - ./sql/primary.sql:/docker-entrypoint-initdb.d/10-node.sql:ro healthcheck: test: "pg_isready -U $${POSTGRES_USER} -d $${POSTGRES_DB}" interval: 2s @@ -27,7 +26,6 @@ services: - "55433:5432" volumes: - replica-one-data:/var/lib/postgresql - - ./sql/replica-one.sql:/docker-entrypoint-initdb.d/10-node.sql:ro healthcheck: test: "pg_isready -U $${POSTGRES_USER} -d $${POSTGRES_DB}" interval: 2s @@ -45,7 +43,6 @@ services: - "55434:5432" volumes: - replica-two-data:/var/lib/postgresql - - ./sql/replica-two.sql:/docker-entrypoint-initdb.d/10-node.sql:ro healthcheck: test: "pg_isready -U $${POSTGRES_USER} -d $${POSTGRES_DB}" interval: 2s @@ -55,8 +52,6 @@ services: adminer: image: adminer:standalone - profiles: - - tools environment: ADMINER_DEFAULT_SERVER: postgres-primary ports: diff --git a/examples/cluster/main.go b/examples/cluster/main.go index 8165240..0d614cc 100644 --- a/examples/cluster/main.go +++ b/examples/cluster/main.go @@ -4,25 +4,17 @@ import ( "context" "fmt" "log" - "os" "github.com/jackc/pgx/v5" - "github.com/mkbeh/xpg" "github.com/mkbeh/xpg/cluster" ) -const ( - defaultPrimaryDatabaseURL = "postgres://postgres:postgres@localhost:55432/postgres?sslmode=disable&target_session_attrs=read-write" - defaultReplicaOneDatabaseURL = "postgres://postgres:postgres@localhost:55433/postgres?sslmode=disable&target_session_attrs=read-only" - defaultReplicaTwoDatabaseURL = "postgres://postgres:postgres@localhost:55434/postgres?sslmode=disable&target_session_attrs=read-only" -) - type nodeInfo struct { Name string Role string } -type queryRower interface { +type rowQuerier interface { QueryRow(context.Context, string, ...any) pgx.Row } @@ -50,112 +42,18 @@ func run(ctx context.Context) error { return nil } -func openCluster(ctx context.Context) (*cluster.Cluster, error) { - nodes := []struct { - databaseURL string - name string - role string - }{ - { - databaseURL: environment( - "XPG_PRIMARY_DATABASE_URL", - defaultPrimaryDatabaseURL, - ), - name: "cluster.primary", - role: "primary", - }, - { - databaseURL: environment( - "XPG_REPLICA_ONE_DATABASE_URL", - defaultReplicaOneDatabaseURL, - ), - name: "cluster.replica-one", - role: "replica", - }, - { - databaseURL: environment( - "XPG_REPLICA_TWO_DATABASE_URL", - defaultReplicaTwoDatabaseURL, - ), - name: "cluster.replica-two", - role: "replica", - }, - } - - pools := make([]*xpg.Pool, 0, len(nodes)) - - for _, node := range nodes { - pool, err := openPool( - ctx, - node.databaseURL, - node.name, - node.role, - ) - if err != nil { - closePools(pools) - - return nil, fmt.Errorf("open %s pool: %w", node.name, err) - } - - pools = append(pools, pool) - } - - // Build a replicated cluster and explicitly use round-robin selection for - // replica reads. - dbCluster, err := cluster.New( - cluster.Config{ - Primary: pools[0], - Replicas: pools[1:], - Selector: cluster.RoundRobinSelector(), - }, - ) - if err != nil { - return nil, fmt.Errorf("create cluster: %w", err) - } - - return dbCluster, nil -} - -func openPool( - ctx context.Context, - databaseURL string, - name string, - role string, -) (*xpg.Pool, error) { - pool, err := xpg.Open( - ctx, - databaseURL, - xpg.WithName(name), - xpg.WithLabel("role", role), - ) +func showRouting(ctx context.Context, dbCluster *cluster.Cluster) error { + primary, err := dbCluster.ReadPool(ctx, cluster.ReadPrimary) if err != nil { - return nil, err - } - - if err := pool.Ping(ctx); err != nil { - pool.Close() - - return nil, fmt.Errorf("ping %s: %w", name, err) - } - - return pool, nil -} - -func closePools(pools []*xpg.Pool) { - for index := len(pools) - 1; index >= 0; index-- { - pools[index].Close() + return fmt.Errorf("resolve primary read: %w", err) } -} - -func showRouting(ctx context.Context, dbCluster *cluster.Cluster) error { - primary := dbCluster.Primary() node, err := loadNode(ctx, primary) if err != nil { return fmt.Errorf("read primary node: %w", err) } - fmt.Println("primary:") + fmt.Println("primary read:") fmt.Printf( "- pool=%s node=%s role=%s\n", primary.Name(), @@ -165,10 +63,8 @@ func showRouting(ctx context.Context, dbCluster *cluster.Cluster) error { fmt.Println("replica reads:") - // Read once per registered replica to demonstrate one complete round-robin - // cycle without hard-coding the cluster size. for range dbCluster.ReplicaCount() { - pool, err := dbCluster.ReadPool( + replica, err := dbCluster.ReadPool( ctx, cluster.ReadReplicaRequired, ) @@ -176,14 +72,14 @@ func showRouting(ctx context.Context, dbCluster *cluster.Cluster) error { return fmt.Errorf("resolve replica read: %w", err) } - node, err := loadNode(ctx, pool) + node, err := loadNode(ctx, replica) if err != nil { return fmt.Errorf("read replica node: %w", err) } fmt.Printf( "- pool=%s node=%s role=%s\n", - pool.Name(), + replica.Name(), node.Name, node.Role, ) @@ -198,8 +94,6 @@ func showTransactions(ctx context.Context, dbCluster *cluster.Cluster) error { primaryReadOnly string ) - // Primary transactions use regular pgx transaction options and may perform - // both reads and writes. err := dbCluster.InPrimaryTx( ctx, pgx.TxOptions{}, @@ -226,8 +120,6 @@ func showTransactions(ctx context.Context, dbCluster *cluster.Cluster) error { replicaReadOnly string ) - // Read transactions resolve their pool through ReadPolicy and always start - // PostgreSQL transactions in READ ONLY mode. err = dbCluster.InReadTx( ctx, cluster.ReadReplicaRequired, @@ -267,7 +159,7 @@ func showTransactions(ctx context.Context, dbCluster *cluster.Cluster) error { return nil } -func loadNode(ctx context.Context, db queryRower) (nodeInfo, error) { +func loadNode(ctx context.Context, db rowQuerier) (nodeInfo, error) { var node nodeInfo err := db.QueryRow( @@ -284,11 +176,3 @@ func loadNode(ctx context.Context, db queryRower) (nodeInfo, error) { return node, nil } - -func environment(key, fallback string) string { - if value := os.Getenv(key); value != "" { - return value - } - - return fallback -} diff --git a/examples/cluster/setup.go b/examples/cluster/setup.go new file mode 100644 index 0000000..fe62855 --- /dev/null +++ b/examples/cluster/setup.go @@ -0,0 +1,96 @@ +package main + +import ( + "context" + "fmt" + "os" + + "github.com/mkbeh/xpg" + "github.com/mkbeh/xpg/cluster" +) + +const ( + defaultPrimaryDatabaseURL = "postgres://postgres:postgres@localhost:55432/postgres?sslmode=disable&target_session_attrs=read-write" + defaultReplicaOneURL = "postgres://postgres:postgres@localhost:55433/postgres?sslmode=disable&target_session_attrs=read-only" + defaultReplicaTwoURL = "postgres://postgres:postgres@localhost:55434/postgres?sslmode=disable&target_session_attrs=read-only" +) + +func openCluster(ctx context.Context) (*cluster.Cluster, error) { + primary, err := openPool( + ctx, + environment("XPG_PRIMARY_DATABASE_URL", defaultPrimaryDatabaseURL), + "cluster.primary", + "primary", + ) + if err != nil { + return nil, fmt.Errorf("open primary pool: %w", err) + } + + replicaOne, err := openPool( + ctx, + environment("XPG_REPLICA_ONE_DATABASE_URL", defaultReplicaOneURL), + "cluster.replica-one", + "replica", + ) + if err != nil { + primary.Close() + + return nil, fmt.Errorf("open replica-one pool: %w", err) + } + + replicaTwo, err := openPool( + ctx, + environment("XPG_REPLICA_TWO_DATABASE_URL", defaultReplicaTwoURL), + "cluster.replica-two", + "replica", + ) + if err != nil { + replicaOne.Close() + primary.Close() + + return nil, fmt.Errorf("open replica-two pool: %w", err) + } + + dbCluster, err := cluster.New(cluster.Config{ + ID: "cluster-example", + Primary: primary, + Replicas: []*xpg.Pool{replicaOne, replicaTwo}, + }) + if err != nil { + replicaTwo.Close() + replicaOne.Close() + primary.Close() + + return nil, fmt.Errorf("create cluster: %w", err) + } + + return dbCluster, nil +} + +func openPool(ctx context.Context, databaseURL, name, role string) (*xpg.Pool, error) { + pool, err := xpg.Open( + ctx, + databaseURL, + xpg.WithName(name), + xpg.WithLabel("xpg.pool.role", role), + ) + if err != nil { + return nil, err + } + + if err := pool.Ping(ctx); err != nil { + pool.Close() + + return nil, fmt.Errorf("ping pool: %w", err) + } + + return pool, nil +} + +func environment(key, fallback string) string { + if value := os.Getenv(key); value != "" { + return value + } + + return fallback +} diff --git a/examples/cluster/sql/primary.sql b/examples/cluster/sql/primary.sql index d53fe6b..d265322 100644 --- a/examples/cluster/sql/primary.sql +++ b/examples/cluster/sql/primary.sql @@ -1,3 +1,8 @@ +ALTER DATABASE postgres +SET default_transaction_read_only = off; + +DROP SCHEMA IF EXISTS xpg_cluster_example CASCADE; + CREATE SCHEMA xpg_cluster_example; CREATE TABLE xpg_cluster_example.node_info ( @@ -13,8 +18,3 @@ VALUES ( 'primary', 'primary' ); - -CREATE TABLE xpg_cluster_example.primary_writes ( - id bigint GENERATED ALWAYS AS IDENTITY PRIMARY KEY, - created_at timestamptz NOT NULL DEFAULT now() -); diff --git a/examples/cluster/sql/replica-one.sql b/examples/cluster/sql/replica-one.sql index 0973062..c502e73 100644 --- a/examples/cluster/sql/replica-one.sql +++ b/examples/cluster/sql/replica-one.sql @@ -1,3 +1,7 @@ +SET default_transaction_read_only = off; + +DROP SCHEMA IF EXISTS xpg_cluster_example CASCADE; + CREATE SCHEMA xpg_cluster_example; CREATE TABLE xpg_cluster_example.node_info ( diff --git a/examples/cluster/sql/replica-two.sql b/examples/cluster/sql/replica-two.sql index 968eb46..d60ed58 100644 --- a/examples/cluster/sql/replica-two.sql +++ b/examples/cluster/sql/replica-two.sql @@ -1,3 +1,7 @@ +SET default_transaction_read_only = off; + +DROP SCHEMA IF EXISTS xpg_cluster_example CASCADE; + CREATE SCHEMA xpg_cluster_example; CREATE TABLE xpg_cluster_example.node_info ( From 1bdf20e59ffd24f22fc84867969dae9ec3572909 Mon Sep 17 00:00:00 2001 From: mkbeh Date: Tue, 25 Aug 2026 15:39:39 +0300 Subject: [PATCH 31/41] refactor: simplify shard example --- examples/shard/README.md | 125 +++++++++---------- examples/shard/docker-compose.yml | 52 +++----- examples/shard/go.mod | 5 +- examples/shard/main.go | 77 +++++++++--- examples/shard/operations.go | 160 ------------------------ examples/shard/routing.go | 133 -------------------- examples/shard/setup.go | 162 +++++++------------------ examples/shard/sql/schema.sql | 8 ++ examples/shard/sql/shard-a-primary.sql | 2 - examples/shard/sql/shard-b-primary.sql | 2 - examples/shard/sql/shard-b-replica.sql | 5 - examples/shard/sql/shard.sql | 21 ---- 12 files changed, 188 insertions(+), 564 deletions(-) delete mode 100644 examples/shard/operations.go delete mode 100644 examples/shard/routing.go create mode 100644 examples/shard/sql/schema.sql delete mode 100644 examples/shard/sql/shard-a-primary.sql delete mode 100644 examples/shard/sql/shard-b-primary.sql delete mode 100644 examples/shard/sql/shard-b-replica.sql delete mode 100644 examples/shard/sql/shard.sql diff --git a/examples/shard/README.md b/examples/shard/README.md index cd06850..4a890d1 100644 --- a/examples/shard/README.md +++ b/examples/shard/README.md @@ -1,84 +1,66 @@ # Sharding -This example routes typed application keys across two logical PostgreSQL shards. Resolvers are bound to an immutable -`shard.Topology` and return a `shard.Shard`, which delegates database operations to its `cluster.Cluster`. +This example shows how to distribute application data across PostgreSQL shards: -```text - Topology []Shard - │ -application key ─── Resolver - │ - Shard - │ - Cluster - │ - primary / replicas -``` - -**This example demonstrates:** - -* Building a topology from primary-only and primary/replica clusters -* Routing `uint64` keys through bounded numeric ranges -* Running shard-local primary and read-only replica transactions -* Checking key colocation and grouping keys by shard -* Reading reference-table copies across shards with bounded fan-out - -> [!NOTE] -> The `shard-b` replica is an independent read-only PostgreSQL instance used to -> demonstrate routing. The local setup does not configure streaming replication. - -## Configuration +* Route user IDs with a range-based shard resolver +* Write records to the resolved shard +* Group keys by shard for efficient batch processing -The example uses the following connection strings by default: +The example uses two primary-only shards: ```text -XPG_SHARD_A_PRIMARY_DATABASE_URL=postgres://postgres:postgres@localhost:56431/postgres?sslmode=disable&target_session_attrs=read-write -XPG_SHARD_B_PRIMARY_DATABASE_URL=postgres://postgres:postgres@localhost:56432/postgres?sslmode=disable&target_session_attrs=read-write -XPG_SHARD_B_REPLICA_DATABASE_URL=postgres://postgres:postgres@localhost:56433/postgres?sslmode=disable&target_session_attrs=read-only -``` - -Set the corresponding environment variables to use different PostgreSQL endpoints: - -```shell -export XPG_SHARD_A_PRIMARY_DATABASE_URL='postgres://user:password@shard-a.example.com:5432/database?sslmode=disable&target_session_attrs=read-write' -export XPG_SHARD_B_PRIMARY_DATABASE_URL='postgres://user:password@shard-b-primary.example.com:5432/database?sslmode=disable&target_session_attrs=read-write' -export XPG_SHARD_B_REPLICA_DATABASE_URL='postgres://user:password@shard-b-replica.example.com:5432/database?sslmode=disable&target_session_attrs=read-only' +[0, 100) -> shard-a +[100, 200) -> shard-b ``` ## Local setup -Start all PostgreSQL endpoints and Adminer from the repository root: +From this directory, start both PostgreSQL shards and Adminer: ```shell -docker compose -f examples/shard/docker-compose.yml --profile tools up -d +docker compose up -d ``` -Or from this example directory: +Apply the example schema to both shards: ```shell -docker compose --profile tools up -d +psql 'postgres://postgres:postgres@localhost:56431/postgres?sslmode=disable' \ + < sql/schema.sql + +psql 'postgres://postgres:postgres@localhost:56432/postgres?sslmode=disable' \ + < sql/schema.sql ``` -Services are available at: +The services are available at: ```text -Shard A primary: localhost:56431 -Shard B primary: localhost:56432 -Shard B replica: localhost:56433 -Adminer: http://localhost:58082 +Shard A: localhost:56431 +Shard B: localhost:56432 +Adminer: http://localhost:8080 ``` -Sign in to Adminer with: +To inspect a shard in Adminer, sign in with: ```text System: PostgreSQL -Server: postgres-shard-a-primary +Server: postgres-shard-a Username: postgres Password: postgres Database: postgres ``` -Use `postgres-shard-b-primary` or `postgres-shard-b-replica` in the **Server** field to inspect another shard endpoint. +Use `postgres-shard-b` in the **Server** field to inspect the second shard. + +## Configuration + +By default, the example connects to: + +```text +Shard A: postgres://postgres:postgres@localhost:56431/postgres?sslmode=disable +Shard B: postgres://postgres:postgres@localhost:56432/postgres?sslmode=disable +``` + +To use other PostgreSQL endpoints, set `XPG_SHARD_A_DATABASE_URL` and `XPG_SHARD_B_DATABASE_URL`. ## Run @@ -97,37 +79,38 @@ go run ./examples/shard ## Expected output ```text -range routing and primary transactions: -- user_id=42 shard=shard-a primary_pool=shard.shard-a.primary -- user_id=142 shard=shard-b primary_pool=shard.shard-b.primary +range routing: +- user_id=42 shard=shard-a pool=shard.shard-a.primary +- user_id=142 shard=shard-b pool=shard.shard-b.primary -grouping and colocation: +grouping: - shard=shard-b user_ids=[142 143] - shard=shard-a user_ids=[42 43] -- colocated shard=shard-a -- cross-shard SameShard returns ErrShardMismatch=true +``` -replica routing and read-only transaction: -- user_id=142 shard=shard-b read_pool=shard.shard-b.replica read_node=shard-b-replica tx_node=shard-b-replica role=replica read_only=on +`GroupByShard` preserves the order in which shards first appear in the input and the relative order of keys within each +group. -reference table copies: -- shard=shard-a countries=2 -- shard=shard-b countries=2 -``` +## Cleanup + +To remove the example schema and data: -## Stop services +```shell +psql 'postgres://postgres:postgres@localhost:56431/postgres?sslmode=disable' \ + -c 'DROP SCHEMA IF EXISTS xpg_shard_example CASCADE;' + +psql 'postgres://postgres:postgres@localhost:56432/postgres?sslmode=disable' \ + -c 'DROP SCHEMA IF EXISTS xpg_shard_example CASCADE;' +``` -From the repository root: +Stop the local services: ```shell -docker compose \ - -f examples/shard/docker-compose.yml \ - --profile tools \ - down --remove-orphans -v +docker compose down ``` -Or from this example directory: +To also remove both PostgreSQL data volumes: ```shell -docker compose --profile tools down --remove-orphans -v +docker compose down -v ``` diff --git a/examples/shard/docker-compose.yml b/examples/shard/docker-compose.yml index d379c3e..ae82f0a 100644 --- a/examples/shard/docker-compose.yml +++ b/examples/shard/docker-compose.yml @@ -1,5 +1,5 @@ services: - postgres-shard-a-primary: + postgres-shard-a: image: postgres:18-alpine environment: POSTGRES_DB: postgres @@ -8,17 +8,15 @@ services: ports: - "56431:5432" volumes: - - shard-a-primary-data:/var/lib/postgresql - - ./sql/shard.sql:/docker-entrypoint-initdb.d/10-shard.sql:ro - - ./sql/shard-a-primary.sql:/docker-entrypoint-initdb.d/20-node.sql:ro - healthcheck: &postgres-healthcheck + - shard-a-data:/var/lib/postgresql + healthcheck: test: "pg_isready -U $${POSTGRES_USER} -d $${POSTGRES_DB}" interval: 2s timeout: 5s retries: 10 start_period: 5s - postgres-shard-b-primary: + postgres-shard-b: image: postgres:18-alpine environment: POSTGRES_DB: postgres @@ -27,42 +25,26 @@ services: ports: - "56432:5432" volumes: - - shard-b-primary-data:/var/lib/postgresql - - ./sql/shard.sql:/docker-entrypoint-initdb.d/10-shard.sql:ro - - ./sql/shard-b-primary.sql:/docker-entrypoint-initdb.d/20-node.sql:ro - healthcheck: *postgres-healthcheck - - postgres-shard-b-replica: - image: postgres:18-alpine - environment: - POSTGRES_DB: postgres - POSTGRES_USER: postgres - POSTGRES_PASSWORD: postgres - ports: - - "56433:5432" - volumes: - - shard-b-replica-data:/var/lib/postgresql - - ./sql/shard.sql:/docker-entrypoint-initdb.d/10-shard.sql:ro - - ./sql/shard-b-replica.sql:/docker-entrypoint-initdb.d/20-node.sql:ro - healthcheck: *postgres-healthcheck + - shard-b-data:/var/lib/postgresql + healthcheck: + test: "pg_isready -U $${POSTGRES_USER} -d $${POSTGRES_DB}" + interval: 2s + timeout: 5s + retries: 10 + start_period: 5s adminer: image: adminer:standalone - profiles: - - tools environment: - ADMINER_DEFAULT_SERVER: postgres-shard-a-primary + ADMINER_DEFAULT_SERVER: postgres-shard-a ports: - - "58082:8080" + - "8080:8080" depends_on: - postgres-shard-a-primary: - condition: service_healthy - postgres-shard-b-primary: + postgres-shard-a: condition: service_healthy - postgres-shard-b-replica: + postgres-shard-b: condition: service_healthy volumes: - shard-a-primary-data: - shard-b-primary-data: - shard-b-replica-data: + shard-a-data: + shard-b-data: diff --git a/examples/shard/go.mod b/examples/shard/go.mod index e9cc1be..4a05837 100644 --- a/examples/shard/go.mod +++ b/examples/shard/go.mod @@ -2,7 +2,4 @@ module github.com/mkbeh/xpg/examples/shard go 1.27 -require ( - github.com/jackc/pgx/v5 v5.10.0 - github.com/mkbeh/xpg v0.2.0 -) +require github.com/mkbeh/xpg v0.2.0 diff --git a/examples/shard/main.go b/examples/shard/main.go index 8b7b00a..9f45891 100644 --- a/examples/shard/main.go +++ b/examples/shard/main.go @@ -5,15 +5,21 @@ import ( "fmt" "log" + "github.com/mkbeh/xpg/shard" "github.com/mkbeh/xpg/shard/resolver" ) const ( - userIDRangeStart uint64 = 0 - userIDBoundary uint64 = 100 - userIDRangeEnd uint64 = 200 + shardARangeStart uint64 = 0 + shardBoundary uint64 = 100 + shardBRangeEnd uint64 = 200 ) +type user struct { + ID uint64 + Name string +} + func main() { if err := run(context.Background()); err != nil { log.Fatal(err) @@ -31,13 +37,13 @@ func run(ctx context.Context) error { topology, []resolver.Range[uint64]{ { - Start: userIDRangeStart, - End: userIDBoundary, + Start: shardARangeStart, + End: shardBoundary, ShardID: shardAID, }, { - Start: userIDBoundary, - End: userIDRangeEnd, + Start: shardBoundary, + End: shardBRangeEnd, ShardID: shardBID, }, }, @@ -46,20 +52,61 @@ func run(ctx context.Context) error { return fmt.Errorf("create user resolver: %w", err) } - if err := showRangeRouting(ctx, userResolver); err != nil { - return err + users := []user{ + {ID: 42, Name: "alice"}, + {ID: 142, Name: "bob"}, } - if err := showGrouping(userResolver); err != nil { - return err + fmt.Println("range routing:") + + for _, current := range users { + resolved, err := userResolver.Resolve(current.ID) + if err != nil { + return fmt.Errorf("resolve user %d: %w", current.ID, err) + } + + primary := resolved.Primary() + if primary == nil { + return fmt.Errorf("shard %q has no primary", resolved.ID()) + } + + if _, err := primary.Exec( + ctx, + `INSERT INTO xpg_shard_example.users (id, name) + VALUES ($1, $2) + ON CONFLICT (id) DO UPDATE + SET name = EXCLUDED.name`, + current.ID, + current.Name, + ); err != nil { + return fmt.Errorf("write user %d: %w", current.ID, err) + } + + fmt.Printf( + "- user_id=%d shard=%s pool=%s\n", + current.ID, + resolved.ID(), + primary.Name(), + ) } - if err := showReplicaRead(ctx, userResolver, shardBUserID); err != nil { - return err + groups, err := shard.GroupByShard( + userResolver, + []uint64{142, 42, 143, 43}, + ) + if err != nil { + return fmt.Errorf("group user IDs: %w", err) } - if err := showReferenceTables(ctx, topology); err != nil { - return err + fmt.Println() + fmt.Println("grouping:") + + for _, group := range groups { + fmt.Printf( + "- shard=%s user_ids=%v\n", + group.Shard.ID(), + group.Keys, + ) } return nil diff --git a/examples/shard/operations.go b/examples/shard/operations.go deleted file mode 100644 index 3e2db85..0000000 --- a/examples/shard/operations.go +++ /dev/null @@ -1,160 +0,0 @@ -package main - -import ( - "context" - "fmt" - "sync" - - "github.com/jackc/pgx/v5" - "github.com/mkbeh/xpg/cluster" - "github.com/mkbeh/xpg/shard" -) - -const referenceTableConcurrency = 2 - -type nodeInfo struct { - name string - role string -} - -type rowQuerier interface { - QueryRow(context.Context, string, ...any) pgx.Row -} - -func showReplicaRead( - ctx context.Context, - userResolver shard.Resolver[uint64], - key uint64, -) error { - resolved, err := userResolver.Resolve(key) - if err != nil { - return fmt.Errorf("resolve replica key %d: %w", key, err) - } - - pool, err := resolved.ReadPool( - ctx, - cluster.ReadReplicaRequired, - ) - if err != nil { - return fmt.Errorf("resolve shard replica: %w", err) - } - - node, err := loadNode(ctx, pool) - if err != nil { - return fmt.Errorf("read replica node: %w", err) - } - - var ( - transactionNode nodeInfo - readOnly string - ) - if err := resolved.InReadTx( - ctx, - cluster.ReadReplicaRequired, - cluster.ReadTxOptions{}, - func(ctx context.Context, tx pgx.Tx) error { - var err error - - transactionNode, err = loadNode(ctx, tx) - if err != nil { - return err - } - - return tx.QueryRow( - ctx, - "SHOW transaction_read_only", - ).Scan(&readOnly) - }, - ); err != nil { - return fmt.Errorf("run shard read transaction: %w", err) - } - - fmt.Println() - fmt.Println("replica routing and read-only transaction:") - fmt.Printf( - "- user_id=%d shard=%s read_pool=%s read_node=%s tx_node=%s role=%s read_only=%s\n", - key, - resolved.ID(), - pool.Name(), - node.name, - transactionNode.name, - node.role, - readOnly, - ) - - return nil -} - -func showReferenceTables( - ctx context.Context, - topology *shard.Topology, -) error { - counts := make(map[shard.ID]int, topology.Len()) - var countsMu sync.Mutex - - // Fan out with bounded concurrency while preserving registration-order - // results for deterministic reporting. - results, err := topology.ForEachShard( - ctx, - referenceTableConcurrency, - func(ctx context.Context, resolved shard.Shard) error { - primary := resolved.Primary() - if primary == nil { - return fmt.Errorf("shard %q has no primary", resolved.ID()) - } - - var count int - if err := primary.QueryRow( - ctx, - `SELECT count(*) - FROM xpg_shard_example.countries`, - ).Scan(&count); err != nil { - return err - } - - countsMu.Lock() - counts[resolved.ID()] = count - countsMu.Unlock() - - return nil - }, - ) - if err != nil { - return fmt.Errorf("schedule reference-table reads: %w", err) - } - - if err := results.Err(); err != nil { - return fmt.Errorf("read reference tables: %w", err) - } - - fmt.Println() - fmt.Println("reference table copies:") - for index := 0; index < topology.Len(); index++ { - resolved := topology.At(index) - fmt.Printf( - "- shard=%s countries=%d\n", - resolved.ID(), - counts[resolved.ID()], - ) - } - - return nil -} - -func loadNode( - ctx context.Context, - db rowQuerier, -) (nodeInfo, error) { - var node nodeInfo - - err := db.QueryRow( - ctx, - `SELECT node_name, node_role - FROM xpg_shard_example.node_info`, - ).Scan(&node.name, &node.role) - if err != nil { - return nodeInfo{}, err - } - - return node, nil -} diff --git a/examples/shard/routing.go b/examples/shard/routing.go deleted file mode 100644 index 6079a8d..0000000 --- a/examples/shard/routing.go +++ /dev/null @@ -1,133 +0,0 @@ -package main - -import ( - "context" - "errors" - "fmt" - - "github.com/jackc/pgx/v5" - "github.com/mkbeh/xpg/shard" -) - -const ( - shardAUserID uint64 = 42 - shardASecondUserID uint64 = 43 - shardBUserID uint64 = 142 - shardBSecondUserID uint64 = 143 -) - -func showRangeRouting( - ctx context.Context, - userResolver shard.Resolver[uint64], -) error { - fmt.Println("range routing and primary transactions:") - - for _, userID := range []uint64{ - shardAUserID, - shardBUserID, - } { - resolved, err := userResolver.Resolve(userID) - if err != nil { - return fmt.Errorf("resolve user %d: %w", userID, err) - } - - primary := resolved.Primary() - if primary == nil { - return fmt.Errorf("shard %q has no primary", resolved.ID()) - } - - name := fmt.Sprintf("user-%d", userID) - - if err := resolved.InPrimaryTx( - ctx, - pgx.TxOptions{}, - func(ctx context.Context, tx pgx.Tx) error { - _, err := tx.Exec( - ctx, - `INSERT INTO xpg_shard_example.users (id, name) - VALUES ($1, $2) - ON CONFLICT (id) DO UPDATE - SET name = EXCLUDED.name`, - userID, - name, - ) - - return err - }, - ); err != nil { - return fmt.Errorf("write user %d: %w", userID, err) - } - - fmt.Printf( - "- user_id=%d shard=%s primary_pool=%s\n", - userID, - resolved.ID(), - primary.Name(), - ) - } - - return nil -} - -func showGrouping( - userResolver shard.Resolver[uint64], -) error { - // The first occurrence belongs to shard-b, so GroupByShard returns - // shard-b before shard-a. - keys := []uint64{ - shardBUserID, - shardAUserID, - shardBSecondUserID, - shardASecondUserID, - } - - groups, err := shard.GroupByShard( - userResolver, - keys, - ) - if err != nil { - return fmt.Errorf("group users: %w", err) - } - - colocated, err := shard.SameShard( - userResolver, - shardAUserID, - shardASecondUserID, - ) - if err != nil { - return fmt.Errorf("verify colocated users: %w", err) - } - - _, mismatchErr := shard.SameShard( - userResolver, - shardAUserID, - shardBUserID, - ) - if !errors.Is(mismatchErr, shard.ErrShardMismatch) { - return fmt.Errorf( - "verify cross-shard users: expected ErrShardMismatch, got %v", - mismatchErr, - ) - } - - fmt.Println() - fmt.Println("grouping and colocation:") - - for _, group := range groups { - fmt.Printf( - "- shard=%s user_ids=%v\n", - group.Shard.ID(), - group.Keys, - ) - } - - fmt.Printf( - "- colocated shard=%s\n", - colocated.ID(), - ) - fmt.Println( - "- cross-shard SameShard returns ErrShardMismatch=true", - ) - - return nil -} diff --git a/examples/shard/setup.go b/examples/shard/setup.go index b72cf3f..2a43d9d 100644 --- a/examples/shard/setup.go +++ b/examples/shard/setup.go @@ -11,155 +11,85 @@ import ( ) const ( - defaultShardAPrimaryDatabaseURL = "postgres://postgres:postgres@localhost:56431/postgres?sslmode=disable&target_session_attrs=read-write" - defaultShardBPrimaryDatabaseURL = "postgres://postgres:postgres@localhost:56432/postgres?sslmode=disable&target_session_attrs=read-write" - defaultShardBReplicaDatabaseURL = "postgres://postgres:postgres@localhost:56433/postgres?sslmode=disable&target_session_attrs=read-only" + defaultShardADatabaseURL = "postgres://postgres:postgres@localhost:56431/postgres?sslmode=disable" + defaultShardBDatabaseURL = "postgres://postgres:postgres@localhost:56432/postgres?sslmode=disable" shardAID shard.ID = "shard-a" shardBID shard.ID = "shard-b" ) -type poolSpec struct { - databaseURL string - name string -} - -type shardSpec struct { - id shard.ID - primary poolSpec - replicas []poolSpec -} - func openTopology(ctx context.Context) (*shard.Topology, error) { - specs := []shardSpec{ - { - id: shardAID, - primary: poolSpec{ - databaseURL: environment( - "XPG_SHARD_A_PRIMARY_DATABASE_URL", - defaultShardAPrimaryDatabaseURL, - ), - name: "shard.shard-a.primary", - }, - }, - { - id: shardBID, - primary: poolSpec{ - databaseURL: environment( - "XPG_SHARD_B_PRIMARY_DATABASE_URL", - defaultShardBPrimaryDatabaseURL, - ), - name: "shard.shard-b.primary", - }, - replicas: []poolSpec{ - { - databaseURL: environment( - "XPG_SHARD_B_REPLICA_DATABASE_URL", - defaultShardBReplicaDatabaseURL, - ), - name: "shard.shard-b.replica", - }, - }, - }, - } - - clusters := make([]*cluster.Cluster, 0, len(specs)) - configs := make([]shard.Config, 0, len(specs)) - - for _, spec := range specs { - dbCluster, err := openCluster(ctx, spec) - if err != nil { - closeClusters(clusters) - - return nil, fmt.Errorf("open %s cluster: %w", spec.id, err) - } - - clusters = append(clusters, dbCluster) - configs = append(configs, shard.Config{Cluster: dbCluster}) - } - - // Topology takes ownership of the clusters only after successful creation. - topology, err := shard.NewTopology(configs) + shardA, err := openShard( + ctx, + shardAID, + "shard.shard-a.primary", + environment( + "XPG_SHARD_A_DATABASE_URL", + defaultShardADatabaseURL, + ), + ) if err != nil { - closeClusters(clusters) - - return nil, fmt.Errorf("create topology: %w", err) + return nil, fmt.Errorf("open shard-a: %w", err) } - return topology, nil -} - -func openCluster(ctx context.Context, spec shardSpec) (*cluster.Cluster, error) { - primary, err := openPool(ctx, spec.primary) + shardB, err := openShard( + ctx, + shardBID, + "shard.shard-b.primary", + environment( + "XPG_SHARD_B_DATABASE_URL", + defaultShardBDatabaseURL, + ), + ) if err != nil { - return nil, fmt.Errorf("open primary %s: %w", spec.primary.name, err) - } - - pools := []*xpg.Pool{primary} - replicas := make([]*xpg.Pool, 0, len(spec.replicas)) - - for _, replicaSpec := range spec.replicas { - replica, err := openPool(ctx, replicaSpec) - if err != nil { - closePools(pools) - - return nil, fmt.Errorf("open replica %s: %w", replicaSpec.name, err) - } - - pools = append(pools, replica) - replicas = append(replicas, replica) - } - - config := cluster.Config{ - ID: spec.id, - Primary: primary, - Replicas: replicas, - } + shardA.Close() - if len(replicas) > 0 { - // Keep replica selection explicit in the runnable example. - config.Selector = cluster.RoundRobinSelector() + return nil, fmt.Errorf("open shard-b: %w", err) } - dbCluster, err := cluster.New(config) + topology, err := shard.NewTopology([]shard.Config{ + {Cluster: shardA}, + {Cluster: shardB}, + }) if err != nil { - closePools(pools) + shardB.Close() + shardA.Close() - return nil, fmt.Errorf("create cluster: %w", err) + return nil, fmt.Errorf("create topology: %w", err) } - return dbCluster, nil + return topology, nil } -func openPool(ctx context.Context, spec poolSpec) (*xpg.Pool, error) { +func openShard(ctx context.Context, id shard.ID, name, databaseURL string) (*cluster.Cluster, error) { pool, err := xpg.Open( ctx, - spec.databaseURL, - xpg.WithName(spec.name), + databaseURL, + xpg.WithName(name), + xpg.WithLabel("xpg.shard.id", string(id)), + xpg.WithLabel("xpg.pool.role", "primary"), ) if err != nil { - return nil, fmt.Errorf("open %s: %w", spec.name, err) + return nil, fmt.Errorf("open pool: %w", err) } if err := pool.Ping(ctx); err != nil { pool.Close() - return nil, fmt.Errorf("ping %s: %w", spec.name, err) + return nil, fmt.Errorf("ping pool: %w", err) } - return pool, nil -} + shardCluster, err := cluster.New(cluster.Config{ + ID: id, + Primary: pool, + }) + if err != nil { + pool.Close() -func closePools(pools []*xpg.Pool) { - for index := len(pools) - 1; index >= 0; index-- { - pools[index].Close() + return nil, fmt.Errorf("create cluster: %w", err) } -} -func closeClusters(clusters []*cluster.Cluster) { - for index := len(clusters) - 1; index >= 0; index-- { - clusters[index].Close() - } + return shardCluster, nil } func environment(name, fallback string) string { diff --git a/examples/shard/sql/schema.sql b/examples/shard/sql/schema.sql new file mode 100644 index 0000000..27c3abd --- /dev/null +++ b/examples/shard/sql/schema.sql @@ -0,0 +1,8 @@ +DROP SCHEMA IF EXISTS xpg_shard_example CASCADE; + +CREATE SCHEMA xpg_shard_example; + +CREATE TABLE xpg_shard_example.users ( + id bigint PRIMARY KEY, + name text NOT NULL +); diff --git a/examples/shard/sql/shard-a-primary.sql b/examples/shard/sql/shard-a-primary.sql deleted file mode 100644 index 7baca90..0000000 --- a/examples/shard/sql/shard-a-primary.sql +++ /dev/null @@ -1,2 +0,0 @@ -INSERT INTO xpg_shard_example.node_info (node_name, node_role) -VALUES ('shard-a-primary', 'primary'); diff --git a/examples/shard/sql/shard-b-primary.sql b/examples/shard/sql/shard-b-primary.sql deleted file mode 100644 index 252c162..0000000 --- a/examples/shard/sql/shard-b-primary.sql +++ /dev/null @@ -1,2 +0,0 @@ -INSERT INTO xpg_shard_example.node_info (node_name, node_role) -VALUES ('shard-b-primary', 'primary'); diff --git a/examples/shard/sql/shard-b-replica.sql b/examples/shard/sql/shard-b-replica.sql deleted file mode 100644 index 14a8bc2..0000000 --- a/examples/shard/sql/shard-b-replica.sql +++ /dev/null @@ -1,5 +0,0 @@ -INSERT INTO xpg_shard_example.node_info (node_name, node_role) -VALUES ('shard-b-replica', 'replica'); - -ALTER DATABASE postgres -SET default_transaction_read_only = on; diff --git a/examples/shard/sql/shard.sql b/examples/shard/sql/shard.sql deleted file mode 100644 index ddc8b69..0000000 --- a/examples/shard/sql/shard.sql +++ /dev/null @@ -1,21 +0,0 @@ -CREATE SCHEMA xpg_shard_example; - -CREATE TABLE xpg_shard_example.node_info ( - node_name text PRIMARY KEY, - node_role text NOT NULL -); - -CREATE TABLE xpg_shard_example.users ( - id bigint PRIMARY KEY, - name text NOT NULL -); - -CREATE TABLE xpg_shard_example.countries ( - code text PRIMARY KEY, - name text NOT NULL -); - -INSERT INTO xpg_shard_example.countries (code, name) -VALUES - ('NL', 'Netherlands'), - ('DE', 'Germany'); From 402ad5c28928f59c1ae979ab993dc1aec4ceeaf4 Mon Sep 17 00:00:00 2001 From: mkbeh Date: Tue, 25 Aug 2026 15:58:09 +0300 Subject: [PATCH 32/41] feat: add geographic sharding example --- examples/shard_geo/README.md | 114 +++++++++++++++++++ examples/shard_geo/docker-compose.yml | 50 +++++++++ examples/shard_geo/go.mod | 8 ++ examples/shard_geo/main.go | 151 ++++++++++++++++++++++++++ examples/shard_geo/resolver.go | 50 +++++++++ examples/shard_geo/setup.go | 110 +++++++++++++++++++ examples/shard_geo/sql/schema.sql | 10 ++ 7 files changed, 493 insertions(+) create mode 100644 examples/shard_geo/README.md create mode 100644 examples/shard_geo/docker-compose.yml create mode 100644 examples/shard_geo/go.mod create mode 100644 examples/shard_geo/main.go create mode 100644 examples/shard_geo/resolver.go create mode 100644 examples/shard_geo/setup.go create mode 100644 examples/shard_geo/sql/schema.sql diff --git a/examples/shard_geo/README.md b/examples/shard_geo/README.md new file mode 100644 index 0000000..f013fb4 --- /dev/null +++ b/examples/shard_geo/README.md @@ -0,0 +1,114 @@ +# Geographic sharding + +This example shows how to route tenants by region across PostgreSQL shards: + +* Associate each shard with a region +* Build a custom resolver from shard metadata +* Route tenant operations to the matching shard +* Handle unsupported regions with `shard.ErrNoShard` + +The topology contains two primary-only shards: + +```text +eu -> shard-eu +us -> shard-us +``` + +## Local setup + +From this directory, start both PostgreSQL shards and Adminer: + +```shell +docker compose up -d +``` + +Apply the example schema to both shards: + +```shell +psql 'postgres://postgres:postgres@localhost:57431/postgres?sslmode=disable' \ + < sql/schema.sql + +psql 'postgres://postgres:postgres@localhost:57432/postgres?sslmode=disable' \ + < sql/schema.sql +``` + +The services are available at: + +```text +EU shard: localhost:57431 +US shard: localhost:57432 +Adminer: http://localhost:8080 +``` + +To inspect a shard in Adminer, sign in with: + +```text +System: PostgreSQL +Server: postgres-shard-eu +Username: postgres +Password: postgres +Database: postgres +``` + +Use `postgres-shard-us` in the **Server** field to inspect the US shard. + +## Configuration + +By default, the example connects to: + +```text +EU shard: postgres://postgres:postgres@localhost:57431/postgres?sslmode=disable +US shard: postgres://postgres:postgres@localhost:57432/postgres?sslmode=disable +``` + +To use other PostgreSQL endpoints, set `XPG_SHARD_EU_DATABASE_URL` and `XPG_SHARD_US_DATABASE_URL`. + +## Run + +From this directory: + +```shell +go run . +``` + +Or from the repository root: + +```shell +go run ./examples/shard_geo +``` + +## Expected output + +```text +geo routing: +- tenant=tenant-42 region=eu shard=shard-eu name=Alice +- tenant=tenant-77 region=us shard=shard-us name=Bob +unsupported region: no shard +``` + +The region index is built once from immutable shard metadata. Each subsequent route is a local map lookup followed by +the normal `shard.Shard` execution API. + +## Cleanup + +To remove the example schema and data: + +```shell +psql 'postgres://postgres:postgres@localhost:57431/postgres?sslmode=disable' \ + -c 'DROP SCHEMA IF EXISTS xpg_shard_geo_example CASCADE;' + +psql 'postgres://postgres:postgres@localhost:57432/postgres?sslmode=disable' \ + -c 'DROP SCHEMA IF EXISTS xpg_shard_geo_example CASCADE;' +``` + +Stop the local services: + +```shell +docker compose down +``` + +To also remove both PostgreSQL data volumes: + +```shell +docker compose down -v +``` diff --git a/examples/shard_geo/docker-compose.yml b/examples/shard_geo/docker-compose.yml new file mode 100644 index 0000000..10c7cb0 --- /dev/null +++ b/examples/shard_geo/docker-compose.yml @@ -0,0 +1,50 @@ +services: + postgres-shard-eu: + image: postgres:18-alpine + environment: + POSTGRES_DB: postgres + POSTGRES_USER: postgres + POSTGRES_PASSWORD: postgres + ports: + - "57431:5432" + volumes: + - shard-eu-data:/var/lib/postgresql + healthcheck: + test: "pg_isready -U $${POSTGRES_USER} -d $${POSTGRES_DB}" + interval: 2s + timeout: 5s + retries: 10 + start_period: 5s + + postgres-shard-us: + image: postgres:18-alpine + environment: + POSTGRES_DB: postgres + POSTGRES_USER: postgres + POSTGRES_PASSWORD: postgres + ports: + - "57432:5432" + volumes: + - shard-us-data:/var/lib/postgresql + healthcheck: + test: "pg_isready -U $${POSTGRES_USER} -d $${POSTGRES_DB}" + interval: 2s + timeout: 5s + retries: 10 + start_period: 5s + + adminer: + image: adminer:standalone + environment: + ADMINER_DEFAULT_SERVER: postgres-shard-eu + ports: + - "8080:8080" + depends_on: + postgres-shard-eu: + condition: service_healthy + postgres-shard-us: + condition: service_healthy + +volumes: + shard-eu-data: + shard-us-data: diff --git a/examples/shard_geo/go.mod b/examples/shard_geo/go.mod new file mode 100644 index 0000000..320f239 --- /dev/null +++ b/examples/shard_geo/go.mod @@ -0,0 +1,8 @@ +module github.com/mkbeh/xpg/examples/shard_geo + +go 1.27 + +require ( + github.com/jackc/pgx/v5 v5.10.0 + github.com/mkbeh/xpg v0.2.0 +) diff --git a/examples/shard_geo/main.go b/examples/shard_geo/main.go new file mode 100644 index 0000000..f2d75fd --- /dev/null +++ b/examples/shard_geo/main.go @@ -0,0 +1,151 @@ +package main + +import ( + "context" + "errors" + "fmt" + "log" + + "github.com/jackc/pgx/v5" + "github.com/mkbeh/xpg/shard" +) + +type tenantKey struct { + TenantID string + Region string +} + +type tenant struct { + ID string + Name string + Region string +} + +func main() { + if err := run(context.Background()); err != nil { + log.Fatal(err) + } +} + +func run(ctx context.Context) error { + topology, err := openTopology(ctx) + if err != nil { + return err + } + defer topology.Close() + + tenantResolver, err := newTenantResolver(topology) + if err != nil { + return fmt.Errorf("create tenant resolver: %w", err) + } + + tenants := []tenant{ + { + ID: "tenant-42", + Name: "Alice", + Region: "eu", + }, + { + ID: "tenant-77", + Name: "Bob", + Region: "us", + }, + } + + fmt.Println("geo routing:") + + for _, current := range tenants { + key := tenantKey{ + TenantID: current.ID, + Region: current.Region, + } + + resolved, err := tenantResolver.Resolve(key) + if err != nil { + return fmt.Errorf("resolve tenant %q: %w", current.ID, err) + } + + if err := upsertTenant(ctx, resolved, current); err != nil { + return fmt.Errorf("upsert tenant %q: %w", current.ID, err) + } + + stored, err := loadTenant(ctx, resolved, current.ID) + if err != nil { + return fmt.Errorf("load tenant %q: %w", current.ID, err) + } + + fmt.Printf( + "- tenant=%s region=%s shard=%s name=%s\n", + stored.ID, + stored.Region, + resolved.ID(), + stored.Name, + ) + } + + _, err = tenantResolver.Resolve(tenantKey{ + TenantID: "tenant-99", + Region: "apac", + }) + if !errors.Is(err, shard.ErrNoShard) { + return fmt.Errorf("resolve unsupported region: got %v, want shard.ErrNoShard", err) + } + + fmt.Println("unsupported region: no shard") + + return nil +} + +func upsertTenant(ctx context.Context, resolved shard.Shard, current tenant) error { + return resolved.InPrimaryTx( + ctx, + pgx.TxOptions{}, + func(ctx context.Context, tx pgx.Tx) error { + _, err := tx.Exec( + ctx, + `INSERT INTO xpg_shard_geo_example.tenants ( + id, + name, + region, + last_seen_at + ) + VALUES ($1, $2, $3, now()) + ON CONFLICT (id) DO UPDATE SET + name = EXCLUDED.name, + region = EXCLUDED.region, + last_seen_at = EXCLUDED.last_seen_at`, + current.ID, + current.Name, + current.Region, + ) + + return err + }, + ) +} + +func loadTenant(ctx context.Context, resolved shard.Shard, tenantID string) (tenant, error) { + primary := resolved.Primary() + if primary == nil { + return tenant{}, fmt.Errorf("shard %q has no primary", resolved.ID()) + } + + var stored tenant + + err := primary.QueryRow( + ctx, + `SELECT id, name, region + FROM xpg_shard_geo_example.tenants + WHERE id = $1`, + tenantID, + ).Scan( + &stored.ID, + &stored.Name, + &stored.Region, + ) + if err != nil { + return tenant{}, err + } + + return stored, nil +} diff --git a/examples/shard_geo/resolver.go b/examples/shard_geo/resolver.go new file mode 100644 index 0000000..dca5e14 --- /dev/null +++ b/examples/shard_geo/resolver.go @@ -0,0 +1,50 @@ +package main + +import ( + "fmt" + + "github.com/mkbeh/xpg/shard" + "github.com/mkbeh/xpg/shard/resolver" +) + +func newTenantResolver(topology *shard.Topology) (shard.Resolver[tenantKey], error) { + shardByRegion := make(map[string]shard.ID, topology.Len()) + + // Build the routing index once instead of scanning the topology on every + // tenant lookup. + for _, candidate := range topology.Shards() { + region, ok := candidate.Label("region") + if !ok || region == "" { + return nil, fmt.Errorf("shard %q has no region label", candidate.ID()) + } + + if existing, exists := shardByRegion[region]; exists { + return nil, fmt.Errorf( + "region %q is assigned to both shard %q and shard %q", + region, + existing, + candidate.ID(), + ) + } + + shardByRegion[region] = candidate.ID() + } + + resolve := resolver.ResolveFunc[tenantKey]( + func(key tenantKey, _ *shard.Topology) (shard.ID, error) { + id, ok := shardByRegion[key.Region] + if !ok { + return "", fmt.Errorf( + "tenant %q has unsupported region %q: %w", + key.TenantID, + key.Region, + shard.ErrNoShard, + ) + } + + return id, nil + }, + ) + + return resolver.NewCustom(topology, resolve) +} diff --git a/examples/shard_geo/setup.go b/examples/shard_geo/setup.go new file mode 100644 index 0000000..e189457 --- /dev/null +++ b/examples/shard_geo/setup.go @@ -0,0 +1,110 @@ +package main + +import ( + "context" + "fmt" + "os" + + "github.com/mkbeh/xpg" + "github.com/mkbeh/xpg/cluster" + "github.com/mkbeh/xpg/shard" +) + +const ( + defaultShardEUDatabaseURL = "postgres://postgres:postgres@localhost:57431/postgres?sslmode=disable" + defaultShardUSDatabaseURL = "postgres://postgres:postgres@localhost:57432/postgres?sslmode=disable" + + shardEUID shard.ID = "shard-eu" + shardUSID shard.ID = "shard-us" +) + +func openTopology(ctx context.Context) (*shard.Topology, error) { + shardEU, err := openShard( + ctx, + shardEUID, + "eu", + "geo.shard-eu.primary", + environment( + "XPG_SHARD_EU_DATABASE_URL", + defaultShardEUDatabaseURL, + ), + ) + if err != nil { + return nil, fmt.Errorf("open shard-eu: %w", err) + } + + shardUS, err := openShard( + ctx, + shardUSID, + "us", + "geo.shard-us.primary", + environment( + "XPG_SHARD_US_DATABASE_URL", + defaultShardUSDatabaseURL, + ), + ) + if err != nil { + shardEU.Close() + + return nil, fmt.Errorf("open shard-us: %w", err) + } + + topology, err := shard.NewTopology([]shard.Config{ + {Cluster: shardEU}, + {Cluster: shardUS}, + }) + if err != nil { + shardUS.Close() + shardEU.Close() + + return nil, fmt.Errorf("create topology: %w", err) + } + + return topology, nil +} + +func openShard( + ctx context.Context, + id shard.ID, + region string, + name string, + databaseURL string, +) (*cluster.Cluster, error) { + pool, err := xpg.Open( + ctx, + databaseURL, + xpg.WithName(name), + xpg.WithLabel("xpg.shard.id", string(id)), + xpg.WithLabel("xpg.pool.role", "primary"), + ) + if err != nil { + return nil, fmt.Errorf("open pool: %w", err) + } + + if err := pool.Ping(ctx); err != nil { + pool.Close() + return nil, fmt.Errorf("ping pool: %w", err) + } + + shardCluster, err := cluster.New(cluster.Config{ + ID: id, + Labels: map[string]string{ + "region": region, + }, + Primary: pool, + }) + if err != nil { + pool.Close() + return nil, fmt.Errorf("create cluster: %w", err) + } + + return shardCluster, nil +} + +func environment(name, fallback string) string { + if value := os.Getenv(name); value != "" { + return value + } + + return fallback +} diff --git a/examples/shard_geo/sql/schema.sql b/examples/shard_geo/sql/schema.sql new file mode 100644 index 0000000..1b30d39 --- /dev/null +++ b/examples/shard_geo/sql/schema.sql @@ -0,0 +1,10 @@ +DROP SCHEMA IF EXISTS xpg_shard_geo_example CASCADE; + +CREATE SCHEMA xpg_shard_geo_example; + +CREATE TABLE xpg_shard_geo_example.tenants ( + id text PRIMARY KEY, + name text NOT NULL, + region text NOT NULL, + last_seen_at timestamptz NOT NULL +); From 9a1c1f96ee01f9f890b115b249df90626f3b9be9 Mon Sep 17 00:00:00 2001 From: mkbeh Date: Tue, 25 Aug 2026 16:23:03 +0300 Subject: [PATCH 33/41] docs: update examples overview --- examples/README.md | 40 +++++++++++++--------------------------- 1 file changed, 13 insertions(+), 27 deletions(-) diff --git a/examples/README.md b/examples/README.md index 29cd2e8..d84b6ba 100644 --- a/examples/README.md +++ b/examples/README.md @@ -1,33 +1,19 @@ # Examples -This directory contains runnable examples demonstrating the main features and usage patterns of `xpg`. - -| Example | Demonstrates | -|:---------------------------------|:------------------------------------------------------------------------------------------| -| [`basic`](basic) | Pool lifecycle and common query methods | -| [`transactions`](transactions) | Committing an outer transaction after an optional operation is rolled back to a savepoint | -| [`advisory`](advisory) | Coordinating concurrent transactions with PostgreSQL advisory locks | -| [`observability`](observability) | Logging with `slog`, OpenTelemetry tracing, and Prometheus pool metrics | -| [`cluster`](cluster) | Routing reads and transactions across primary and replica pools | -| [`sharding`](shard) | Typed routing across immutable standalone and cluster shard targets | +This directory contains runnable examples covering the main `xpg` usage patterns. + +| Example | Covers | +|:---------------------------------|:----------------------------------------------------------------------------| +| [`basic`](basic) | Core `xpg.Pool` usage for common PostgreSQL operations | +| [`transactions`](transactions) | Transactions, savepoints, and recovering from an optional operation failure | +| [`advisory`](advisory) | Coordinating concurrent work with transaction-level advisory locks | +| [`observability`](observability) | `slog` logging, OpenTelemetry tracing, and Prometheus pool metrics | +| [`cluster`](cluster) | Primary and replica routing, round-robin reads, and read-only transactions | +| [`shard`](shard) | Range-based shard routing and grouping keys by shard | +| [`shard_geo`](shard_geo) | Custom geographic routing built from shard metadata | ## Running the examples -The examples use Docker Compose to start PostgreSQL and any required supporting services. - -From the `examples` directory, start PostgreSQL and Adminer: - -```shell -docker compose --profile tools up -d -``` - -Then run the example from its directory: - -```shell -cd transactions -go run . -``` +Each example is self-contained and includes its own setup and run instructions. -> [!NOTE] -> Some examples may require different services or configuration. Refer to the README in the corresponding example -> directory for the exact startup command, connection settings, and expected output. \ No newline at end of file +Open the corresponding directory and follow its README. From bd029ea6312ff98496abad27b18fb374449e1ba6 Mon Sep 17 00:00:00 2001 From: mkbeh Date: Tue, 25 Aug 2026 20:44:41 +0300 Subject: [PATCH 34/41] docs: add initial project documentation --- README.md | 169 +++++++++++++++++++++++++++++++++++++++++++++++++++--- 1 file changed, 160 insertions(+), 9 deletions(-) diff --git a/README.md b/README.md index 6ed45c4..a08b136 100644 --- a/README.md +++ b/README.md @@ -1,29 +1,180 @@
-# xpg +# Postgres toolkit for Go **Lightweight PostgreSQL wrapper for Go, built on top of [pgx](https://github.com/jackc/pgx).** -![Go Version](https://img.shields.io/badge/go-1.27%2B-blue) -[![License: MIT](https://img.shields.io/badge/license-MIT-green.svg)](LICENSE) +[![Go Reference](https://pkg.go.dev/badge/github.com/mkbeh/xpg.svg)](https://pkg.go.dev/github.com/mkbeh/xpg) +[![Test](https://github.com/mkbeh/xpg/actions/workflows/test.yml/badge.svg)](https://github.com/mkbeh/xpg/actions/workflows/test.yml) +[![Coverage](https://codecov.io/gh/mkbeh/xpg/graph/badge.svg)](https://codecov.io/gh/mkbeh/xpg)
-`xpg` wraps the excellent [`pgx`](https://github.com/jackc/pgx) PostgreSQL driver with a compact API for common -PostgreSQL workflows: read/write connection pool splitting, transaction helpers, embedded SQL migrations, query -building, normalized errors, and exposing PostgreSQL observability with OpenTelemetry and Prometheus. +`xpg` builds on the `pgx` client with a compact API for common PostgreSQL infrastructure patterns. It adds support for +pool lifecycle, transactions and savepoints, PostgreSQL error classification, advisory locks, primary/replica routing, +application-level sharding, and observability. + +The library uses `pgx` types and query model directly while keeping its core behavior and reducing boilerplate around +connection management, routing, and common production workflows. ## Features -[TODOO] +* **Pool Lifecycle Management:** Thin pool management on top of `pgx` with direct access to the underlying PostgreSQL + client. +* **Transactions and Savepoints:** Managed transactions, savepoints, and helpers for common multi-step transactional + workflows. +* **Error Classification:** Classification of PostgreSQL constraint, transaction, cancellation, connection, and other + common database errors. +* **Advisory Locking:** Transaction-level advisory locks for coordinating concurrent database operations. +* **Primary/Replica Routing:** Explicit read policies, replica selection, primary fallback, and read-only transactions + across PostgreSQL nodes. +* **Application-Level Sharding:** Hash, range, time-based, and custom routing with colocation checks, key grouping, and + bounded concurrent fan-out. +* **Observability:** Structured logging, tracing, pool statistics, and optional OpenTelemetry metrics. ## Installation -[TODO] +This repository contains the core `xpg` module. The core package is released from the repository root: + +```bash +go get github.com/mkbeh/xpg +``` + +Optional integrations are released independently under `extra`: + +```bash +go get github.com/mkbeh/xpg/extra/otelxpg +``` ## Quick start -[TODO] +Open an `xpg` pool and execute a PostgreSQL query: + + +```go +// urlExample := "postgres://username:password@localhost:5432/database_name" +pool, err := xpg.Open( + context.Background(), + os.Getenv("DATABASE_URL"), + xpg.WithName("example-pool"), +) +if err != nil { + log.Fatalf("failed to open pool: %v", err) +} +defer pool.Close() + +var message string +err = pool.QueryRow(context.Background(), "SELECT 'hello from xpg'").Scan(&message) +if err != nil { + log.Fatalf("query failed: %v", err) +} + +fmt.Println(message) // Outputs: hello from xpg +``` + + +## Clustering + +`xpg` groups primary and replica pools into a logical cluster with explicit read routing. + + +```go +orders, err := cluster.New(cluster.Config{ + ID: "orders", + Primary: primary, + Replicas: []*xpg.Pool{replicaA, replicaB}, +}) +if err != nil { + panic(err) +} +defer orders.Close() + +// Route writes explicitly to the primary. +primaryPool := orders.Primary() + +_, err = primaryPool.Exec(ctx, "UPDATE orders SET status = 'processed' WHERE id = $1", orderID) +if err != nil { + panic(err) +} + +// Route reads according to the selected policy. +readPool, err := orders.ReadPool(ctx, cluster.ReadReplicaPreferred) +if err != nil { + panic(err) +} + +var status string +err = readPool.QueryRow(ctx, "SELECT status FROM orders WHERE id = $1", orderID).Scan(&status) +if err != nil { + panic(err) +} +``` + + +Read policies support primary-only, replica-required, and replica-preferred routing with primary fallback when no +replica is available. Replica selection is round-robin by default and can be customized. + +## Sharding + +`xpg` provides application-level sharding with explicit key routing across an immutable shard topology. + + +```go +topology, err := shard.NewTopology([]shard.Config{ + {Cluster: shardA}, + {Cluster: shardB}, +}) +if err != nil { + panic(err) +} +defer topology.Close() + +// Partition user IDs into shard ranges. +users, err := resolver.NewRange( + topology, + []resolver.Range[uint64]{ + {Start: 0, End: 100, ShardID: "shard-a"}, + {Start: 100, End: 200, ShardID: "shard-b"}, + }, +) +if err != nil { + panic(err) +} + +// Resolve the target shard. +shard, err := users.Resolve(userID) +if err != nil { + panic(err) +} + +// Write to the shard primary. +primaryPool := shard.Primary() + +_, err = primaryPool.Exec(ctx, "UPDATE users SET active = true WHERE id = $1", userID) +if err != nil { + panic(err) +} + +// Read from the same shard using the selected read policy. +readPool, err := shard.ReadPool(ctx, cluster.ReadReplicaPreferred) +if err != nil { + panic(err) +} + +var active bool +err = readPool.QueryRow(ctx, "SELECT active FROM users WHERE id = $1", userID).Scan(&active) +if err != nil { + panic(err) +} +``` + + +Built-in routing strategies include rendezvous hashing, ordered ranges, time ranges, and custom resolvers. Sharding +utilities cover key colocation, grouping by shard, and bounded concurrent fan-out. + +## Examples + +See the [examples](examples) directory for runnable examples covering the main `xpg` usage patterns. ## License From 67356c576e66b218490f65439c728bb09fdf5aa9 Mon Sep 17 00:00:00 2001 From: mkbeh Date: Tue, 25 Aug 2026 21:36:37 +0300 Subject: [PATCH 35/41] docs: expand usage documentation --- README.md | 63 ++++++++++++++++++++++++++++++++++++++++++++++++++++--- 1 file changed, 60 insertions(+), 3 deletions(-) diff --git a/README.md b/README.md index a08b136..a2b0b18 100644 --- a/README.md +++ b/README.md @@ -34,7 +34,7 @@ connection management, routing, and common production workflows. ## Installation -This repository contains the core `xpg` module. The core package is released from the repository root: +This repository contains the core `xpg` module. The core module is released from the repository root: ```bash go get github.com/mkbeh/xpg @@ -46,7 +46,7 @@ Optional integrations are released independently under `extra`: go get github.com/mkbeh/xpg/extra/otelxpg ``` -## Quick start +## Usage Open an `xpg` pool and execute a PostgreSQL query: @@ -73,6 +73,63 @@ fmt.Println(message) // Outputs: hello from xpg ``` +### Transactions + +`xpg` provides managed transactions using the native `pgx` transaction API. Returning `nil` commits the transaction; +returning an error rolls it back. + + +```go +err := pool.InTx(ctx, pgx.TxOptions{}, func(ctx context.Context, tx pgx.Tx) error { + _, err := tx.Exec(ctx, "UPDATE users SET active = true WHERE id = $1", userID) + return err +}) +``` + + +Savepoints can isolate optional work without aborting the outer transaction. + +### Advisory Locks + +`xpg` provides transaction-level PostgreSQL advisory locks for coordinating concurrent work. + + +```go +err := pool.InTx(ctx, pgx.TxOptions{}, func(ctx context.Context, tx pgx.Tx) error { + if err := xpg.AdvisoryXactLock(ctx, tx, lockID); err != nil { + return err + } + + _, err := tx.Exec(ctx, "UPDATE jobs SET status = 'running' WHERE id = $1", jobID) + return err +}) +``` + + +The lock is held for the duration of the transaction and released automatically on commit or rollback. + +### Error Handling + +`xpg` provides helpers for classifying PostgreSQL errors and inspecting SQLSTATE codes. + + +```go +_, err := pool.Exec(ctx, "INSERT INTO users (id, email) VALUES ($1, $2)", userID, email) + +switch { +case xpg.IsUniqueViolation(err): + // Handle duplicate data. +case xpg.IsRetryableTransaction(err): + // Retry the transaction when the operation is safe to replay. +case err != nil: + return err +} +``` + + +The underlying SQLSTATE code is also available through `xpg.SQLState`. Helpers cover constraint violations, +serialization failures, deadlocks, lock errors, query cancellation, and connection failures. + ## Clustering `xpg` groups primary and replica pools into a logical cluster with explicit read routing. @@ -170,7 +227,7 @@ if err != nil { Built-in routing strategies include rendezvous hashing, ordered ranges, time ranges, and custom resolvers. Sharding -utilities cover key colocation, grouping by shard, and bounded concurrent fan-out. +utilities cover key colocation, grouping by shard, and parallel operations across shards. ## Examples From 4ed8cd29d4aa140f5baee42bfa15c7dc952b4fbb Mon Sep 17 00:00:00 2001 From: mkbeh Date: Tue, 25 Aug 2026 22:03:15 +0300 Subject: [PATCH 36/41] docs: add changelog --- CHANGELOG.md | 52 +++++++++++++++++++++++++++++++++++++++++++++++++++- 1 file changed, 51 insertions(+), 1 deletion(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index c2c8c48..621ec21 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1 +1,51 @@ -[TODO] \ No newline at end of file +# Changelog + +All notable changes to this project will be documented in this file. + +## v0.2.0 + +Initial production release of `xpg`, built around `pgx` with PostgreSQL transaction helpers, primary/replica clustering, +application-level sharding, and observability integrations. + +### Added + +* **Pool Lifecycle Management:** Added a thin pool layer over `pgxpool` with pgx query operations, pool metadata, + runtime statistics, configuration options, and direct access to the underlying pool. +* **Transactions and Savepoints:** Added managed transactions and savepoints for isolating optional work within a + transaction. +* **Advisory Locks:** Added transaction-level PostgreSQL advisory locks for coordinating concurrent work. +* **PostgreSQL Error Classification:** Added SQLSTATE inspection and helpers for constraint violations, serialization + failures, deadlocks, lock errors, query cancellation, and connection failures. +* **Primary/Replica Clustering:** Added logical clusters with explicit read policies, round-robin and custom replica + selection, replica-preferred fallback when no replica can be selected, and enforced read-only transactions. +* **Application-Level Sharding:** Added immutable shard topologies with rendezvous hashing, ordered ranges, time ranges, + and custom resolvers, together with colocation checks, grouping by shard, and bounded parallel operations. +* **Logging and Tracing:** Added pgx-compatible logging and tracing hooks with support for combining multiple query + tracers. +* **Examples:** Added runnable examples covering pool usage, transactions, advisory locks, observability, clustering, + range-based sharding, and custom geographic routing. + +--- + +## extra/otelxpg/v0.1.0 + +Initial release of the `otelxpg` integration module. + +### Added + +* **OpenTelemetry Pool Metrics:** Added metrics for PostgreSQL connection-pool state, usage, acquisition behavior, and + lifecycle counters. +* **Metric Attributes:** Added pool names and custom labels as metric attributes for filtering and aggregation across + standalone pools, cluster nodes, and shard pools. +* **Meter Provider Integration:** Added support for an application-provided `metric.MeterProvider`. + +--- + +## extra/slogxpg/v0.1.0 + +Initial release of the `slogxpg` integration module. + +### Added + +* **slog Adapter:** Added an adapter from the standard library `log/slog` logger to `pgx/tracelog.Logger`. +* **Log Level Mapping:** Added deterministic mapping from pgx trace log levels to `slog` levels. \ No newline at end of file From 3ac53fcbe5d7f6e9a2e77dc1655f9ac974c2bd8f Mon Sep 17 00:00:00 2001 From: mkbeh Date: Tue, 25 Aug 2026 22:19:50 +0300 Subject: [PATCH 37/41] docs: add otelxpg usage guide --- extra/otelxpg/README.md | 46 +++++++++++++++++++++++++++++++++++++++++ 1 file changed, 46 insertions(+) create mode 100644 extra/otelxpg/README.md diff --git a/extra/otelxpg/README.md b/extra/otelxpg/README.md new file mode 100644 index 0000000..492f65c --- /dev/null +++ b/extra/otelxpg/README.md @@ -0,0 +1,46 @@ +# OpenTelemetry Metrics for xpg + +`otelxpg` provides optional OpenTelemetry metrics integration for `xpg`. + +The package is exporter-agnostic: applications own the OpenTelemetry SDK lifecycle and exporter configuration, while +`otelxpg` uses the configured `MeterProvider` to expose connection pool metrics. + +## Installation + +```bash +go get github.com/mkbeh/xpg/extra/otelxpg +``` + +## Usage + + +```go +import ( + "context" + + "github.com/mkbeh/xpg" + "github.com/mkbeh/xpg/extra/otelxpg" + sdkmetric "go.opentelemetry.io/otel/sdk/metric" +) + +meterProvider := sdkmetric.NewMeterProvider() +defer meterProvider.Shutdown(context.Background()) + +// Create a reusable OpenTelemetry metrics integration. +metrics := otelxpg.NewMetrics( + otelxpg.WithMeterProvider(meterProvider), +) + +// Attach metrics when creating the pool. +pool, _ := xpg.Open( + context.Background(), + databaseURL, + xpg.WithName("example-pool"), + xpg.WithMetrics(metrics), +) +defer pool.Close() +``` + + +For a complete runnable setup using OpenTelemetry and Prometheus, see the +[observability example](../../examples/observability). \ No newline at end of file From d4c256c7bd066ddc152e884bae05195bab5be7c5 Mon Sep 17 00:00:00 2001 From: mkbeh Date: Wed, 26 Aug 2026 00:39:26 +0300 Subject: [PATCH 38/41] refactor: polish core APIs and documentation --- advisory.go | 2 +- advisory_test.go | 218 ++++++-- cluster/cluster.go | 34 +- cluster/cluster_test.go | 2 +- cluster/doc.go | 7 +- cluster/errors.go | 8 +- cluster/helpers_test.go | 3 +- cluster/resolver.go | 16 +- cluster/{resolver_test.go => routing_test.go} | 55 +- cluster/selector.go | 34 +- cluster/selector_test.go | 13 +- cluster/tx.go | 15 +- cluster/tx_test.go | 22 +- doc.go | 2 +- errors.go | 22 +- errors_test.go | 162 +++++- examples/observability/go.mod | 8 +- extra/otelxpg/go.mod | 4 +- extra/otelxpg/metrics.go | 11 +- extra/otelxpg/metrics_test.go | 508 ++++++++++++++++++ extra/otelxpg/options.go | 32 +- extra/otelxpg/options_test.go | 66 ++- extra/otelxpg/pool.go | 220 ++++---- extra/slogxpg/logger_test.go | 4 +- go.sum | 9 +- metrics.go | 9 +- options.go | 197 +++---- pool.go | 28 +- pool_test.go | 274 ++++++++++ shard/doc.go | 16 +- shard/errors.go | 15 +- shard/foreach.go | 22 +- shard/foreach_test.go | 43 +- shard/group.go | 16 +- shard/group_test.go | 2 +- shard/helpers_test.go | 3 +- shard/resolver.go | 2 +- shard/resolver/custom.go | 16 +- shard/resolver/custom_test.go | 46 +- shard/resolver/doc.go | 5 +- shard/resolver/encoder.go | 44 +- shard/resolver/encoder_test.go | 2 + shard/resolver/hash.go | 59 +- shard/resolver/hash_test.go | 65 ++- shard/resolver/helpers_test.go | 3 +- shard/resolver/range.go | 28 +- shard/resolver/range_test.go | 133 +++-- shard/resolver/time_range.go | 34 +- shard/resolver/time_range_test.go | 87 +-- shard/resolver/validation.go | 8 +- shard/shard.go | 14 +- shard/shard_test.go | 9 +- shard/topology.go | 63 ++- shard/topology_test.go | 14 +- stats.go | 10 +- tx.go | 25 +- tx_test.go | 165 +++++- 57 files changed, 2142 insertions(+), 792 deletions(-) rename cluster/{resolver_test.go => routing_test.go} (82%) create mode 100644 extra/otelxpg/metrics_test.go create mode 100644 pool_test.go diff --git a/advisory.go b/advisory.go index a06bef4..369c236 100644 --- a/advisory.go +++ b/advisory.go @@ -15,7 +15,7 @@ const ( // AdvisoryXactLock acquires an exclusive transaction-level advisory lock. // -// The call waits until the lock is available or ctx is canceled. PostgreSQL +// The call blocks until the lock is acquired or ctx is canceled. PostgreSQL // releases the lock automatically when tx is committed or rolled back. func AdvisoryXactLock(ctx context.Context, tx pgx.Tx, key int64) error { if tx == nil { diff --git a/advisory_test.go b/advisory_test.go index df8eb68..24cbcb8 100644 --- a/advisory_test.go +++ b/advisory_test.go @@ -9,51 +9,41 @@ import ( "github.com/jackc/pgx/v5/pgconn" ) -type advisoryTestTx struct { - pgx.Tx +func TestAdvisoryXactLockNilTx(t *testing.T) { + t.Parallel() - execErr error - row pgx.Row -} + err := AdvisoryXactLock( + t.Context(), + nil, + 1, + ) -func (tx *advisoryTestTx) Exec(context.Context, string, ...any) (pgconn.CommandTag, error) { - return pgconn.CommandTag{}, tx.execErr + assertErrorMessage(t, err, "xpg: transaction is nil") } -func (tx *advisoryTestTx) QueryRow(context.Context, string, ...any) pgx.Row { - return tx.row -} +func TestAdvisoryXactLock(t *testing.T) { + t.Parallel() -type advisoryTestRow struct { - acquired bool - err error -} + tx := &advisoryTestTx{} -func (row advisoryTestRow) Scan(dest ...any) error { - if row.err != nil { - return row.err + err := AdvisoryXactLock( + t.Context(), + tx, + 42, + ) + if err != nil { + t.Fatalf("AdvisoryXactLock returned an error: %v", err) } - acquired, ok := dest[0].(*bool) - if !ok { - return errors.New("unexpected destination type") + if tx.execSQL != advisoryXactLockSQL { + t.Fatalf( + "SQL = %q, want %q", + tx.execSQL, + advisoryXactLockSQL, + ) } - *acquired = row.acquired - - return nil -} - -func TestAdvisoryXactLockNilTx(t *testing.T) { - t.Parallel() - - err := AdvisoryXactLock( - context.Background(), - nil, - 1, - ) - - assertErrorMessage(t, err, "xpg: transaction is nil") + assertAdvisoryKey(t, tx.execArgs, 42) } func TestAdvisoryXactLockPreservesError(t *testing.T) { @@ -62,8 +52,10 @@ func TestAdvisoryXactLockPreservesError(t *testing.T) { expectedErr := errors.New("lock failed") err := AdvisoryXactLock( - context.Background(), - &advisoryTestTx{execErr: expectedErr}, + t.Context(), + &advisoryTestTx{ + execErr: expectedErr, + }, 1, ) if !errors.Is(err, expectedErr) { @@ -75,7 +67,7 @@ func TestTryAdvisoryXactLockNilTx(t *testing.T) { t.Parallel() _, err := TryAdvisoryXactLock( - context.Background(), + t.Context(), nil, 1, ) @@ -86,19 +78,52 @@ func TestTryAdvisoryXactLockNilTx(t *testing.T) { func TestTryAdvisoryXactLock(t *testing.T) { t.Parallel() - acquired, err := TryAdvisoryXactLock( - context.Background(), - &advisoryTestTx{ - row: advisoryTestRow{acquired: true}, - }, - 1, - ) - if err != nil { - t.Fatalf("TryAdvisoryXactLock returned an error: %v", err) - } - - if !acquired { - t.Fatal("TryAdvisoryXactLock returned false") + for _, acquired := range []bool{false, true} { + t.Run( + map[bool]string{ + false: "not acquired", + true: "acquired", + }[acquired], + func(t *testing.T) { + t.Parallel() + + tx := &advisoryTestTx{ + row: advisoryTestRow{ + acquired: acquired, + }, + } + + got, err := TryAdvisoryXactLock( + t.Context(), + tx, + 42, + ) + if err != nil { + t.Fatalf( + "TryAdvisoryXactLock returned an error: %v", + err, + ) + } + + if got != acquired { + t.Fatalf( + "acquired = %v, want %v", + got, + acquired, + ) + } + + if tx.querySQL != tryAdvisoryXactLockSQL { + t.Fatalf( + "SQL = %q, want %q", + tx.querySQL, + tryAdvisoryXactLockSQL, + ) + } + + assertAdvisoryKey(t, tx.queryArgs, 42) + }, + ) } } @@ -108,9 +133,11 @@ func TestTryAdvisoryXactLockPreservesError(t *testing.T) { expectedErr := errors.New("try lock failed") _, err := TryAdvisoryXactLock( - context.Background(), + t.Context(), &advisoryTestTx{ - row: advisoryTestRow{err: expectedErr}, + row: advisoryTestRow{ + err: expectedErr, + }, }, 1, ) @@ -118,3 +145,88 @@ func TestTryAdvisoryXactLockPreservesError(t *testing.T) { t.Fatalf("original error was not preserved: %v", err) } } + +type advisoryTestTx struct { + pgx.Tx + + execSQL string + execArgs []any + execErr error + + querySQL string + queryArgs []any + row pgx.Row +} + +func (tx *advisoryTestTx) Exec( + _ context.Context, + sql string, + args ...any, +) (pgconn.CommandTag, error) { + tx.execSQL = sql + tx.execArgs = args + + return pgconn.CommandTag{}, tx.execErr +} + +func (tx *advisoryTestTx) QueryRow( + _ context.Context, + sql string, + args ...any, +) pgx.Row { + tx.querySQL = sql + tx.queryArgs = args + + return tx.row +} + +type advisoryTestRow struct { + acquired bool + err error +} + +func (row advisoryTestRow) Scan(dest ...any) error { + if row.err != nil { + return row.err + } + + acquired, ok := dest[0].(*bool) + if !ok { + return errors.New("unexpected destination type") + } + + *acquired = row.acquired + + return nil +} + +func assertAdvisoryKey( + t *testing.T, + args []any, + want int64, +) { + t.Helper() + + if len(args) != 1 { + t.Fatalf( + "argument count = %d, want 1", + len(args), + ) + } + + key, ok := args[0].(int64) + if !ok { + t.Fatalf( + "argument type = %T, want int64", + args[0], + ) + } + + if key != want { + t.Fatalf( + "key = %d, want %d", + key, + want, + ) + } +} diff --git a/cluster/cluster.go b/cluster/cluster.go index 1aa9648..dbe925d 100644 --- a/cluster/cluster.go +++ b/cluster/cluster.go @@ -3,6 +3,7 @@ package cluster import ( "errors" "fmt" + "maps" "slices" "sync" @@ -14,9 +15,9 @@ type ID string // Config configures a Cluster from independently created pools. // -// ID and Labels are optional cluster metadata. New takes ownership of Primary -// and Replicas only after it returns successfully. Cluster.Close closes the -// owned pools. +// ID and Labels are optional metadata. New takes ownership of Primary and +// Replicas only after it returns successfully. Cluster.Close closes all owned +// pools. type Config struct { ID ID Labels map[string]string @@ -26,8 +27,8 @@ type Config struct { Selector ReplicaSelector } -// Cluster routes operations across an optional primary pool and zero or more -// replica pools. +// Cluster represents a logical PostgreSQL cluster composed of an optional +// primary pool and zero or more replica pools. // // A deployment with one PostgreSQL endpoint is represented by a Cluster with // one Primary and no Replicas. A read-only deployment may omit Primary and @@ -95,7 +96,7 @@ func New(config Config) (*Cluster, error) { }, nil } -// ID returns the stable logical cluster ID. +// ID returns the logical cluster ID. func (c *Cluster) ID() ID { if c == nil { return "" @@ -148,24 +149,21 @@ func (c *Cluster) ReplicaCount() int { // ReplicaAt returns the replica at index in registration order. // // The returned pool is borrowed and must not be closed separately. ReplicaAt -// panics when c is nil or index is outside the replica set, matching ordinary -// slice indexing semantics. +// panics when c is nil or index is out of range. func (c *Cluster) ReplicaAt(index int) *xpg.Pool { return c.replicas[index] } -// Close closes all replica pools in reverse registration order and then closes -// the primary pool when one is configured. -// -// Close is safe to call multiple times. +// Close closes replicas in reverse registration order, then closes the primary +// when one is configured. Close is safe to call multiple times. func (c *Cluster) Close() { if c == nil { return } c.closeOnce.Do(func() { - for _, v := range slices.Backward(c.replicas) { - v.Close() + for _, replica := range slices.Backward(c.replicas) { + replica.Close() } if c.primary != nil { @@ -174,6 +172,14 @@ func (c *Cluster) Close() { }) } +func cloneLabels(labels map[string]string) map[string]string { + if len(labels) == 0 { + return nil + } + + return maps.Clone(labels) +} + func validateLabels(labels map[string]string) error { for key := range labels { if key == "" { diff --git a/cluster/cluster_test.go b/cluster/cluster_test.go index 491a45e..403d64b 100644 --- a/cluster/cluster_test.go +++ b/cluster/cluster_test.go @@ -242,7 +242,7 @@ func TestNewCapturesReplicaMetadata(t *testing.T) { } t.Cleanup(cluster.Close) - resolved, err := cluster.ReadPool(context.Background(), ReadReplicaRequired) + resolved, err := cluster.ReadPool(t.Context(), ReadReplicaRequired) if err != nil { t.Fatalf("ReadPool() error = %v", err) } diff --git a/cluster/doc.go b/cluster/doc.go index 1286c8b..d47cd57 100644 --- a/cluster/doc.go +++ b/cluster/doc.go @@ -1,6 +1,3 @@ -// Package cluster provides explicit routing between a PostgreSQL primary pool, -// when configured, and zero or more replica pools. -// -// The package does not inspect SQL, retry failed queries, promote replicas, -// or discover PostgreSQL nodes. +// Package cluster provides primary/replica routing for PostgreSQL connection +// pools. package cluster diff --git a/cluster/errors.go b/cluster/errors.go index 8267d4a..03adc17 100644 --- a/cluster/errors.go +++ b/cluster/errors.go @@ -3,11 +3,11 @@ package cluster import "errors" var ( - // ErrNoPrimary is returned when an operation requires a primary pool but - // the cluster has no primary configured. + // ErrNoPrimary is returned when an operation requires a primary but none + // is configured. ErrNoPrimary = errors.New("xpg/cluster: no primary available") - // ErrNoReplica is returned when an operation requires a replica but no - // replica can be selected. + // ErrNoReplica is returned when an operation requires a replica but none + // can be selected. ErrNoReplica = errors.New("xpg/cluster: no replica available") ) diff --git a/cluster/helpers_test.go b/cluster/helpers_test.go index 77f5411..092a235 100644 --- a/cluster/helpers_test.go +++ b/cluster/helpers_test.go @@ -1,7 +1,6 @@ package cluster import ( - "context" "testing" "github.com/jackc/pgx/v5/pgxpool" @@ -29,7 +28,7 @@ func newTestPool(t *testing.T, name string, labels map[string]string) *xpg.Pool options = append(options, xpg.WithLabels(labels)) } - pool, err := xpg.New(context.Background(), config, options...) + pool, err := xpg.New(t.Context(), config, options...) if err != nil { t.Fatalf("xpg.New() error = %v", err) } diff --git a/cluster/resolver.go b/cluster/resolver.go index 3294ba5..0d0a902 100644 --- a/cluster/resolver.go +++ b/cluster/resolver.go @@ -30,8 +30,8 @@ const ( readPolicyReplicaRequired = "replica_required" ) -// ParsePolicy parses a ReadPolicy from its string representation. -func ParsePolicy(value string) (ReadPolicy, error) { +// ParseReadPolicy parses a ReadPolicy from its string representation. +func ParseReadPolicy(value string) (ReadPolicy, error) { switch value { case readPolicyPrimary: return ReadPrimary, nil @@ -40,7 +40,10 @@ func ParsePolicy(value string) (ReadPolicy, error) { case readPolicyReplicaRequired: return ReadReplicaRequired, nil default: - return 0, fmt.Errorf("xpg/cluster: unknown read policy %q", value) + return 0, fmt.Errorf( + "xpg/cluster: unknown read policy %q", + value, + ) } } @@ -58,11 +61,10 @@ func (policy ReadPolicy) String() string { } } -// ReadPool returns a pool for a read operation according to policy. +// ReadPool returns a pool according to policy. // -// ReadReplicaPreferred falls back to the primary when no replica can be -// selected. If the cluster has no primary, it returns ErrNoPrimary. -// Other selector errors are returned to the caller. +// ReadReplicaPreferred falls back to the primary only when no replica can be +// selected. Other selector errors are returned to the caller. func (c *Cluster) ReadPool(ctx context.Context, policy ReadPolicy) (*xpg.Pool, error) { if c == nil { return nil, errors.New("xpg/cluster: cluster is nil") diff --git a/cluster/resolver_test.go b/cluster/routing_test.go similarity index 82% rename from cluster/resolver_test.go rename to cluster/routing_test.go index 7d2fc6c..bd08329 100644 --- a/cluster/resolver_test.go +++ b/cluster/routing_test.go @@ -3,13 +3,13 @@ package cluster import ( "context" "errors" - "strings" + "fmt" "testing" "github.com/mkbeh/xpg" ) -func TestParsePolicy(t *testing.T) { +func TestParseReadPolicy(t *testing.T) { t.Parallel() tests := []struct { @@ -25,22 +25,22 @@ func TestParsePolicy(t *testing.T) { t.Run(test.value, func(t *testing.T) { t.Parallel() - got, err := ParsePolicy(test.value) + got, err := ParseReadPolicy(test.value) if err != nil { - t.Fatalf("ParsePolicy() error = %v", err) + t.Fatalf("ParseReadPolicy() error = %v", err) } if got != test.want { - t.Fatalf("ParsePolicy(%q) = %v, want %v", test.value, got, test.want) + t.Fatalf("ParseReadPolicy(%q) = %v, want %v", test.value, got, test.want) } }) } } -func TestParsePolicyRejectsUnknown(t *testing.T) { +func TestParseReadPolicyRejectsUnknown(t *testing.T) { t.Parallel() - _, err := ParsePolicy("nearest") + _, err := ParseReadPolicy("nearest") if err == nil { t.Fatal("expected error") } @@ -75,7 +75,7 @@ func TestReadPoolNilCluster(t *testing.T) { var cluster *Cluster - pool, err := cluster.ReadPool(context.Background(), ReadPrimary) + pool, err := cluster.ReadPool(t.Context(), ReadPrimary) if pool != nil { t.Fatalf("ReadPool() pool = %p, want nil", pool) } @@ -95,7 +95,7 @@ func TestReadPoolPrimary(t *testing.T) { primary := newTestPool(t, "primary", nil) cluster := newTestCluster(t, Config{Primary: primary}) - got, err := cluster.ReadPool(context.Background(), ReadPrimary) + got, err := cluster.ReadPool(t.Context(), ReadPrimary) if err != nil { t.Fatalf("ReadPool() error = %v", err) } @@ -113,7 +113,7 @@ func TestReadPoolPrimaryWithoutPrimary(t *testing.T) { Replicas: []*xpg.Pool{replica}, }) - pool, err := cluster.ReadPool(context.Background(), ReadPrimary) + pool, err := cluster.ReadPool(t.Context(), ReadPrimary) if pool != nil { t.Fatalf("ReadPool() pool = %p, want nil", pool) } @@ -133,7 +133,7 @@ func TestReadPoolReplicaPreferredUsesReplica(t *testing.T) { Replicas: []*xpg.Pool{replica}, }) - got, err := cluster.ReadPool(context.Background(), ReadReplicaPreferred) + got, err := cluster.ReadPool(t.Context(), ReadReplicaPreferred) if err != nil { t.Fatalf("ReadPool() error = %v", err) } @@ -149,7 +149,7 @@ func TestReadPoolReplicaPreferredFallsBackWithoutReplicas(t *testing.T) { primary := newTestPool(t, "primary", nil) cluster := newTestCluster(t, Config{Primary: primary}) - got, err := cluster.ReadPool(context.Background(), ReadReplicaPreferred) + got, err := cluster.ReadPool(t.Context(), ReadReplicaPreferred) if err != nil { t.Fatalf("ReadPool() error = %v", err) } @@ -174,7 +174,7 @@ func TestReadPoolReplicaPreferredFallsBackOnErrNoReplica(t *testing.T) { Selector: selector, }) - got, err := cluster.ReadPool(context.Background(), ReadReplicaPreferred) + got, err := cluster.ReadPool(t.Context(), ReadReplicaPreferred) if err != nil { t.Fatalf("ReadPool() error = %v", err) } @@ -200,7 +200,7 @@ func TestReadPoolReplicaPreferredDoesNotFallbackOnSelectorError(t *testing.T) { Selector: selector, }) - pool, err := cluster.ReadPool(context.Background(), ReadReplicaPreferred) + pool, err := cluster.ReadPool(t.Context(), ReadReplicaPreferred) if pool != nil { t.Fatalf("ReadPool() pool = %p, want nil", pool) } @@ -223,7 +223,7 @@ func TestReadPoolReplicaPreferredWithoutPrimary(t *testing.T) { Selector: selector, }) - pool, err := cluster.ReadPool(context.Background(), ReadReplicaPreferred) + pool, err := cluster.ReadPool(t.Context(), ReadReplicaPreferred) if pool != nil { t.Fatalf("ReadPool() pool = %p, want nil", pool) } @@ -241,7 +241,7 @@ func TestReadPoolReplicaRequired(t *testing.T) { Replicas: []*xpg.Pool{replica}, }) - got, err := cluster.ReadPool(context.Background(), ReadReplicaRequired) + got, err := cluster.ReadPool(t.Context(), ReadReplicaRequired) if err != nil { t.Fatalf("ReadPool() error = %v", err) } @@ -257,7 +257,7 @@ func TestReadPoolReplicaRequiredWithoutReplicas(t *testing.T) { primary := newTestPool(t, "primary", nil) cluster := newTestCluster(t, Config{Primary: primary}) - pool, err := cluster.ReadPool(context.Background(), ReadReplicaRequired) + pool, err := cluster.ReadPool(t.Context(), ReadReplicaRequired) if pool != nil { t.Fatalf("ReadPool() pool = %p, want nil", pool) } @@ -284,7 +284,7 @@ func TestReadPoolDefaultSelectorRoundRobin(t *testing.T) { } for call, wantPool := range want { - got, err := cluster.ReadPool(context.Background(), ReadReplicaRequired) + got, err := cluster.ReadPool(t.Context(), ReadReplicaRequired) if err != nil { t.Fatalf("ReadPool() call %d error = %v", call, err) } @@ -320,7 +320,7 @@ func TestReadPoolRejectsSelectorIndex(t *testing.T) { Selector: selector, }) - pool, err := cluster.ReadPool(context.Background(), ReadReplicaRequired) + pool, err := cluster.ReadPool(t.Context(), ReadReplicaRequired) if pool != nil { t.Fatalf("ReadPool() pool = %p, want nil", pool) } @@ -329,8 +329,13 @@ func TestReadPoolRejectsSelectorIndex(t *testing.T) { t.Fatal("expected error") } - if !strings.Contains(err.Error(), "replica selector returned invalid index") { - t.Fatalf("error = %q, want invalid replica index error", err) + want := fmt.Sprintf( + "xpg/cluster: replica selector returned invalid index %d for 1 replicas", + test.index, + ) + + if got := err.Error(); got != want { + t.Fatalf("error = %q, want %q", got, want) } }) } @@ -350,7 +355,7 @@ func TestReadPoolPreservesSelectorError(t *testing.T) { Selector: selector, }) - pool, err := cluster.ReadPool(context.Background(), ReadReplicaRequired) + pool, err := cluster.ReadPool(t.Context(), ReadReplicaRequired) if pool != nil { t.Fatalf("ReadPool() pool = %p, want nil", pool) } @@ -359,8 +364,8 @@ func TestReadPoolPreservesSelectorError(t *testing.T) { t.Fatalf("ReadPool() error = %v, want wrapped selector error", err) } - if !strings.Contains(err.Error(), "xpg/cluster: select replica") { - t.Fatalf("error = %q, want selector context", err) + if got, want := err.Error(), "xpg/cluster: select replica: boom"; got != want { + t.Fatalf("error = %q, want %q", got, want) } } @@ -370,7 +375,7 @@ func TestReadPoolRejectsUnsupportedPolicy(t *testing.T) { primary := newTestPool(t, "primary", nil) cluster := newTestCluster(t, Config{Primary: primary}) - pool, err := cluster.ReadPool(context.Background(), ReadPolicy(255)) + pool, err := cluster.ReadPool(t.Context(), ReadPolicy(255)) if pool != nil { t.Fatalf("ReadPool() pool = %p, want nil", pool) } diff --git a/cluster/selector.go b/cluster/selector.go index 930abf9..ab43dd6 100644 --- a/cluster/selector.go +++ b/cluster/selector.go @@ -3,18 +3,17 @@ package cluster import ( "context" "errors" - "maps" "sync/atomic" ) -// ReplicaInfo contains immutable metadata captured from one replica pool when -// the Cluster is created. +// ReplicaInfo contains immutable metadata captured from a replica pool when the +// Cluster is created. type ReplicaInfo struct { name string labels map[string]string } -// Name returns the stable logical pool name. +// Name returns the logical replica pool name. func (info ReplicaInfo) Name() string { return info.name } @@ -22,6 +21,7 @@ func (info ReplicaInfo) Name() string { // Label returns one replica label without allocating a copy of all labels. func (info ReplicaInfo) Label(key string) (string, bool) { value, ok := info.labels[key] + return value, ok } @@ -46,8 +46,12 @@ func (replicas replicaMetadata) At(index int) ReplicaInfo { return replicas[index] } -// ReplicaSelector selects one replica index from the supplied metadata. -// Implementations used by concurrent callers must be concurrency-safe. +// ReplicaSelector selects one replica from the supplied metadata. +// +// Implementations must be safe for concurrent use. Select must return +// ErrNoReplica when no replica is eligible for selection. Other errors are +// propagated to the caller. A successful call must return a valid replica +// index. type ReplicaSelector interface { Select(ctx context.Context, replicas ReplicaSet) (index int, err error) } @@ -55,7 +59,7 @@ type ReplicaSelector interface { // ReplicaSelectorFunc adapts a function to ReplicaSelector. type ReplicaSelectorFunc func(context.Context, ReplicaSet) (int, error) -// Select calls selector. +// Select calls the wrapped selector function. func (selector ReplicaSelectorFunc) Select(ctx context.Context, replicas ReplicaSet) (int, error) { if selector == nil { return -1, errors.New("xpg/cluster: replica selector function is nil") @@ -64,16 +68,16 @@ func (selector ReplicaSelectorFunc) Select(ctx context.Context, replicas Replica return selector(ctx, replicas) } -type roundRobinSelector struct { - next atomic.Uint64 -} - // RoundRobinSelector returns a concurrency-safe selector that distributes // selections across replicas in registration order. func RoundRobinSelector() ReplicaSelector { return &roundRobinSelector{} } +type roundRobinSelector struct { + next atomic.Uint64 +} + func (selector *roundRobinSelector) Select(_ context.Context, replicas ReplicaSet) (int, error) { length := replicas.Len() @@ -88,11 +92,3 @@ func (selector *roundRobinSelector) Select(_ context.Context, replicas ReplicaSe return int(next % uint64(length)), nil } - -func cloneLabels(labels map[string]string) map[string]string { - if len(labels) == 0 { - return nil - } - - return maps.Clone(labels) -} diff --git a/cluster/selector_test.go b/cluster/selector_test.go index de5005b..348eefa 100644 --- a/cluster/selector_test.go +++ b/cluster/selector_test.go @@ -52,7 +52,7 @@ func TestReplicaSelectorFunc(t *testing.T) { t.Parallel() type contextKey struct{} - ctx := context.WithValue(context.Background(), contextKey{}, "value") + ctx := context.WithValue(t.Context(), contextKey{}, "value") replicas := replicaMetadata{{name: "replica"}} selector := ReplicaSelectorFunc(func(gotCtx context.Context, gotReplicas ReplicaSet) (int, error) { @@ -82,7 +82,7 @@ func TestReplicaSelectorFuncNil(t *testing.T) { var selector ReplicaSelectorFunc - index, err := selector.Select(context.Background(), nil) + index, err := selector.Select(t.Context(), nil) if index != -1 { t.Fatalf("Select() index = %d, want -1", index) } @@ -101,7 +101,7 @@ func TestRoundRobinSelectorEmpty(t *testing.T) { selector := RoundRobinSelector() - index, err := selector.Select(context.Background(), replicaMetadata(nil)) + index, err := selector.Select(t.Context(), replicaMetadata(nil)) if index != -1 { t.Fatalf("Select() index = %d, want -1", index) } @@ -118,7 +118,7 @@ func TestRoundRobinSelectorSingleReplica(t *testing.T) { replicas := replicaMetadata{{name: "replica"}} for range 10 { - index, err := selector.Select(context.Background(), replicas) + index, err := selector.Select(t.Context(), replicas) if err != nil { t.Fatalf("Select() error = %v", err) } @@ -141,7 +141,7 @@ func TestRoundRobinSelectorSequence(t *testing.T) { want := []int{0, 1, 2, 0, 1, 2, 0} for call, wantIndex := range want { - index, err := selector.Select(context.Background(), replicas) + index, err := selector.Select(t.Context(), replicas) if err != nil { t.Fatalf("Select() call %d error = %v", call, err) } @@ -163,6 +163,7 @@ func TestRoundRobinSelectorConcurrent(t *testing.T) { selector := RoundRobinSelector() replicas := make(replicaMetadata, replicaCount) results := make(chan int, callCount) + ctx := t.Context() var waitGroup sync.WaitGroup waitGroup.Add(callCount) @@ -171,7 +172,7 @@ func TestRoundRobinSelectorConcurrent(t *testing.T) { go func() { defer waitGroup.Done() - index, err := selector.Select(context.Background(), replicas) + index, err := selector.Select(ctx, replicas) if err != nil { results <- -1 return diff --git a/cluster/tx.go b/cluster/tx.go index 6a925fa..f2d5818 100644 --- a/cluster/tx.go +++ b/cluster/tx.go @@ -9,8 +9,8 @@ import ( // ReadTxOptions configures a read-only transaction. // -// AccessMode, BeginQuery, and CommitQuery are intentionally controlled by the -// cluster. +// AccessMode is always pgx.ReadOnly. BeginQuery and CommitQuery are not exposed +// so callers cannot override the read-only transaction semantics. type ReadTxOptions struct { IsoLevel pgx.TxIsoLevel DeferrableMode pgx.TxDeferrableMode @@ -26,15 +26,18 @@ func (c *Cluster) InPrimaryTx( return errors.New("xpg/cluster: cluster is nil") } - if c.primary == nil { - return ErrNoPrimary + pool, err := c.resolvePrimary() + if err != nil { + return err } - return c.primary.InTx(ctx, options, fn) + return pool.InTx(ctx, options, fn) } // InReadTx selects a pool according to policy and executes fn in a read-only -// transaction on that pool. +// transaction. +// +// The transaction remains read-only when policy resolves to the primary. func (c *Cluster) InReadTx( ctx context.Context, policy ReadPolicy, diff --git a/cluster/tx_test.go b/cluster/tx_test.go index e839f27..8478e22 100644 --- a/cluster/tx_test.go +++ b/cluster/tx_test.go @@ -16,7 +16,7 @@ func TestInPrimaryTxNilCluster(t *testing.T) { called := false err := cluster.InPrimaryTx( - context.Background(), + t.Context(), pgx.TxOptions{}, func(context.Context, pgx.Tx) error { called = true @@ -46,7 +46,7 @@ func TestInPrimaryTxWithoutPrimary(t *testing.T) { called := false err := cluster.InPrimaryTx( - context.Background(), + t.Context(), pgx.TxOptions{}, func(context.Context, pgx.Tx) error { called = true @@ -63,13 +63,13 @@ func TestInPrimaryTxWithoutPrimary(t *testing.T) { } } -func TestInPrimaryTxDelegatesToPool(t *testing.T) { +func TestInPrimaryTxContextCancellation(t *testing.T) { t.Parallel() primary := newTestPool(t, "primary", nil) cluster := newTestCluster(t, Config{Primary: primary}) - ctx, cancel := context.WithCancel(context.Background()) + ctx, cancel := context.WithCancel(t.Context()) cancel() called := false @@ -82,8 +82,8 @@ func TestInPrimaryTxDelegatesToPool(t *testing.T) { }, ) - if err == nil { - t.Fatal("expected error") + if !errors.Is(err, context.Canceled) { + t.Fatalf("InPrimaryTx() error = %v, want context.Canceled", err) } if called { @@ -99,7 +99,7 @@ func TestInReadTxRoutingError(t *testing.T) { called := false err := cluster.InReadTx( - context.Background(), + t.Context(), ReadReplicaRequired, ReadTxOptions{}, func(context.Context, pgx.Tx) error { @@ -117,7 +117,7 @@ func TestInReadTxRoutingError(t *testing.T) { } } -func TestInReadTxDelegatesToResolvedPool(t *testing.T) { +func TestInReadTxContextCancellation(t *testing.T) { t.Parallel() replica := newTestPool(t, "replica", nil) @@ -125,7 +125,7 @@ func TestInReadTxDelegatesToResolvedPool(t *testing.T) { Replicas: []*xpg.Pool{replica}, }) - ctx, cancel := context.WithCancel(context.Background()) + ctx, cancel := context.WithCancel(t.Context()) cancel() called := false @@ -142,8 +142,8 @@ func TestInReadTxDelegatesToResolvedPool(t *testing.T) { }, ) - if err == nil { - t.Fatal("expected error") + if !errors.Is(err, context.Canceled) { + t.Fatalf("InReadTx() error = %v, want context.Canceled", err) } if called { diff --git a/doc.go b/doc.go index 0b4c1db..0dec990 100644 --- a/doc.go +++ b/doc.go @@ -1,2 +1,2 @@ -// Package xpg provides pgx-first infrastructure primitives for PostgreSQL. +// Package xpg provides PostgreSQL infrastructure utilities built on pgx. package xpg diff --git a/errors.go b/errors.go index ef374d7..66c963d 100644 --- a/errors.go +++ b/errors.go @@ -10,14 +10,15 @@ import ( ) const ( - sqlStateUniqueViolation = "23505" - sqlStateForeignKeyViolation = "23503" - sqlStateNotNullViolation = "23502" - sqlStateCheckViolation = "23514" - sqlStateSerializationFailure = "40001" - sqlStateDeadlockDetected = "40P01" - sqlStateLockNotAvailable = "55P03" - sqlStateQueryCanceled = "57014" + sqlStateUniqueViolation = "23505" + sqlStateForeignKeyViolation = "23503" + sqlStateNotNullViolation = "23502" + sqlStateCheckViolation = "23514" + sqlStateSerializationFailure = "40001" + sqlStateDeadlockDetected = "40P01" + sqlStateLockNotAvailable = "55P03" + sqlStateQueryCanceled = "57014" + sqlStateConnectionExceptionClass = "08" ) // SQLState returns the PostgreSQL SQLSTATE code carried by err. @@ -29,7 +30,7 @@ func SQLState(err error) string { return "" } - return pgErr.Code + return pgErr.SQLState() } // IsNoRows reports whether err indicates that a query returned no rows. @@ -106,7 +107,8 @@ func IsConnectionError(err error) bool { state := SQLState(err) - return len(state) >= 2 && state[:2] == "08" + return len(state) >= 2 && + state[:2] == sqlStateConnectionExceptionClass } // IsRetryableTransaction reports whether PostgreSQL aborted the transaction diff --git a/errors_test.go b/errors_test.go index 6f74f34..a21c04e 100644 --- a/errors_test.go +++ b/errors_test.go @@ -1,8 +1,10 @@ package xpg import ( + "context" "errors" "fmt" + "io" "net" "testing" @@ -13,17 +15,43 @@ import ( func TestSQLState(t *testing.T) { t.Parallel() - err := fmt.Errorf( - "wrapped: %w", - &pgconn.PgError{Code: sqlStateUniqueViolation}, - ) - - if state := SQLState(err); state != sqlStateUniqueViolation { - t.Fatalf( - "unexpected SQLSTATE: got %q, want %q", - state, - sqlStateUniqueViolation, - ) + tests := []struct { + name string + err error + want string + }{ + { + name: "PostgreSQL error", + err: fmt.Errorf( + "wrapped: %w", + &pgconn.PgError{Code: sqlStateUniqueViolation}, + ), + want: sqlStateUniqueViolation, + }, + { + name: "generic error", + err: errors.New("generic"), + want: "", + }, + { + name: "nil", + err: nil, + want: "", + }, + } + + for _, test := range tests { + t.Run(test.name, func(t *testing.T) { + t.Parallel() + + if got := SQLState(test.err); got != test.want { + t.Fatalf( + "SQLState() = %q, want %q", + got, + test.want, + ) + } + }) } } @@ -31,7 +59,11 @@ func TestIsNoRows(t *testing.T) { t.Parallel() if !IsNoRows(fmt.Errorf("wrapped: %w", pgx.ErrNoRows)) { - t.Fatal("IsNoRows returned false") + t.Fatal("IsNoRows returned false for pgx.ErrNoRows") + } + + if IsNoRows(errors.New("generic")) { + t.Fatal("IsNoRows returned true for generic error") } } @@ -100,24 +132,57 @@ func TestErrorClassifiers(t *testing.T) { test.state, ) } + + if test.classifier( + &pgconn.PgError{Code: "00000"}, + ) { + t.Fatalf( + "classifier returned true for unrelated SQLSTATE", + ) + } }) } } +func TestIsQueryCanceledDoesNotClassifyContextCancellation(t *testing.T) { + t.Parallel() + + if IsQueryCanceled(context.Canceled) { + t.Fatal("IsQueryCanceled returned true for context.Canceled") + } + + if IsQueryCanceled(context.DeadlineExceeded) { + t.Fatal("IsQueryCanceled returned true for context.DeadlineExceeded") + } +} + func TestIsConnectionError(t *testing.T) { t.Parallel() tests := []struct { name string err error + want bool }{ { name: "SQLSTATE connection exception", err: &pgconn.PgError{Code: "08006"}, + want: true, }, { name: "closed connection", err: fmt.Errorf("wrapped: %w", pgconn.ErrConnClosed), + want: true, + }, + { + name: "EOF", + err: io.EOF, + want: true, + }, + { + name: "unexpected EOF", + err: io.ErrUnexpectedEOF, + want: true, }, { name: "network operation", @@ -126,6 +191,17 @@ func TestIsConnectionError(t *testing.T) { Net: "tcp", Err: errors.New("connection reset"), }, + want: true, + }, + { + name: "generic error", + err: errors.New("generic"), + want: false, + }, + { + name: "nil", + err: nil, + want: false, }, } @@ -133,8 +209,12 @@ func TestIsConnectionError(t *testing.T) { t.Run(test.name, func(t *testing.T) { t.Parallel() - if !IsConnectionError(test.err) { - t.Fatal("IsConnectionError returned false") + if got := IsConnectionError(test.err); got != test.want { + t.Fatalf( + "IsConnectionError() = %v, want %v", + got, + test.want, + ) } }) } @@ -143,22 +223,46 @@ func TestIsConnectionError(t *testing.T) { func TestIsRetryableTransaction(t *testing.T) { t.Parallel() - for _, state := range []string{ - sqlStateSerializationFailure, - sqlStateDeadlockDetected, - } { - err := &pgconn.PgError{Code: state} - if !IsRetryableTransaction(err) { - t.Fatalf( - "IsRetryableTransaction returned false for SQLSTATE %q", - state, - ) - } + tests := []struct { + name string + state string + want bool + }{ + { + name: "serialization failure", + state: sqlStateSerializationFailure, + want: true, + }, + { + name: "deadlock", + state: sqlStateDeadlockDetected, + want: true, + }, + { + name: "unique violation", + state: sqlStateUniqueViolation, + want: false, + }, + { + name: "connection exception", + state: "08006", + want: false, + }, } - if IsRetryableTransaction( - &pgconn.PgError{Code: sqlStateUniqueViolation}, - ) { - t.Fatal("IsRetryableTransaction returned true for unique violation") + for _, test := range tests { + t.Run(test.name, func(t *testing.T) { + t.Parallel() + + err := &pgconn.PgError{Code: test.state} + + if got := IsRetryableTransaction(err); got != test.want { + t.Fatalf( + "IsRetryableTransaction() = %v, want %v", + got, + test.want, + ) + } + }) } } diff --git a/examples/observability/go.mod b/examples/observability/go.mod index 190c8cb..0176b64 100644 --- a/examples/observability/go.mod +++ b/examples/observability/go.mod @@ -10,9 +10,9 @@ require ( github.com/mkbeh/xpg/extra/slogxpg v0.1.0 github.com/prometheus/client_golang v1.24.1 go.opentelemetry.io/otel v1.45.0 - go.opentelemetry.io/otel/exporters/prometheus v0.67.0 - go.opentelemetry.io/otel/exporters/stdout/stdouttrace v1.45.0 - go.opentelemetry.io/otel/sdk v1.45.0 - go.opentelemetry.io/otel/sdk/metric v1.45.0 + go.opentelemetry.io/otel/exporters/prometheus v0.68.0 + go.opentelemetry.io/otel/exporters/stdout/stdouttrace v1.46.0 + go.opentelemetry.io/otel/sdk v1.46.0 + go.opentelemetry.io/otel/sdk/metric v1.46.0 go.opentelemetry.io/otel/trace v1.45.0 ) diff --git a/extra/otelxpg/go.mod b/extra/otelxpg/go.mod index 859994a..38c5329 100644 --- a/extra/otelxpg/go.mod +++ b/extra/otelxpg/go.mod @@ -4,6 +4,6 @@ go 1.27 require ( github.com/mkbeh/xpg v0.2.0 - go.opentelemetry.io/otel v1.45.0 - go.opentelemetry.io/otel/metric v1.45.0 + go.opentelemetry.io/otel v1.46.0 + go.opentelemetry.io/otel/metric v1.46.0 ) diff --git a/extra/otelxpg/metrics.go b/extra/otelxpg/metrics.go index 7ae85ec..7427388 100644 --- a/extra/otelxpg/metrics.go +++ b/extra/otelxpg/metrics.go @@ -10,15 +10,14 @@ import ( const instrumentationName = "github.com/mkbeh/xpg/extra/otelxpg" -// Metrics exports xpg statistics through OpenTelemetry. +// Metrics exports xpg pool statistics through OpenTelemetry. // -// Metrics is immutable after construction and may be reused for multiple -// pool registrations. +// Metrics is safe for concurrent use and reuse across multiple pools. type Metrics struct { meterProvider metric.MeterProvider } -// metricsRegistration owns one OpenTelemetry callback registration. +// metricsRegistration represents one OpenTelemetry callback registration. type metricsRegistration struct { registration metric.Registration closeOnce sync.Once @@ -31,9 +30,7 @@ func (m *metricsRegistration) Close() { m.closeOnce.Do(func() { if err := m.registration.Unregister(); err != nil { - otel.Handle( - fmt.Errorf("otelxpg: unregister metrics: %w", err), - ) + otel.Handle(fmt.Errorf("otelxpg: unregister metrics: %w", err)) } }) } diff --git a/extra/otelxpg/metrics_test.go b/extra/otelxpg/metrics_test.go new file mode 100644 index 0000000..4a50451 --- /dev/null +++ b/extra/otelxpg/metrics_test.go @@ -0,0 +1,508 @@ +package otelxpg + +import ( + "slices" + "testing" + "time" + + "github.com/jackc/pgx/v5/pgxpool" + "github.com/mkbeh/xpg" + "go.opentelemetry.io/otel/attribute" + "go.opentelemetry.io/otel/metric" + "go.opentelemetry.io/otel/metric/noop" +) + +func TestMetricsRegistration(t *testing.T) { + t.Parallel() + + metrics := NewMetrics( + WithMeterProvider(noop.NewMeterProvider()), + ) + + pool := newTestPool(t, metrics) + + pool.Close() + pool.Close() +} + +func TestMetricsRegisterNilReceiver(t *testing.T) { + t.Parallel() + + var metrics *Metrics + + registration, err := metrics.Register(nil) + if err == nil { + t.Fatal("expected error") + } + + if registration != nil { + t.Fatal("expected nil registration") + } +} + +func TestMetricsRegisterNilPool(t *testing.T) { + t.Parallel() + + metrics := NewMetrics( + WithMeterProvider(noop.NewMeterProvider()), + ) + + registration, err := metrics.Register(nil) + if err == nil { + t.Fatal("expected error") + } + + if registration != nil { + t.Fatal("expected nil registration") + } +} + +func TestWithMeterProviderNilUsesGlobalProvider(t *testing.T) { + metrics := NewMetrics( + WithMeterProvider(nil), + ) + + pool := newTestPool(t, metrics) + pool.Close() +} + +func TestPoolMetricInstruments(t *testing.T) { + t.Parallel() + + meter := &recordingMeter{} + + instruments, err := newPoolMetricInstruments(meter) + if err != nil { + t.Fatalf("create instruments: %v", err) + } + + want := []instrumentSpec{ + { + name: connectionCountMetricName, + kind: "int64_observable_up_down_counter", + description: "The number of connections currently in the state described by db.client.connection.state.", + unit: "{connection}", + }, + { + name: connectionMaxMetricName, + kind: "int64_observable_up_down_counter", + description: "The maximum number of open connections allowed by the pool.", + unit: "{connection}", + }, + { + name: connectionConstructingMetricName, + kind: "int64_observable_gauge", + description: "The number of connections currently being created by the pool.", + unit: "{connection}", + }, + { + name: connectionAcquireCountMetricName, + kind: "int64_observable_counter", + description: "The cumulative number of successful connection acquires.", + unit: "{request}", + }, + { + name: connectionAcquireTimeMetricName, + kind: "float64_observable_counter", + description: "The cumulative time spent acquiring connections.", + unit: "s", + }, + { + name: connectionAcquireCanceledCountMetricName, + kind: "int64_observable_counter", + description: "The cumulative number of connection acquires canceled by context.", + unit: "{request}", + }, + { + name: connectionAcquireEmptyCountMetricName, + kind: "int64_observable_counter", + description: "The cumulative number of successful acquires that waited because the pool was empty.", + unit: "{request}", + }, + { + name: connectionAcquireEmptyWaitTimeMetricName, + kind: "float64_observable_counter", + description: "The cumulative time spent waiting for a connection while the pool was empty.", + unit: "s", + }, + { + name: connectionCreateCountMetricName, + kind: "int64_observable_counter", + description: "The cumulative number of connections created by the pool.", + unit: "{connection}", + }, + { + name: connectionDestroyCountMetricName, + kind: "int64_observable_counter", + description: "The cumulative number of connections destroyed by pool lifecycle limits.", + unit: "{connection}", + }, + } + + if !slices.Equal(meter.instruments, want) { + t.Fatalf( + "instruments = %#v, want %#v", + meter.instruments, + want, + ) + } + + if got, want := len(instruments.observables()), len(want); got != want { + t.Fatalf("observable count = %d, want %d", got, want) + } +} + +func TestPoolMetricAttributes(t *testing.T) { + t.Parallel() + + attributes := newPoolMetricAttributes( + "test-pool", + map[string]string{ + "region": "eu", + dbSystemNameAttribute: "custom-system", + poolNameAttribute: "custom-pool", + connectionStateAttribute: "custom-state", + destroyReasonAttribute: "custom-reason", + }, + ) + + base := observeAttributes(attributes.base) + + requireStringAttribute(t, base, dbSystemNameAttribute, dbSystemPostgreSQL) + requireStringAttribute(t, base, poolNameAttribute, "test-pool") + requireStringAttribute(t, base, "region", "eu") + requireStringAttribute(t, base, connectionStateAttribute, "custom-state") + requireStringAttribute(t, base, destroyReasonAttribute, "custom-reason") + + idle := observeAttributes(attributes.idle) + requireStringAttribute(t, idle, connectionStateAttribute, connectionStateIdle) + + used := observeAttributes(attributes.used) + requireStringAttribute(t, used, connectionStateAttribute, connectionStateUsed) + + destroyedIdle := observeAttributes(attributes.destroyedIdle) + requireStringAttribute(t, destroyedIdle, destroyReasonAttribute, destroyReasonIdleTimeout) + + destroyedLifetime := observeAttributes(attributes.destroyedLifetime) + requireStringAttribute(t, destroyedLifetime, destroyReasonAttribute, destroyReasonLifetime) +} + +func TestPoolMetricInstrumentsObserve(t *testing.T) { + t.Parallel() + + stats := xpg.PoolStats{ + AcquiredConns: 1, + ConstructingConns: 2, + IdleConns: 3, + MaxConns: 4, + AcquireCount: 5, + AcquireDuration: 6 * time.Second, + CanceledAcquireCount: 7, + EmptyAcquireCount: 8, + EmptyAcquireWaitTime: 9 * time.Second, + NewConnsCount: 10, + MaxIdleDestroyCount: 11, + MaxLifetimeDestroyCount: 12, + } + + attributes := newPoolMetricAttributes( + "test-pool", + map[string]string{ + "region": "eu", + }, + ) + + observer := &recordingObserver{} + + var instruments poolMetricInstruments + + instruments.observe( + observer, + stats, + attributes, + ) + + wantInt64 := []int64{ + 3, + 1, + 4, + 2, + 5, + 7, + 8, + 10, + 11, + 12, + } + + if !slices.Equal(observer.int64Values, wantInt64) { + t.Fatalf( + "int64 observations = %v, want %v", + observer.int64Values, + wantInt64, + ) + } + + wantFloat64 := []float64{ + 6, + 9, + } + + if !slices.Equal(observer.float64Values, wantFloat64) { + t.Fatalf( + "float64 observations = %v, want %v", + observer.float64Values, + wantFloat64, + ) + } + + requireStringAttribute( + t, + observer.int64Attributes[0], + connectionStateAttribute, + connectionStateIdle, + ) + requireStringAttribute( + t, + observer.int64Attributes[1], + connectionStateAttribute, + connectionStateUsed, + ) + requireStringAttribute( + t, + observer.int64Attributes[8], + destroyReasonAttribute, + destroyReasonIdleTimeout, + ) + requireStringAttribute( + t, + observer.int64Attributes[9], + destroyReasonAttribute, + destroyReasonLifetime, + ) + + for _, attributes := range observer.int64Attributes { + requireStringAttribute(t, attributes, poolNameAttribute, "test-pool") + requireStringAttribute(t, attributes, "region", "eu") + } + + for _, attributes := range observer.float64Attributes { + requireStringAttribute(t, attributes, poolNameAttribute, "test-pool") + requireStringAttribute(t, attributes, "region", "eu") + } +} + +func TestMetricsRegistrationCloseOnce(t *testing.T) { + t.Parallel() + + registration := ®istrationStub{} + + metricsRegistration := &metricsRegistration{ + registration: registration, + } + + metricsRegistration.Close() + metricsRegistration.Close() + + if registration.calls != 1 { + t.Fatalf( + "unregister calls = %d, want 1", + registration.calls, + ) + } +} + +func newTestPool( + t *testing.T, + metrics xpg.Metrics, +) *xpg.Pool { + t.Helper() + + poolConfig, err := pgxpool.ParseConfig("") + if err != nil { + t.Fatalf("parse pool config: %v", err) + } + + pool, err := xpg.New( + t.Context(), + poolConfig, + xpg.WithName("test-pool"), + xpg.WithMetrics(metrics), + ) + if err != nil { + t.Fatalf("create pool: %v", err) + } + + return pool +} + +func observeAttributes(option metric.ObserveOption) attribute.Set { + return metric.NewObserveConfig( + []metric.ObserveOption{option}, + ).Attributes() +} + +func requireStringAttribute( + t *testing.T, + attributes attribute.Set, + key string, + want string, +) { + t.Helper() + + value, ok := attributes.Value(attribute.Key(key)) + if !ok { + t.Fatalf("attribute %q is missing", key) + } + + if got := value.AsString(); got != want { + t.Fatalf( + "attribute %q = %q, want %q", + key, + got, + want, + ) + } +} + +type instrumentSpec struct { + name string + kind string + description string + unit string +} + +type recordingMeter struct { + noop.Meter + + instruments []instrumentSpec +} + +func (m *recordingMeter) Int64ObservableUpDownCounter( + name string, + options ...metric.Int64ObservableUpDownCounterOption, +) (metric.Int64ObservableUpDownCounter, error) { + config := metric.NewInt64ObservableUpDownCounterConfig(options...) + + m.instruments = append( + m.instruments, + instrumentSpec{ + name: name, + kind: "int64_observable_up_down_counter", + description: config.Description(), + unit: config.Unit(), + }, + ) + + return m.Meter.Int64ObservableUpDownCounter(name, options...) +} + +func (m *recordingMeter) Int64ObservableGauge( + name string, + options ...metric.Int64ObservableGaugeOption, +) (metric.Int64ObservableGauge, error) { + config := metric.NewInt64ObservableGaugeConfig(options...) + + m.instruments = append( + m.instruments, + instrumentSpec{ + name: name, + kind: "int64_observable_gauge", + description: config.Description(), + unit: config.Unit(), + }, + ) + + return m.Meter.Int64ObservableGauge(name, options...) +} + +func (m *recordingMeter) Int64ObservableCounter( + name string, + options ...metric.Int64ObservableCounterOption, +) (metric.Int64ObservableCounter, error) { + config := metric.NewInt64ObservableCounterConfig(options...) + + m.instruments = append( + m.instruments, + instrumentSpec{ + name: name, + kind: "int64_observable_counter", + description: config.Description(), + unit: config.Unit(), + }, + ) + + return m.Meter.Int64ObservableCounter(name, options...) +} + +func (m *recordingMeter) Float64ObservableCounter( + name string, + options ...metric.Float64ObservableCounterOption, +) (metric.Float64ObservableCounter, error) { + config := metric.NewFloat64ObservableCounterConfig(options...) + + m.instruments = append( + m.instruments, + instrumentSpec{ + name: name, + kind: "float64_observable_counter", + description: config.Description(), + unit: config.Unit(), + }, + ) + + return m.Meter.Float64ObservableCounter(name, options...) +} + +type recordingObserver struct { + noop.Observer + + int64Values []int64 + int64Attributes []attribute.Set + float64Values []float64 + float64Attributes []attribute.Set +} + +func (o *recordingObserver) ObserveInt64( + _ metric.Int64Observable, + value int64, + options ...metric.ObserveOption, +) { + o.int64Values = append( + o.int64Values, + value, + ) + + o.int64Attributes = append( + o.int64Attributes, + metric.NewObserveConfig(options).Attributes(), + ) +} + +func (o *recordingObserver) ObserveFloat64( + _ metric.Float64Observable, + value float64, + options ...metric.ObserveOption, +) { + o.float64Values = append( + o.float64Values, + value, + ) + + o.float64Attributes = append( + o.float64Attributes, + metric.NewObserveConfig(options).Attributes(), + ) +} + +type registrationStub struct { + noop.Registration + + calls int +} + +func (r *registrationStub) Unregister() error { + r.calls++ + + return nil +} diff --git a/extra/otelxpg/options.go b/extra/otelxpg/options.go index b35f0f6..786960f 100644 --- a/extra/otelxpg/options.go +++ b/extra/otelxpg/options.go @@ -11,20 +11,11 @@ type MetricsOption interface { apply(*metricsSettings) } -type metricsOptionFunc func(*metricsSettings) - -func (option metricsOptionFunc) apply(settings *metricsSettings) { - option(settings) -} - -type metricsSettings struct { - meterProvider metric.MeterProvider -} - -// NewMetrics creates an OpenTelemetry metrics implementation. +// NewMetrics creates an OpenTelemetry metrics integration. // -// By default, metrics use the global OpenTelemetry MeterProvider. The returned -// value is immutable and may be reused for multiple pools. +// When no MeterProvider is configured, the global OpenTelemetry MeterProvider +// is used. The returned value is safe for concurrent use and reuse across +// multiple pools. func NewMetrics(options ...MetricsOption) *Metrics { settings := metricsSettings{} @@ -43,8 +34,9 @@ func NewMetrics(options ...MetricsOption) *Metrics { // WithMeterProvider configures the MeterProvider used for metrics. // -// The caller owns the provider and must shut it down after all instrumented -// pools have been closed. +// A nil provider leaves the global OpenTelemetry MeterProvider in use. The +// caller retains ownership of a non-nil provider and is responsible for +// shutting it down after all instrumented pools have been closed. func WithMeterProvider(provider metric.MeterProvider) MetricsOption { return metricsOptionFunc(func(settings *metricsSettings) { if provider != nil { @@ -52,3 +44,13 @@ func WithMeterProvider(provider metric.MeterProvider) MetricsOption { } }) } + +type metricsOptionFunc func(*metricsSettings) + +func (option metricsOptionFunc) apply(settings *metricsSettings) { + option(settings) +} + +type metricsSettings struct { + meterProvider metric.MeterProvider +} diff --git a/extra/otelxpg/options_test.go b/extra/otelxpg/options_test.go index 47b213d..df39469 100644 --- a/extra/otelxpg/options_test.go +++ b/extra/otelxpg/options_test.go @@ -1,54 +1,66 @@ package otelxpg import ( - "context" "testing" - "github.com/jackc/pgx/v5/pgxpool" - "github.com/mkbeh/xpg" "go.opentelemetry.io/otel/metric/noop" ) -func TestMetricsRegistration(t *testing.T) { +func TestNewMetrics(t *testing.T) { t.Parallel() + metrics := NewMetrics(nil) + + if metrics == nil { + t.Fatal("expected metrics") + } + + if metrics.meterProvider != nil { + t.Fatal("expected nil meter provider") + } +} + +func TestWithMeterProvider(t *testing.T) { + t.Parallel() + + provider := noop.NewMeterProvider() + metrics := NewMetrics( - WithMeterProvider(noop.NewMeterProvider()), + WithMeterProvider(provider), ) - pool := newTestPool(t, metrics) - pool.Close() - pool.Close() + if metrics.meterProvider != provider { + t.Fatal("unexpected meter provider") + } } -func TestWithMeterProviderNilUsesGlobalProvider(t *testing.T) { +func TestWithMeterProviderLastWins(t *testing.T) { t.Parallel() + first := noop.NewMeterProvider() + second := noop.NewMeterProvider() + metrics := NewMetrics( - WithMeterProvider(nil), + WithMeterProvider(first), + WithMeterProvider(second), ) - pool := newTestPool(t, metrics) - pool.Close() + if metrics.meterProvider != second { + t.Fatal("expected last meter provider to win") + } } -func newTestPool(t *testing.T, metrics xpg.Metrics) *xpg.Pool { - t.Helper() +func TestWithMeterProviderNilIgnored(t *testing.T) { + t.Parallel() - poolConfig, err := pgxpool.ParseConfig("") - if err != nil { - t.Fatalf("parse pool config: %v", err) - } + provider := noop.NewMeterProvider() - pool, err := xpg.New( - context.Background(), - poolConfig, - xpg.WithName("test-pool"), - xpg.WithMetrics(metrics), + metrics := NewMetrics( + WithMeterProvider(provider), + WithMeterProvider(nil), ) - if err != nil { - t.Fatalf("create pool: %v", err) - } - return pool + if metrics.meterProvider != provider { + t.Fatal("expected nil meter provider option to be ignored") + } } diff --git a/extra/otelxpg/pool.go b/extra/otelxpg/pool.go index f1bc29c..634de96 100644 --- a/extra/otelxpg/pool.go +++ b/extra/otelxpg/pool.go @@ -65,20 +65,22 @@ type poolMetricAttributes struct { destroyedLifetime metric.ObserveOption } -var _ xpg.Metrics = (*Metrics)(nil) - // Register registers metrics for one xpg Pool. func (m *Metrics) Register(pool *xpg.Pool) (xpg.MetricsRegistration, error) { if m == nil { return nil, errors.New("otelxpg: metrics is nil") } - provider := m.meterProvider - if provider == nil { - provider = otel.GetMeterProvider() + if pool == nil { + return nil, errors.New("otelxpg: pool is nil") + } + + meterProvider := m.meterProvider + if meterProvider == nil { + meterProvider = otel.GetMeterProvider() } - return registerPoolMetrics(pool, provider) + return registerPoolMetrics(pool, meterProvider) } func registerPoolMetrics(pool *xpg.Pool, provider metric.MeterProvider) (xpg.MetricsRegistration, error) { @@ -93,12 +95,7 @@ func registerPoolMetrics(pool *xpg.Pool, provider metric.MeterProvider) (xpg.Met registration, err := meter.RegisterCallback( func(_ context.Context, observer metric.Observer) error { - instruments.observe( - observer, - pool.Stats(), - attributes, - ) - + instruments.observe(observer, pool.Stats(), attributes) return nil }, instruments.observables()..., @@ -112,100 +109,9 @@ func registerPoolMetrics(pool *xpg.Pool, provider metric.MeterProvider) (xpg.Met }, nil } -func (i poolMetricInstruments) observe( - observer metric.Observer, - stats xpg.PoolStats, - attributes poolMetricAttributes, -) { - observer.ObserveInt64( - i.connectionCount, - int64(stats.IdleConns), - attributes.idle, - ) - - observer.ObserveInt64( - i.connectionCount, - int64(stats.AcquiredConns), - attributes.used, - ) - - observer.ObserveInt64( - i.connectionMax, - int64(stats.MaxConns), - attributes.base, - ) - - observer.ObserveInt64( - i.constructingConnections, - int64(stats.ConstructingConns), - attributes.base, - ) - - observer.ObserveInt64( - i.acquireCount, - stats.AcquireCount, - attributes.base, - ) - - observer.ObserveFloat64( - i.acquireTime, - stats.AcquireDuration.Seconds(), - attributes.base, - ) - - observer.ObserveInt64( - i.canceledAcquireCount, - stats.CanceledAcquireCount, - attributes.base, - ) - - observer.ObserveInt64( - i.emptyAcquireCount, - stats.EmptyAcquireCount, - attributes.base, - ) - - observer.ObserveFloat64( - i.emptyAcquireWaitTime, - stats.EmptyAcquireWaitTime.Seconds(), - attributes.base, - ) - - observer.ObserveInt64( - i.createdConnections, - stats.NewConnsCount, - attributes.base, - ) - - observer.ObserveInt64( - i.destroyedConnections, - stats.MaxIdleDestroyCount, - attributes.destroyedIdle, - ) - - observer.ObserveInt64( - i.destroyedConnections, - stats.MaxLifetimeDestroyCount, - attributes.destroyedLifetime, - ) -} - -func (i poolMetricInstruments) observables() []metric.Observable { - return []metric.Observable{ - i.connectionCount, - i.connectionMax, - i.constructingConnections, - i.acquireCount, - i.acquireTime, - i.canceledAcquireCount, - i.emptyAcquireCount, - i.emptyAcquireWaitTime, - i.createdConnections, - i.destroyedConnections, - } -} - -func newPoolMetricInstruments(meter metric.Meter) (poolMetricInstruments, error) { +func newPoolMetricInstruments( + meter metric.Meter, +) (poolMetricInstruments, error) { var instruments poolMetricInstruments var err error @@ -213,7 +119,7 @@ func newPoolMetricInstruments(meter metric.Meter) (poolMetricInstruments, error) instruments.connectionCount, err = meter.Int64ObservableUpDownCounter( connectionCountMetricName, metric.WithDescription( - "The number of connections currently used or idle in the pool.", + "The number of connections currently in the state described by db.client.connection.state.", ), metric.WithUnit("{connection}"), ) @@ -333,7 +239,7 @@ func newPoolMetricInstruments(meter metric.Meter) (poolMetricInstruments, error) instruments.createdConnections, err = meter.Int64ObservableCounter( connectionCreateCountMetricName, metric.WithDescription( - "The cumulative number of connections opened by the pool.", + "The cumulative number of connections created by the pool.", ), metric.WithUnit("{connection}"), ) @@ -367,10 +273,7 @@ func newPoolMetricAttributes(name string, labels map[string]string) poolMetricAt var base []attribute.KeyValue for key, value := range labels { - base = append( - base, - attribute.String(key, value), - ) + base = append(base, attribute.String(key, value)) } // System attributes are appended last, so xpg-controlled values win when @@ -427,3 +330,96 @@ func newPoolMetricAttributes(name string, labels map[string]string) poolMetricAt ), } } + +func (i poolMetricInstruments) observe( + observer metric.Observer, + stats xpg.PoolStats, + attributes poolMetricAttributes, +) { + observer.ObserveInt64( + i.connectionCount, + int64(stats.IdleConns), + attributes.idle, + ) + + observer.ObserveInt64( + i.connectionCount, + int64(stats.AcquiredConns), + attributes.used, + ) + + observer.ObserveInt64( + i.connectionMax, + int64(stats.MaxConns), + attributes.base, + ) + + observer.ObserveInt64( + i.constructingConnections, + int64(stats.ConstructingConns), + attributes.base, + ) + + observer.ObserveInt64( + i.acquireCount, + stats.AcquireCount, + attributes.base, + ) + + observer.ObserveFloat64( + i.acquireTime, + stats.AcquireDuration.Seconds(), + attributes.base, + ) + + observer.ObserveInt64( + i.canceledAcquireCount, + stats.CanceledAcquireCount, + attributes.base, + ) + + observer.ObserveInt64( + i.emptyAcquireCount, + stats.EmptyAcquireCount, + attributes.base, + ) + + observer.ObserveFloat64( + i.emptyAcquireWaitTime, + stats.EmptyAcquireWaitTime.Seconds(), + attributes.base, + ) + + observer.ObserveInt64( + i.createdConnections, + stats.NewConnsCount, + attributes.base, + ) + + observer.ObserveInt64( + i.destroyedConnections, + stats.MaxIdleDestroyCount, + attributes.destroyedIdle, + ) + + observer.ObserveInt64( + i.destroyedConnections, + stats.MaxLifetimeDestroyCount, + attributes.destroyedLifetime, + ) +} + +func (i poolMetricInstruments) observables() []metric.Observable { + return []metric.Observable{ + i.connectionCount, + i.connectionMax, + i.constructingConnections, + i.acquireCount, + i.acquireTime, + i.canceledAcquireCount, + i.emptyAcquireCount, + i.emptyAcquireWaitTime, + i.createdConnections, + i.destroyedConnections, + } +} diff --git a/extra/slogxpg/logger_test.go b/extra/slogxpg/logger_test.go index cd8cbab..ccc8350 100644 --- a/extra/slogxpg/logger_test.go +++ b/extra/slogxpg/logger_test.go @@ -67,7 +67,7 @@ func TestSlogLevel(t *testing.T) { } } -func TestLoggerLog(t *testing.T) { +func TestAdapterLog(t *testing.T) { t.Parallel() handler := &captureHandler{} @@ -104,7 +104,7 @@ func TestLoggerLog(t *testing.T) { } } -func TestLoggerUnknownLevel(t *testing.T) { +func TestAdapterUnknownLevel(t *testing.T) { t.Parallel() handler := &captureHandler{} diff --git a/go.sum b/go.sum index 31890d5..a277497 100644 --- a/go.sum +++ b/go.sum @@ -1,6 +1,4 @@ github.com/davecgh/go-spew v1.1.0/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38= -github.com/davecgh/go-spew v1.1.1 h1:vj9j/u1bqnvCEfJOwUhtlOARqs3+rkHYY13jYWTU97c= -github.com/davecgh/go-spew v1.1.1/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38= github.com/jackc/pgpassfile v1.0.0 h1:/6Hmqy13Ss2zCq62VdNG8tM1wchn8zjSGOBJ6icpsIM= github.com/jackc/pgpassfile v1.0.0/go.mod h1:CEx0iS5ambNFdcRtxPj5JhEz+xB6uRky5eyVu/W2HEg= github.com/jackc/pgservicefile v0.0.0-20240606120523-5a60cdf6a761 h1:iCEnooe7UlwOQYpKFhBabPMi4aNAfoODPEFNiAnClxo= @@ -9,18 +7,15 @@ github.com/jackc/pgx/v5 v5.10.0 h1:VhSvgU2jSli8o3AqIEOTJr7rZwAEUVo4E4XhR94Zfr0= github.com/jackc/pgx/v5 v5.10.0/go.mod h1:mal1tBGAFfLHvZzaYh77YS/eC6IX9OWbRV1QIIM0Jn4= github.com/jackc/puddle/v2 v2.2.2 h1:PR8nw+E/1w0GLuRFSmiioY6UooMp6KJv0/61nB7icHo= github.com/jackc/puddle/v2 v2.2.2/go.mod h1:vriiEXHvEE654aYKXXjOvZM39qJ0q+azkZFrfEOc3H4= -github.com/pmezard/go-difflib v1.0.0 h1:4DBwDE0NGyQoBHbLQYPwSUPoCMWR5BEzIk/f1lZbAQM= github.com/pmezard/go-difflib v1.0.0/go.mod h1:iKH77koFhYxTK1pcRnkKkqfTogsbg7gZNVY4sRDYZ/4= github.com/stretchr/objx v0.1.0/go.mod h1:HFkY916IF+rwdDfMAkV7OtwuqBVzrE8GR6GFx+wExME= github.com/stretchr/testify v1.3.0/go.mod h1:M5WIy9Dh21IEIfnGCwXGc5bZfKNJtfHm1UVUgZn+9EI= github.com/stretchr/testify v1.7.0/go.mod h1:6Fq8oRcR53rry900zMqJjRRixrwX3KX962/h/Wwjteg= -github.com/stretchr/testify v1.11.1 h1:7s2iGBzp5EwR7/aIZr8ao5+dra3wiQyKjjFuvgVKu7U= -github.com/stretchr/testify v1.11.1/go.mod h1:wZwfW3scLgRK+23gO65QZefKpKQRnfz6sD981Nm4B6U= +github.com/stretchr/testify v1.12.1 h1:EuwCh5fleGS7H32xRwO3wRGT7DxrDhLAT6FF8MpWDWE= +go.yaml.in/yaml/v3 v3.0.5 h1:N6y/pJk8buWs9NY5ERU2HSMfm+IuD/OtfdAnq6kESPw= golang.org/x/sync v0.22.0 h1:SZjpbeLmrCk4xhRSZFNZW5gFUeCeFgjekvI/+gfScek= golang.org/x/sync v0.22.0/go.mod h1:9xrNwdLfx4jkKbNva9FpL6vEN7evnE43NNNJQ2LF3+0= golang.org/x/text v0.41.0 h1:vz/seA0lnX87Othu2f/0L24RcgrXD9/YFTSuGjj3rH8= golang.org/x/text v0.41.0/go.mod h1:jvf1O8ajNzZqhSrQBPbutR/EB83Cc0CFrezNQIwbb5M= gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0= gopkg.in/yaml.v3 v3.0.0-20200313102051-9f266ea9e77c/go.mod h1:K4uyk7z7BCEPqu6E+C64Yfv1cQ7kz7rIZviUmN+EgEM= -gopkg.in/yaml.v3 v3.0.1 h1:fxVm/GzAzEWqLHuvctI91KS9hhNmmWOoWu0XTYJS7CA= -gopkg.in/yaml.v3 v3.0.1/go.mod h1:K4uyk7z7BCEPqu6E+C64Yfv1cQ7kz7rIZviUmN+EgEM= diff --git a/metrics.go b/metrics.go index be69822..3d08aa8 100644 --- a/metrics.go +++ b/metrics.go @@ -2,14 +2,15 @@ package xpg // Metrics registers metrics for a Pool. // -// Implementations are expected to be immutable and safe to reuse for multiple -// pools. Register is called after the underlying pgxpool.Pool has been created. +// Implementations must be safe to reuse across multiple pools. Register is +// called after the underlying pgxpool.Pool has been created. type Metrics interface { Register(pool *Pool) (MetricsRegistration, error) } -// MetricsRegistration owns a metrics registration associated with one -// Pool. Close is called once before the underlying pgxpool.Pool is closed. +// MetricsRegistration represents a metrics registration for one Pool. +// +// Close is called once before the underlying pgxpool.Pool is closed. type MetricsRegistration interface { Close() } diff --git a/options.go b/options.go index 6f44ed1..c41dbd9 100644 --- a/options.go +++ b/options.go @@ -20,63 +20,6 @@ type Option interface { apply(*settings) error } -type optionFunc func(*settings) error - -func (option optionFunc) apply(settings *settings) error { - return option(settings) -} - -type settings struct { - name string - labels map[string]string - metrics Metrics - tracers []pgx.QueryTracer -} - -func (s *settings) poolName(host string, port uint16, database string) string { - if s.name != "" { - return s.name - } - - address := net.JoinHostPort( - host, - strconv.Itoa(int(port)), - ) - if database == "" { - return address - } - - return address + "/" + database -} - -func (s *settings) buildTracer() pgx.QueryTracer { - if len(s.tracers) == 0 { - return nil - } - - return multitracer.New(s.tracers...) -} - -func defaultSettings() *settings { - return &settings{ - labels: make(map[string]string), - } -} - -func applyOptions(settings *settings, opts ...Option) error { - for _, opt := range opts { - if opt == nil { - return errors.New("xpg: option is nil") - } - - if err := opt.apply(settings); err != nil { - return fmt.Errorf("xpg: apply option: %w", err) - } - } - - return nil -} - // WithName assigns a stable logical name to the pool. // // Name is metadata for diagnostics and observability. It does not change the @@ -95,26 +38,6 @@ func WithName(name string) Option { }) } -// WithLabels merges labels into the pool metadata. -// -// Labels are defensively copied. When the same key is configured more than -// once, the last value wins. -func WithLabels(labels map[string]string) Option { - labels = cloneLabels(labels) - - return optionFunc(func(settings *settings) error { - for key, value := range labels { - if key == "" { - return errors.New("label key must not be empty") - } - - settings.labels[key] = value - } - - return nil - }) -} - // WithLabel adds or replaces one pool label. func WithLabel(key, value string) Option { return optionFunc(func(settings *settings) error { @@ -128,17 +51,21 @@ func WithLabel(key, value string) Option { }) } -// WithMetrics attaches one metrics implementation to the pool. +// WithLabels merges labels into the pool metadata. // -// Metrics are registered during New and unregistered automatically when the -// Pool is closed. -func WithMetrics(metrics Metrics) Option { +// Labels are defensively copied. When the same key is configured more than +// once, the last value wins. +func WithLabels(labels map[string]string) Option { + labels = cloneLabels(labels) + return optionFunc(func(settings *settings) error { - if metrics == nil { - return errors.New("pool metrics is nil") - } + for key, value := range labels { + if key == "" { + return errors.New("label key must not be empty") + } - settings.metrics = metrics + settings.labels[key] = value + } return nil }) @@ -146,11 +73,10 @@ func WithMetrics(metrics Metrics) Option { // WithLogger attaches a pgx-compatible logger to the pool. // -// Logging is implemented through pgx tracelog and participates in the same -// tracing pipeline as custom tracers. If other tracers are configured, xpg -// combines them automatically with pgx multitracer. pgx tracelog may include -// SQL text and query arguments in log records; applications are responsible for -// choosing an appropriate level and handling sensitive values. +// Logging uses pgx tracelog and participates in the same tracing pipeline as +// tracers configured through xpg. pgx tracelog may include SQL text and query +// arguments in log records; applications are responsible for choosing an +// appropriate level and handling sensitive values. func WithLogger(logger tracelog.Logger, level tracelog.LogLevel) Option { return optionFunc(func(settings *settings) error { if logger == nil { @@ -166,18 +92,22 @@ func WithLogger(logger tracelog.Logger, level tracelog.LogLevel) Option { }) } -// WithTracer attaches a pgx query tracer to the pool. +// WithTracer attaches one pgx query tracer to the pool. // -// The option may be specified multiple times. Tracers are invoked in the order -// they are configured, after any tracer already present in -// config.ConnConfig.Tracer. When more than one tracer is present, xpg combines -// them with pgx multitracer. Additional pgx tracing capabilities implemented by -// a tracer, such as batch, copy, prepare, connect, acquire, and release tracing, -// are preserved by multitracer. +// The option may be specified multiple times. Configured loggers and tracers +// are combined through pgx multitracer. When xpg logging or tracing options are +// configured, the resulting tracing pipeline replaces any tracer already +// configured on the pgx connection config. func WithTracer(tracer pgx.QueryTracer) Option { return WithTracers(tracer) } +// WithTracers attaches multiple pgx query tracers to the pool. +// +// Configured loggers and tracers are invoked in configuration order and +// combined through pgx multitracer. When xpg logging or tracing options are +// configured, the resulting tracing pipeline replaces any tracer already +// configured on the pgx connection config. func WithTracers(tracers ...pgx.QueryTracer) Option { return optionFunc(func(settings *settings) error { for _, tracer := range tracers { @@ -192,6 +122,79 @@ func WithTracers(tracers ...pgx.QueryTracer) Option { }) } +// WithMetrics attaches one metrics implementation to the pool. +// +// Metrics are registered when the pool is created and unregistered +// automatically when the Pool is closed. +func WithMetrics(metrics Metrics) Option { + return optionFunc(func(settings *settings) error { + if metrics == nil { + return errors.New("pool metrics is nil") + } + + settings.metrics = metrics + + return nil + }) +} + +type optionFunc func(*settings) error + +func (option optionFunc) apply(settings *settings) error { + return option(settings) +} + +type settings struct { + name string + labels map[string]string + metrics Metrics + tracers []pgx.QueryTracer +} + +func defaultSettings() *settings { + return &settings{ + labels: make(map[string]string), + } +} + +func applyOptions(settings *settings, opts ...Option) error { + for _, opt := range opts { + if opt == nil { + return errors.New("xpg: option is nil") + } + + if err := opt.apply(settings); err != nil { + return fmt.Errorf("xpg: apply option: %w", err) + } + } + + return nil +} + +func (s *settings) poolName(host string, port uint16, database string) string { + if s.name != "" { + return s.name + } + + address := net.JoinHostPort( + host, + strconv.Itoa(int(port)), + ) + if database == "" { + return address + } + + return address + "/" + database +} + +func (s *settings) buildTracer() pgx.QueryTracer { + if len(s.tracers) == 0 { + return nil + } + + return multitracer.New(s.tracers...) +} + func cloneLabels(labels map[string]string) map[string]string { if len(labels) == 0 { return nil diff --git a/pool.go b/pool.go index 84317ef..72d5a29 100644 --- a/pool.go +++ b/pool.go @@ -22,9 +22,9 @@ type Pool struct { closeOnce sync.Once } -// Open parses a DSN and creates a Pool. -func Open(ctx context.Context, dsn string, options ...Option) (*Pool, error) { - config, err := pgxpool.ParseConfig(dsn) +// Open parses a PostgreSQL connection string and creates a Pool. +func Open(ctx context.Context, connString string, options ...Option) (*Pool, error) { + config, err := pgxpool.ParseConfig(connString) if err != nil { return nil, fmt.Errorf("xpg: parse pool config: %w", err) } @@ -35,8 +35,7 @@ func Open(ctx context.Context, dsn string, options ...Option) (*Pool, error) { // New creates a Pool from config. // // Config must have been created by pgxpool.ParseConfig. New passes a defensive -// copy to pgxpool, so subsequent changes to the original config do not affect -// the created Pool. +// copy to pgxpool, so subsequent changes to config do not affect the Pool. // // As with pgxpool.Config.Copy, the referenced tls.Config remains shared and // must not be modified after it has been used to create connections. @@ -75,32 +74,40 @@ func New(ctx context.Context, config *pgxpool.Config, options ...Option) (*Pool, if err := pool.registerMetrics(settings.metrics); err != nil { pool.Close() + return nil, fmt.Errorf("xpg: register pool metrics: %w", err) } return pool, nil } -// Name returns the logical pool name configured with WithName. +// Name returns the logical pool name. +// +// If WithName is not configured, the name is derived from the connection host, +// port, and database. func (p *Pool) Name() string { return p.name } +// Labels returns a copy of the pool labels. func (p *Pool) Labels() map[string]string { return cloneLabels(p.labels) } // Raw returns the underlying pgxpool.Pool. +// +// The returned pool is owned by Pool and must not be closed directly. func (p *Pool) Raw() *pgxpool.Pool { return p.pool } +// Ping verifies connectivity to PostgreSQL. func (p *Pool) Ping(ctx context.Context) error { return p.pool.Ping(ctx) } -// Close closes the underlying pool and waits for acquired connections to be -// returned. Close is safe to call multiple times. +// Close closes the pool and waits for acquired connections to be returned. +// Close is safe to call multiple times. func (p *Pool) Close() { p.closeOnce.Do(func() { if p.metrics != nil { @@ -111,22 +118,27 @@ func (p *Pool) Close() { }) } +// Exec executes SQL against the pool. func (p *Pool) Exec(ctx context.Context, sql string, arguments ...any) (pgconn.CommandTag, error) { return p.pool.Exec(ctx, sql, arguments...) } +// Query executes SQL and returns the resulting rows. func (p *Pool) Query(ctx context.Context, sql string, args ...any) (pgx.Rows, error) { return p.pool.Query(ctx, sql, args...) } +// QueryRow executes SQL that is expected to return at most one row. func (p *Pool) QueryRow(ctx context.Context, sql string, args ...any) pgx.Row { return p.pool.QueryRow(ctx, sql, args...) } +// SendBatch sends a batch of queries through the pool. func (p *Pool) SendBatch(ctx context.Context, batch *pgx.Batch) pgx.BatchResults { return p.pool.SendBatch(ctx, batch) } +// CopyFrom copies rows into the specified table. func (p *Pool) CopyFrom(ctx context.Context, tableName pgx.Identifier, columnNames []string, rowSrc pgx.CopyFromSource) (int64, error) { return p.pool.CopyFrom(ctx, tableName, columnNames, rowSrc) } diff --git a/pool_test.go b/pool_test.go new file mode 100644 index 0000000..b68d91b --- /dev/null +++ b/pool_test.go @@ -0,0 +1,274 @@ +package xpg + +import ( + "errors" + "testing" + + "github.com/jackc/pgx/v5/pgxpool" +) + +func TestNewRejectsNilConfig(t *testing.T) { + t.Parallel() + + tests := []struct { + name string + config *pgxpool.Config + }{ + { + name: "nil config", + config: nil, + }, + { + name: "nil connection config", + config: &pgxpool.Config{}, + }, + } + + for _, test := range tests { + t.Run(test.name, func(t *testing.T) { + t.Parallel() + + _, err := New( + t.Context(), + test.config, + ) + + assertErrorMessage( + t, + err, + "xpg: pool config is nil", + ) + }) + } +} + +func TestPoolMetadata(t *testing.T) { + t.Parallel() + + config := newPoolTestConfig(t) + + labels := map[string]string{ + "region": "eu", + "role": "reader", + } + + labelsOption := WithLabels(labels) + labels["region"] = "changed" + + pool, err := New( + t.Context(), + config, + WithName(" orders "), + labelsOption, + WithLabel("role", "primary"), + ) + if err != nil { + t.Fatalf("create pool: %v", err) + } + defer pool.Close() + + if got, want := pool.Name(), "orders"; got != want { + t.Fatalf( + "Name() = %q, want %q", + got, + want, + ) + } + + gotLabels := pool.Labels() + + if got, want := gotLabels["region"], "eu"; got != want { + t.Fatalf( + "region label = %q, want %q", + got, + want, + ) + } + + if got, want := gotLabels["role"], "primary"; got != want { + t.Fatalf( + "role label = %q, want %q", + got, + want, + ) + } + + gotLabels["region"] = "mutated" + + if got, want := pool.Labels()["region"], "eu"; got != want { + t.Fatalf( + "region label after mutation = %q, want %q", + got, + want, + ) + } + + if pool.Raw() == nil { + t.Fatal("Raw returned nil") + } +} + +func TestPoolDerivedName(t *testing.T) { + t.Parallel() + + pool, err := New( + t.Context(), + newPoolTestConfig(t), + ) + if err != nil { + t.Fatalf("create pool: %v", err) + } + defer pool.Close() + + if got, want := pool.Name(), "localhost:5432/testdb"; got != want { + t.Fatalf( + "Name() = %q, want %q", + got, + want, + ) + } +} + +func TestPoolStats(t *testing.T) { + t.Parallel() + + config := newPoolTestConfig(t) + config.MaxConns = 17 + + pool, err := New( + t.Context(), + config, + ) + if err != nil { + t.Fatalf("create pool: %v", err) + } + defer pool.Close() + + stats := pool.Stats() + + if got, want := stats.MaxConns, int32(17); got != want { + t.Fatalf( + "MaxConns = %d, want %d", + got, + want, + ) + } + + if stats.TotalConns != 0 { + t.Fatalf( + "TotalConns = %d, want 0", + stats.TotalConns, + ) + } +} + +func TestPoolMetricsLifecycle(t *testing.T) { + t.Parallel() + + registration := &poolTestMetricsRegistration{} + metrics := &poolTestMetrics{ + registration: registration, + } + + pool, err := New( + t.Context(), + newPoolTestConfig(t), + WithMetrics(metrics), + ) + if err != nil { + t.Fatalf("create pool: %v", err) + } + + if metrics.registerCalls != 1 { + t.Fatalf( + "register calls = %d, want 1", + metrics.registerCalls, + ) + } + + if metrics.pool != pool { + t.Fatal("metrics registered with unexpected pool") + } + + pool.Close() + pool.Close() + + if registration.closeCalls != 1 { + t.Fatalf( + "registration close calls = %d, want 1", + registration.closeCalls, + ) + } +} + +func TestNewPreservesMetricsRegistrationError(t *testing.T) { + t.Parallel() + + expectedErr := errors.New("register failed") + metrics := &poolTestMetrics{ + err: expectedErr, + } + + pool, err := New( + t.Context(), + newPoolTestConfig(t), + WithMetrics(metrics), + ) + + if pool != nil { + t.Fatal("expected nil pool") + } + + if !errors.Is(err, expectedErr) { + t.Fatalf("original error was not preserved: %v", err) + } + + if metrics.registerCalls != 1 { + t.Fatalf( + "register calls = %d, want 1", + metrics.registerCalls, + ) + } +} + +func newPoolTestConfig(t *testing.T) *pgxpool.Config { + t.Helper() + + config, err := pgxpool.ParseConfig( + "postgres://user:password@localhost:5432/testdb?sslmode=disable", + ) + if err != nil { + t.Fatalf("parse pool config: %v", err) + } + + return config +} + +type poolTestMetrics struct { + registration MetricsRegistration + err error + + pool *Pool + registerCalls int +} + +func (m *poolTestMetrics) Register( + pool *Pool, +) (MetricsRegistration, error) { + m.registerCalls++ + m.pool = pool + + if m.err != nil { + return nil, m.err + } + + return m.registration, nil +} + +type poolTestMetricsRegistration struct { + closeCalls int +} + +func (r *poolTestMetricsRegistration) Close() { + r.closeCalls++ +} diff --git a/shard/doc.go b/shard/doc.go index 9fa129f..aa07ba8 100644 --- a/shard/doc.go +++ b/shard/doc.go @@ -1,15 +1,9 @@ -// Package shard provides explicit application-level routing across PostgreSQL -// clusters. +// Package shard provides application-level routing across PostgreSQL clusters. // -// A Topology owns an ordered set of logical shards. Every Shard is backed by a -// cluster.Cluster, and typed resolvers map application keys directly to shards. -// Built-in resolvers support rendezvous hashing, ordered ranges, and time -// ranges; applications may also provide custom routing logic. Resolvers borrow -// their topology and do not own its clusters. +// A Topology owns an immutable set of logical shards backed by cluster.Cluster +// values. Resolvers map application keys to shards using rendezvous hashing, +// ordered ranges, time ranges, or custom routing logic. // // The package also provides shard grouping, colocation checks, and bounded -// fan-out. -// -// The package does not inspect SQL, hide shard keys in contexts, move data, -// replicate reference tables, or provide distributed transactions. +// parallel operations across shards. package shard diff --git a/shard/errors.go b/shard/errors.go index fcf0411..908af96 100644 --- a/shard/errors.go +++ b/shard/errors.go @@ -6,19 +6,20 @@ import ( ) var ( - // ErrNoShard indicates that a resolver could not map a key to any shard. + // ErrNoShard is returned when a resolver cannot map a key to any shard. ErrNoShard = errors.New("xpg/shard: no shard resolved") - // ErrUnknownShard indicates that routing configuration or custom routing - // logic referenced a shard that does not exist in the topology. + // ErrUnknownShard is returned when routing references a shard that does not + // exist in the topology. ErrUnknownShard = errors.New("xpg/shard: unknown shard") - // ErrShardMismatch indicates that keys expected to be colocated resolved to + // ErrShardMismatch is returned when keys expected to be colocated resolve to // different shards. ErrShardMismatch = errors.New("xpg/shard: keys resolve to different shards") ) -// UnknownShardError identifies a shard that does not exist in a topology. +// UnknownShardError identifies a shard referenced by routing that does not +// exist in the topology. type UnknownShardError struct { ShardID ID } @@ -31,8 +32,8 @@ func (e *UnknownShardError) Unwrap() error { return ErrUnknownShard } -// MismatchError describes the first key that resolved to a different shard -// than the first key. +// MismatchError describes the first key whose resolved shard differs from the +// shard of the first key. type MismatchError struct { Expected ID Actual ID diff --git a/shard/foreach.go b/shard/foreach.go index fc6d5c3..9effadc 100644 --- a/shard/foreach.go +++ b/shard/foreach.go @@ -7,7 +7,7 @@ import ( "sync" ) -// ForEachShardResult contains the result of one shard callback invocation. +// ForEachShardResult contains the result associated with one shard. type ForEachShardResult struct { ShardID ID Err error @@ -28,7 +28,7 @@ func (results ForEachShardResults) Err() error { errs = append( errs, fmt.Errorf( - "xpg/shard: shard %q callback: %w", + "xpg/shard: shard %q: %w", result.ShardID, result.Err, ), @@ -38,21 +38,31 @@ func (results ForEachShardResults) Err() error { return errors.Join(errs...) } -// ForEachShard invokes fn for each shard with at most concurrency callbacks -// running at once. Results are returned in topology registration order. +// ForEachShard invokes fn across the topology with at most concurrency +// callbacks running at once. Results are returned in topology registration +// order; callback execution order is not guaranteed. +// +// Callback failures and context cancellation are stored in the corresponding +// results and can be joined with ForEachShardResults.Err. The returned error is +// reserved for invalid invocation arguments. // // Once context cancellation is observed, callbacks that have not started are // skipped and their results contain ctx.Err(). Callbacks already running are -// responsible for observing ctx. +// responsible for observing ctx. ForEachShard waits for all started callbacks +// to finish before returning. func (t *Topology) ForEachShard( ctx context.Context, concurrency int, fn func(context.Context, Shard) error, ) (ForEachShardResults, error) { - if t == nil || len(t.shards) == 0 { + if t == nil { return nil, errors.New("xpg/shard: topology is nil") } + if len(t.shards) == 0 { + return nil, errors.New("xpg/shard: topology is empty") + } + if concurrency <= 0 { return nil, errors.New("xpg/shard: concurrency must be positive") } diff --git a/shard/foreach_test.go b/shard/foreach_test.go index 2c304e7..ceb8060 100644 --- a/shard/foreach_test.go +++ b/shard/foreach_test.go @@ -3,7 +3,6 @@ package shard import ( "context" "errors" - "strings" "sync/atomic" "testing" "time" @@ -26,28 +25,35 @@ func TestForEachShardValidatesArguments(t *testing.T) { topology: nil, concurrency: 1, fn: func(context.Context, Shard) error { return nil }, - wantError: "topology is nil", + wantError: "xpg/shard: topology is nil", + }, + { + name: "empty topology", + topology: &Topology{}, + concurrency: 1, + fn: func(context.Context, Shard) error { return nil }, + wantError: "xpg/shard: topology is empty", }, { name: "zero concurrency", topology: topology, concurrency: 0, fn: func(context.Context, Shard) error { return nil }, - wantError: "concurrency must be positive", + wantError: "xpg/shard: concurrency must be positive", }, { name: "nil callback", topology: topology, concurrency: 1, fn: nil, - wantError: "callback is nil", + wantError: "xpg/shard: callback is nil", }, } for _, test := range tests { t.Run(test.name, func(t *testing.T) { _, err := test.topology.ForEachShard( - context.Background(), + t.Context(), test.concurrency, test.fn, ) @@ -55,8 +61,8 @@ func TestForEachShardValidatesArguments(t *testing.T) { t.Fatal("expected error") } - if !strings.Contains(err.Error(), test.wantError) { - t.Fatalf("error = %q, want substring %q", err, test.wantError) + if got := err.Error(); got != test.wantError { + t.Fatalf("error = %q, want %q", got, test.wantError) } }) } @@ -68,7 +74,7 @@ func TestForEachShardPreservesRegistrationOrder(t *testing.T) { topology := newTestTopology(t, "shard-c", "shard-a", "shard-b") results, err := topology.ForEachShard( - context.Background(), + t.Context(), 2, func(context.Context, Shard) error { return nil }, ) @@ -101,7 +107,7 @@ func TestForEachShardHonorsConcurrencyLimit(t *testing.T) { "shard-f", ) - ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second) + ctx, cancel := context.WithTimeout(t.Context(), 5*time.Second) defer cancel() started := make(chan struct{}, topology.Len()) @@ -200,7 +206,7 @@ func TestForEachShardCallbackErrorsDoNotStopOtherShards(t *testing.T) { var calls atomic.Int32 results, err := topology.ForEachShard( - context.Background(), + t.Context(), 2, func(_ context.Context, current Shard) error { calls.Add(1) @@ -228,8 +234,9 @@ func TestForEachShardCallbackErrorsDoNotStopOtherShards(t *testing.T) { t.Fatalf("results.Err() = %v, want wrapped sentinel", joined) } - if !strings.Contains(joined.Error(), `shard "shard-b" callback`) { - t.Fatalf("results.Err() = %q, want shard context", joined) + if got, want := joined.Error(), + `xpg/shard: shard "shard-b": callback failed`; got != want { + t.Fatalf("results.Err() = %q, want %q", got, want) } } @@ -237,7 +244,7 @@ func TestForEachShardCanceledBeforeScheduling(t *testing.T) { t.Parallel() topology := newTestTopology(t, "shard-a", "shard-b", "shard-c") - ctx, cancel := context.WithCancel(context.Background()) + ctx, cancel := context.WithCancel(t.Context()) cancel() var calls atomic.Int32 @@ -269,7 +276,7 @@ func TestForEachShardCancellationSkipsCallbacksNotStarted(t *testing.T) { t.Parallel() topology := newTestTopology(t, "shard-a", "shard-b", "shard-c", "shard-d") - ctx, cancel := context.WithCancel(context.Background()) + ctx, cancel := context.WithCancel(t.Context()) defer cancel() started := make(chan struct{}) @@ -336,9 +343,11 @@ func TestForEachShardResultsErr(t *testing.T) { t.Fatalf("Err() = %v, want both failures", err) } - if got := err.Error(); !strings.Contains(got, `shard "shard-a" callback`) || - !strings.Contains(got, `shard "shard-c" callback`) { - t.Fatalf("Err() = %q, want shard context", got) + want := "xpg/shard: shard \"shard-a\": first\n" + + "xpg/shard: shard \"shard-c\": second" + + if got := err.Error(); got != want { + t.Fatalf("Err() = %q, want %q", got, want) } if err := (ForEachShardResults{{ShardID: "shard-a"}}).Err(); err != nil { diff --git a/shard/group.go b/shard/group.go index 59116ec..5d04c0d 100644 --- a/shard/group.go +++ b/shard/group.go @@ -5,8 +5,10 @@ import ( "fmt" ) -// SameShard resolves the keys and verifies that they all belong to the same -// shard. It returns the resolved shard when all keys are colocated. +// SameShard resolves keys and verifies that they all belong to the same shard. +// +// It returns ErrNoShard when no keys are provided and MismatchError when a key +// resolves to a different shard. func SameShard[K any](resolver Resolver[K], keys ...K) (Shard, error) { if resolver == nil { return Shard{}, errors.New("xpg/shard: resolver is nil") @@ -44,15 +46,17 @@ func SameShard[K any](resolver Resolver[K], keys ...K) (Shard, error) { return expected, nil } -// Group contains input keys that resolve to one shard. Keys preserve their -// original relative order. +// Group contains keys that resolve to the same shard. +// Keys preserve their original relative order. type Group[K any] struct { Shard Shard Keys []K } -// GroupByShard resolves every key once and returns groups in order of each -// shard's first appearance in the input. +// GroupByShard resolves each key once and groups keys by shard. +// +// Groups are returned in order of each shard's first appearance in keys. +// Keys within each group preserve their original relative order. func GroupByShard[K any](resolver Resolver[K], keys []K) ([]Group[K], error) { if resolver == nil { return nil, errors.New("xpg/shard: resolver is nil") diff --git a/shard/group_test.go b/shard/group_test.go index 9a65457..5d548f5 100644 --- a/shard/group_test.go +++ b/shard/group_test.go @@ -128,7 +128,7 @@ func TestSameShardReturnsMismatchDetails(t *testing.T) { } } -func TestGroupByShardPreservesStableOrder(t *testing.T) { +func TestGroupByShardPreservesGroupAndKeyOrder(t *testing.T) { t.Parallel() topology := newTestTopology(t, "shard-a", "shard-b") diff --git a/shard/helpers_test.go b/shard/helpers_test.go index 553a0f5..17aab99 100644 --- a/shard/helpers_test.go +++ b/shard/helpers_test.go @@ -1,7 +1,6 @@ package shard import ( - "context" "testing" "github.com/jackc/pgx/v5/pgxpool" @@ -23,7 +22,7 @@ func newTestCluster(t *testing.T, id ID, labels map[string]string) *cluster.Clus config.MaxConns = 1 pool, err := xpg.New( - context.Background(), + t.Context(), config, xpg.WithName("shard."+string(id)+".primary"), ) diff --git a/shard/resolver.go b/shard/resolver.go index e7e6927..500fcb8 100644 --- a/shard/resolver.go +++ b/shard/resolver.go @@ -3,7 +3,7 @@ package shard // Resolver maps a typed application key to one shard. // // Resolve should return ErrNoShard when the key cannot be mapped to a shard. -// Implementations shared by concurrent callers must be concurrency-safe. +// Implementations must be safe for concurrent use. type Resolver[K any] interface { Resolve(key K) (Shard, error) } diff --git a/shard/resolver/custom.go b/shard/resolver/custom.go index b5e04f4..cae2096 100644 --- a/shard/resolver/custom.go +++ b/shard/resolver/custom.go @@ -8,18 +8,20 @@ import ( // ResolveFunc maps an application key to a shard ID within topology. // -// Resolve functions should return shard.ErrNoShard when a key cannot be mapped -// to a shard. Implementations shared by concurrent callers must be deterministic -// and concurrency-safe. They should not perform hidden I/O or modify topology. +// Resolve functions must be deterministic and safe for concurrent use. They +// should return shard.ErrNoShard when a key cannot be mapped and should not +// perform hidden I/O or modify topology. type ResolveFunc[K any] func(key K, topology *shard.Topology) (shard.ID, error) // CustomResolver adapts ResolveFunc to shard.Resolver. +// +// CustomResolver borrows its topology and must not outlive it. type CustomResolver[K any] struct { topology *shard.Topology resolve ResolveFunc[K] } -// NewCustom binds custom routing logic to one immutable topology. +// NewCustom binds custom routing logic to an immutable topology. func NewCustom[K any](topology *shard.Topology, resolve ResolveFunc[K]) (*CustomResolver[K], error) { if err := requireTopology(topology); err != nil { return nil, err @@ -35,7 +37,7 @@ func NewCustom[K any](topology *shard.Topology, resolve ResolveFunc[K]) (*Custom }, nil } -// Resolve maps key to a shard and rejects IDs absent from the bound topology. +// Resolve maps key to a shard in the bound topology. func (resolver *CustomResolver[K]) Resolve(key K) (shard.Shard, error) { if resolver == nil || resolver.topology == nil || resolver.resolve == nil { return shard.Shard{}, errors.New("xpg/shard/resolver: custom resolver is not initialized") @@ -48,7 +50,9 @@ func (resolver *CustomResolver[K]) Resolve(key K) (shard.Shard, error) { resolved, ok := resolver.topology.Shard(id) if !ok { - return shard.Shard{}, &shard.UnknownShardError{ShardID: id} + return shard.Shard{}, &shard.UnknownShardError{ + ShardID: id, + } } return resolved, nil diff --git a/shard/resolver/custom_test.go b/shard/resolver/custom_test.go index a02d4b4..764c4f4 100644 --- a/shard/resolver/custom_test.go +++ b/shard/resolver/custom_test.go @@ -12,16 +12,46 @@ func TestNewCustomValidatesArguments(t *testing.T) { topology := newTestTopology(t, "shard-a") - if resolver, err := NewCustom[int](nil, func(int, *shard.Topology) (shard.ID, error) { - return "shard-a", nil - }); err == nil { - _ = resolver - t.Fatal("expected topology error") + validResolve := ResolveFunc[int]( + func(int, *shard.Topology) (shard.ID, error) { + return "shard-a", nil + }, + ) + + tests := []struct { + name string + topology *shard.Topology + resolve ResolveFunc[int] + wantError string + }{ + { + name: "nil topology", + resolve: validResolve, + wantError: "xpg/shard/resolver: topology is nil or empty", + }, + { + name: "nil resolve function", + topology: topology, + wantError: "xpg/shard/resolver: custom resolve function is nil", + }, } - if resolver, err := NewCustom[int](topology, nil); err == nil { - _ = resolver - t.Fatal("expected resolve function error") + for _, test := range tests { + t.Run(test.name, func(t *testing.T) { + t.Parallel() + + _, err := NewCustom( + test.topology, + test.resolve, + ) + if err == nil { + t.Fatal("expected error") + } + + if got := err.Error(); got != test.wantError { + t.Fatalf("error = %q, want %q", got, test.wantError) + } + }) } } diff --git a/shard/resolver/doc.go b/shard/resolver/doc.go index 7138e4e..a565f8f 100644 --- a/shard/resolver/doc.go +++ b/shard/resolver/doc.go @@ -1,8 +1,7 @@ // Package resolver provides routing strategies for shard.Topology. // -// Resolvers are bound to an immutable topology and map application keys to -// shard.Shard values. The package provides rendezvous hashing, ordered ranges, -// time ranges, and an adapter for custom routing functions. +// Resolvers map application keys to shards using rendezvous hashing, ordered +// ranges, time ranges, or custom routing logic. // // Resolvers borrow their topology and must not outlive it. package resolver diff --git a/shard/resolver/encoder.go b/shard/resolver/encoder.go index 1009493..995357e 100644 --- a/shard/resolver/encoder.go +++ b/shard/resolver/encoder.go @@ -8,10 +8,8 @@ import ( // KeyEncoder converts a typed key into stable canonical bytes. // -// Implementations used for persistent shard placement must remain deterministic -// across processes and releases. Implementations shared by concurrent Resolve -// calls must be concurrency-safe. Changing an encoder changes hash placement -// and may require data migration. +// Implementations must be deterministic and safe for concurrent use. Changing +// an encoder changes persistent hash placement and may require data migration. type KeyEncoder[K any] interface { Encode(K) ([]byte, error) } @@ -19,7 +17,7 @@ type KeyEncoder[K any] interface { // KeyEncoderFunc adapts a function to KeyEncoder. type KeyEncoderFunc[K any] func(K) ([]byte, error) -// Encode encodes key using the adapted function. +// Encode calls the wrapped encoder function. func (encoder KeyEncoderFunc[K]) Encode(key K) ([]byte, error) { if encoder == nil { return nil, errors.New("xpg/shard/resolver: key encoder function is nil") @@ -47,6 +45,24 @@ func BytesKeyEncoder() KeyEncoder[[]byte] { ) } +// Bytes16KeyEncoder encodes a 16-byte key exactly. +func Bytes16KeyEncoder() KeyEncoder[[16]byte] { + return KeyEncoderFunc[[16]byte]( + func(key [16]byte) ([]byte, error) { + return bytes.Clone(key[:]), nil + }, + ) +} + +// Bytes32KeyEncoder encodes a 32-byte key exactly. +func Bytes32KeyEncoder() KeyEncoder[[32]byte] { + return KeyEncoderFunc[[32]byte]( + func(key [32]byte) ([]byte, error) { + return bytes.Clone(key[:]), nil + }, + ) +} + // Int64KeyEncoder encodes a signed integer as big-endian two's-complement. func Int64KeyEncoder() KeyEncoder[int64] { return KeyEncoderFunc[int64]( @@ -72,21 +88,3 @@ func Uint64KeyEncoder() KeyEncoder[uint64] { }, ) } - -// Bytes16KeyEncoder encodes a 16-byte key exactly. -func Bytes16KeyEncoder() KeyEncoder[[16]byte] { - return KeyEncoderFunc[[16]byte]( - func(key [16]byte) ([]byte, error) { - return bytes.Clone(key[:]), nil - }, - ) -} - -// Bytes32KeyEncoder encodes a 32-byte key exactly. -func Bytes32KeyEncoder() KeyEncoder[[32]byte] { - return KeyEncoderFunc[[32]byte]( - func(key [32]byte) ([]byte, error) { - return bytes.Clone(key[:]), nil - }, - ) -} diff --git a/shard/resolver/encoder_test.go b/shard/resolver/encoder_test.go index c2285ed..067187b 100644 --- a/shard/resolver/encoder_test.go +++ b/shard/resolver/encoder_test.go @@ -103,6 +103,8 @@ func TestIntegerKeyEncoders(t *testing.T) { for _, test := range tests { t.Run(test.name, func(t *testing.T) { + t.Parallel() + encoded, err := test.got() if err != nil { t.Fatalf("Encode() error = %v", err) diff --git a/shard/resolver/hash.go b/shard/resolver/hash.go index 08cac0a..a631afe 100644 --- a/shard/resolver/hash.go +++ b/shard/resolver/hash.go @@ -20,19 +20,22 @@ const ( rendezvousLengthSize = 4 ) -// HashResolver implements rendezvous/HRW routing with SHA-256 and stable named -// shard IDs. +// HashResolver routes keys using rendezvous/HRW hashing with SHA-256. +// +// HashResolver captures the shard set when it is created. The shards remain +// borrowed from the topology, so the resolver must not outlive it. type HashResolver[K any] struct { - shards []shard.Shard - prefix []byte - encoder KeyEncoder[K] - maxIDLength int + shards []shard.Shard + prefix []byte + encoder KeyEncoder[K] + maxShardIDLength int } -// NewHash creates the version-1 rendezvous resolver bound to topology. +// NewHash creates a rendezvous hash resolver bound to topology. // -// Namespace is part of the persistent placement contract. Changing it changes -// shard placement and may require data migration. +// Namespace is an opaque non-empty string and part of the persistent placement +// contract. Changing the namespace, key encoder, shard IDs, or placement format +// changes shard placement and may require data migration. func NewHash[K any]( topology *shard.Topology, namespace string, @@ -55,7 +58,7 @@ func NewHash[K any]( } shards := topology.Shards() - maxIDLength := 0 + maxShardIDLength := 0 for _, candidate := range shards { id := candidate.ID() @@ -64,16 +67,9 @@ func NewHash[K any]( return nil, errors.New("xpg/shard/resolver: shard ID is too large") } - // Resolve reuses one score-input buffer for all candidates, sized for - // the largest shard ID in the topology. - maxIDLength = max(maxIDLength, len(id)) + maxShardIDLength = max(maxShardIDLength, len(id)) } - // Prefix is part of the persistent placement format: - // - // domain || namespace_length || namespace - // - // Changing this layout changes shard placement and requires data migration. prefixSize := len(rendezvousDomain) + rendezvousLengthSize + len(namespace) prefix := make([]byte, prefixSize) @@ -90,13 +86,14 @@ func NewHash[K any]( copy(prefix[namespaceOffset:], namespace) return &HashResolver[K]{ - shards: shards, - prefix: prefix, - encoder: encoder, - maxIDLength: maxIDLength, + shards: shards, + prefix: prefix, + encoder: encoder, + maxShardIDLength: maxShardIDLength, }, nil } +// Resolve maps key to a shard using rendezvous hashing. func (resolver *HashResolver[K]) Resolve(key K) (shard.Shard, error) { if resolver == nil || len(resolver.shards) == 0 || resolver.encoder == nil { return shard.Shard{}, errors.New("xpg/shard/resolver: hash resolver is not initialized") @@ -116,9 +113,17 @@ func (resolver *HashResolver[K]) Resolve(key K) (shard.Shard, error) { idLengthOffset := keyOffset + len(encoded) idOffset := idLengthOffset + rendezvousLengthSize - // The candidate-independent part is written once. Only the shard ID suffix - // is overwritten while evaluating rendezvous scores. - scoreInput := make([]byte, idOffset+resolver.maxIDLength) + // Persistent placement format: + // + // domain || namespace_length || namespace || + // key_length || key || shard_id_length || shard_id + // + // The candidate-independent prefix and key are written once. Only the shard + // ID suffix is overwritten while evaluating candidates. + scoreInput := make( + []byte, + idOffset+resolver.maxShardIDLength, + ) copy(scoreInput, resolver.prefix) @@ -152,7 +157,9 @@ func (resolver *HashResolver[K]) Resolve(key K) (shard.Shard, error) { // Shard ID is the deterministic tie-breaker, so placement does not // depend on topology registration order when scores are equal. - if !hasBest || comparison > 0 || (comparison == 0 && candidateID < bestID) { + if !hasBest || + comparison > 0 || + (comparison == 0 && candidateID < bestID) { selected = candidate best = score bestID = candidateID diff --git a/shard/resolver/hash_test.go b/shard/resolver/hash_test.go index 6fdf064..f29c5c4 100644 --- a/shard/resolver/hash_test.go +++ b/shard/resolver/hash_test.go @@ -13,19 +13,50 @@ func TestNewHashValidatesArguments(t *testing.T) { topology := newTestTopology(t, "shard-a") - if resolver, err := NewHash[string](nil, "users", StringKeyEncoder()); err == nil { - _ = resolver - t.Fatal("expected topology error") + tests := []struct { + name string + topology *shard.Topology + namespace string + encoder KeyEncoder[string] + wantError string + }{ + { + name: "nil topology", + namespace: "users", + encoder: StringKeyEncoder(), + wantError: "xpg/shard/resolver: topology is nil or empty", + }, + { + name: "nil encoder", + topology: topology, + namespace: "users", + wantError: "xpg/shard/resolver: key encoder is nil", + }, + { + name: "empty namespace", + topology: topology, + encoder: StringKeyEncoder(), + wantError: "xpg/shard/resolver: hash namespace must not be empty", + }, } - if resolver, err := NewHash[string](topology, "users", nil); err == nil { - _ = resolver - t.Fatal("expected encoder error") - } + for _, test := range tests { + t.Run(test.name, func(t *testing.T) { + t.Parallel() - if resolver, err := NewHash(topology, "", StringKeyEncoder()); err == nil { - _ = resolver - t.Fatal("expected namespace error") + _, err := NewHash( + test.topology, + test.namespace, + test.encoder, + ) + if err == nil { + t.Fatal("expected error") + } + + if got := err.Error(); got != test.wantError { + t.Fatalf("error = %q, want %q", got, test.wantError) + } + }) } } @@ -76,13 +107,20 @@ func TestHashResolverStablePlacementVectors(t *testing.T) { for _, test := range tests { t.Run(test.key, func(t *testing.T) { + t.Parallel() + resolved, err := resolver.Resolve(test.key) if err != nil { t.Fatalf("Resolve() error = %v", err) } if got := resolved.ID(); got != test.want { - t.Fatalf("Resolve(%q).ID() = %q, want %q", test.key, got, test.want) + t.Fatalf( + "Resolve(%q).ID() = %q, want %q", + test.key, + got, + test.want, + ) } }) } @@ -98,6 +136,7 @@ func TestHashResolverPlacementDoesNotDependOnTopologyOrder(t *testing.T) { if err != nil { t.Fatalf("NewHash(first) error = %v", err) } + secondResolver, err := NewHash(second, "users", StringKeyEncoder()) if err != nil { t.Fatalf("NewHash(second) error = %v", err) @@ -108,6 +147,7 @@ func TestHashResolverPlacementDoesNotDependOnTopologyOrder(t *testing.T) { if err != nil { t.Fatalf("first Resolve(%q) error = %v", key, err) } + secondShard, err := secondResolver.Resolve(key) if err != nil { t.Fatalf("second Resolve(%q) error = %v", key, err) @@ -134,6 +174,7 @@ func TestHashResolverAddingShardOnlyMovesKeysToNewShard(t *testing.T) { if err != nil { t.Fatalf("NewHash(before) error = %v", err) } + afterResolver, err := NewHash(after, "users", StringKeyEncoder()) if err != nil { t.Fatalf("NewHash(after) error = %v", err) @@ -148,6 +189,7 @@ func TestHashResolverAddingShardOnlyMovesKeysToNewShard(t *testing.T) { if err != nil { t.Fatalf("before Resolve(%q) error = %v", key, err) } + current, err := afterResolver.Resolve(key) if err != nil { t.Fatalf("after Resolve(%q) error = %v", key, err) @@ -158,6 +200,7 @@ func TestHashResolverAddingShardOnlyMovesKeysToNewShard(t *testing.T) { } moved++ + if current.ID() != "shard-c" { t.Fatalf( "Resolve(%q) moved from %q to existing shard %q", diff --git a/shard/resolver/helpers_test.go b/shard/resolver/helpers_test.go index cfeabf3..a12991b 100644 --- a/shard/resolver/helpers_test.go +++ b/shard/resolver/helpers_test.go @@ -1,7 +1,6 @@ package resolver import ( - "context" "testing" "github.com/jackc/pgx/v5/pgxpool" @@ -27,7 +26,7 @@ func newTestTopology(t *testing.T, ids ...shard.ID) *shard.Topology { poolConfig.MaxConns = 1 pool, err := xpg.New( - context.Background(), + t.Context(), poolConfig, xpg.WithName("shard."+string(id)+".primary"), ) diff --git a/shard/resolver/range.go b/shard/resolver/range.go index 67595fe..9714777 100644 --- a/shard/resolver/range.go +++ b/shard/resolver/range.go @@ -20,24 +20,15 @@ type Range[K cmp.Ordered] struct { ShardID shard.ID } -// RangeResolver resolves ordered keys through bounded, non-overlapping ranges. +// RangeResolver routes ordered keys through non-overlapping ranges. type RangeResolver[K cmp.Ordered] struct { ranges []rangeEntry[K] } -type rangeEntry[K cmp.Ordered] struct { - start K - end K - shard shard.Shard - - sourceIndex int -} - -// NewRange creates a resolver from bounded, non-overlapping half-open ranges. +// NewRange creates a resolver from non-overlapping half-open ranges. // -// NewRange copies the supplied ranges into an internal representation, sorts -// them by Start, and validates that they do not overlap. The caller's slice is -// not modified. +// The supplied ranges may be unordered. NewRange copies and sorts them by Start, +// validates their boundaries and overlap, and leaves the caller's slice unchanged. func NewRange[K cmp.Ordered](topology *shard.Topology, ranges []Range[K]) (*RangeResolver[K], error) { if err := requireTopology(topology); err != nil { return nil, err @@ -112,9 +103,6 @@ func NewRange[K cmp.Ordered](topology *shard.Topology, ranges []Range[K]) (*Rang } // Resolve returns the shard whose configured range contains key. -// -// Resolve performs only an in-memory lookup. It does not acquire a connection -// or execute a PostgreSQL query. func (resolver *RangeResolver[K]) Resolve(key K) (shard.Shard, error) { if resolver == nil || len(resolver.ranges) == 0 { return shard.Shard{}, errors.New("xpg/shard/resolver: range resolver is not initialized") @@ -143,3 +131,11 @@ func (resolver *RangeResolver[K]) Resolve(key K) (shard.Shard, error) { return entry.shard, nil } + +type rangeEntry[K cmp.Ordered] struct { + start K + end K + shard shard.Shard + + sourceIndex int +} diff --git a/shard/resolver/range_test.go b/shard/resolver/range_test.go index c4b4ae2..d001386 100644 --- a/shard/resolver/range_test.go +++ b/shard/resolver/range_test.go @@ -4,7 +4,6 @@ import ( "errors" "math" "slices" - "strings" "testing" "github.com/mkbeh/xpg/shard" @@ -15,53 +14,72 @@ func TestNewRangeValidatesArguments(t *testing.T) { topology := newTestTopology(t, "shard-a") - if resolver, err := NewRange[int](nil, []Range[int]{{Start: 0, End: 10, ShardID: "shard-a"}}); err == nil { - _ = resolver - t.Fatal("expected topology error") - } - - if resolver, err := NewRange[int](topology, nil); err == nil { - _ = resolver - t.Fatal("expected empty ranges error") - } - tests := []struct { - name string - valueRange Range[int] - want string + name string + topology *shard.Topology + ranges []Range[int] + wantError string }{ { - name: "empty shard ID", - valueRange: Range[int]{Start: 0, End: 10}, - want: "shard ID must not be empty", + name: "nil topology", + ranges: []Range[int]{ + {Start: 0, End: 10, ShardID: "shard-a"}, + }, + wantError: "xpg/shard/resolver: topology is nil or empty", + }, + { + name: "empty ranges", + topology: topology, + wantError: "xpg/shard/resolver: range resolver requires at least one range", + }, + { + name: "empty shard ID", + topology: topology, + ranges: []Range[int]{ + {Start: 0, End: 10}, + }, + wantError: "xpg/shard/resolver: range 0: shard ID must not be empty", }, { - name: "empty interval", - valueRange: Range[int]{Start: 10, End: 10, ShardID: "shard-a"}, - want: "must satisfy start < end", + name: "empty interval", + topology: topology, + ranges: []Range[int]{ + {Start: 10, End: 10, ShardID: "shard-a"}, + }, + wantError: "xpg/shard/resolver: range 0 must satisfy start < end", }, { - name: "reversed interval", - valueRange: Range[int]{Start: 20, End: 10, ShardID: "shard-a"}, - want: "must satisfy start < end", + name: "reversed interval", + topology: topology, + ranges: []Range[int]{ + {Start: 20, End: 10, ShardID: "shard-a"}, + }, + wantError: "xpg/shard/resolver: range 0 must satisfy start < end", }, { - name: "unknown shard", - valueRange: Range[int]{Start: 0, End: 10, ShardID: "missing"}, - want: `unknown shard "missing"`, + name: "unknown shard", + topology: topology, + ranges: []Range[int]{ + {Start: 0, End: 10, ShardID: "missing"}, + }, + wantError: `xpg/shard/resolver: range 0: xpg/shard: unknown shard "missing"`, }, } for _, test := range tests { t.Run(test.name, func(t *testing.T) { - resolver, err := NewRange(topology, []Range[int]{test.valueRange}) + t.Parallel() + + _, err := NewRange( + test.topology, + test.ranges, + ) if err == nil { - _ = resolver t.Fatal("expected error") } - if !strings.Contains(err.Error(), test.want) { - t.Fatalf("error = %q, want substring %q", err, test.want) + if got := err.Error(); got != test.wantError { + t.Fatalf("error = %q, want %q", got, test.wantError) } }) } @@ -72,17 +90,45 @@ func TestNewRangeRejectsNaNBoundaries(t *testing.T) { topology := newTestTopology(t, "shard-a") - tests := []Range[float64]{ - {Start: math.NaN(), End: 10, ShardID: "shard-a"}, - {Start: 0, End: math.NaN(), ShardID: "shard-a"}, + tests := []struct { + name string + valueRange Range[float64] + }{ + { + name: "NaN start", + valueRange: Range[float64]{ + Start: math.NaN(), + End: 10, + ShardID: "shard-a", + }, + }, + { + name: "NaN end", + valueRange: Range[float64]{ + Start: 0, + End: math.NaN(), + ShardID: "shard-a", + }, + }, } - for index, valueRange := range tests { - resolver, err := NewRange(topology, []Range[float64]{valueRange}) - if err == nil { - _ = resolver - t.Fatalf("range %d: expected error", index) - } + for _, test := range tests { + t.Run(test.name, func(t *testing.T) { + t.Parallel() + + _, err := NewRange( + topology, + []Range[float64]{test.valueRange}, + ) + if err == nil { + t.Fatal("expected error") + } + + if got, want := err.Error(), + "xpg/shard/resolver: range 0 must satisfy start < end"; got != want { + t.Fatalf("error = %q, want %q", got, want) + } + }) } } @@ -91,12 +137,11 @@ func TestNewRangeRejectsOverlapUsingSourceIndexes(t *testing.T) { topology := newTestTopology(t, "shard-a", "shard-b") - resolver, err := NewRange(topology, []Range[int]{ + _, err := NewRange(topology, []Range[int]{ {Start: 100, End: 200, ShardID: "shard-b"}, {Start: 50, End: 150, ShardID: "shard-a"}, }) if err == nil { - _ = resolver t.Fatal("expected overlap error") } @@ -113,7 +158,7 @@ func TestNewRangeDoesNotModifyInputOrder(t *testing.T) { {Start: 100, End: 200, ShardID: "shard-b"}, {Start: 0, End: 100, ShardID: "shard-a"}, } - want := append([]Range[int](nil), ranges...) + want := slices.Clone(ranges) if _, err := NewRange(topology, ranges); err != nil { t.Fatalf("NewRange() error = %v", err) @@ -156,12 +201,14 @@ func TestRangeResolverHalfOpenBoundariesAndGaps(t *testing.T) { if !errors.Is(err, test.wantErr) { t.Fatalf("Resolve(%d) error = %v, want %v", test.key, err, test.wantErr) } + continue } if err != nil { t.Fatalf("Resolve(%d) error = %v", test.key, err) } + if got := resolved.ID(); got != test.wantID { t.Fatalf("Resolve(%d).ID() = %q, want %q", test.key, got, test.wantID) } @@ -184,6 +231,7 @@ func TestRangeResolverAllowsAdjacentRanges(t *testing.T) { if err != nil { t.Fatalf("Resolve() error = %v", err) } + if got, want := resolved.ID(), shard.ID("shard-b"); got != want { t.Fatalf("Resolve().ID() = %q, want %q", got, want) } @@ -205,6 +253,7 @@ func TestRangeResolverSupportsStrings(t *testing.T) { if err != nil { t.Fatalf("Resolve() error = %v", err) } + if got, want := resolved.ID(), shard.ID("shard-b"); got != want { t.Fatalf("Resolve().ID() = %q, want %q", got, want) } diff --git a/shard/resolver/time_range.go b/shard/resolver/time_range.go index 4e28a80..84d329b 100644 --- a/shard/resolver/time_range.go +++ b/shard/resolver/time_range.go @@ -20,25 +20,16 @@ type TimeRange struct { ShardID shard.ID } -// TimeRangeResolver resolves time instants through bounded, non-overlapping -// ranges. +// TimeRangeResolver routes time instants through non-overlapping ranges. type TimeRangeResolver struct { ranges []timeRangeEntry } -type timeRangeEntry struct { - start time.Time - end time.Time - shard shard.Shard - - sourceIndex int -} - -// NewTimeRange creates a resolver from bounded, non-overlapping half-open time -// ranges. +// NewTimeRange creates a resolver from non-overlapping half-open time ranges. // -// Range boundaries are normalized to UTC. The caller's slice and time values -// are not modified. +// The supplied ranges may be unordered. NewTimeRange normalizes boundaries to +// UTC, sorts ranges by Start, validates their overlap, and leaves the caller's +// slice unchanged. func NewTimeRange(topology *shard.Topology, ranges []TimeRange) (*TimeRangeResolver, error) { if err := requireTopology(topology); err != nil { return nil, err @@ -112,9 +103,6 @@ func NewTimeRange(topology *shard.Topology, ranges []TimeRange) (*TimeRangeResol } // Resolve returns the shard whose configured time range contains key. -// -// Resolve performs only an in-memory lookup. It does not acquire a connection -// or execute a PostgreSQL query. func (resolver *TimeRangeResolver) Resolve(key time.Time) (shard.Shard, error) { if resolver == nil || len(resolver.ranges) == 0 { return shard.Shard{}, errors.New("xpg/shard/resolver: time range resolver is not initialized") @@ -146,6 +134,14 @@ func (resolver *TimeRangeResolver) Resolve(key time.Time) (shard.Shard, error) { return entry.shard, nil } -func timeToUTC(t time.Time) time.Time { - return t.UTC() +type timeRangeEntry struct { + start time.Time + end time.Time + shard shard.Shard + + sourceIndex int +} + +func timeToUTC(value time.Time) time.Time { + return value.UTC() } diff --git a/shard/resolver/time_range_test.go b/shard/resolver/time_range_test.go index 6611611..246c34d 100644 --- a/shard/resolver/time_range_test.go +++ b/shard/resolver/time_range_test.go @@ -3,7 +3,6 @@ package resolver import ( "errors" "slices" - "strings" "testing" "time" @@ -17,53 +16,72 @@ func TestNewTimeRangeValidatesArguments(t *testing.T) { start := time.Date(2026, 1, 1, 0, 0, 0, 0, time.UTC) end := start.Add(time.Hour) - if resolver, err := NewTimeRange(nil, []TimeRange{{Start: start, End: end, ShardID: "shard-a"}}); err == nil { - _ = resolver - t.Fatal("expected topology error") - } - - if resolver, err := NewTimeRange(topology, nil); err == nil { - _ = resolver - t.Fatal("expected empty ranges error") - } - tests := []struct { - name string - valueRange TimeRange - want string + name string + topology *shard.Topology + ranges []TimeRange + wantError string }{ { - name: "empty shard ID", - valueRange: TimeRange{Start: start, End: end}, - want: "shard ID must not be empty", + name: "nil topology", + ranges: []TimeRange{ + {Start: start, End: end, ShardID: "shard-a"}, + }, + wantError: "xpg/shard/resolver: topology is nil or empty", + }, + { + name: "empty ranges", + topology: topology, + wantError: "xpg/shard/resolver: time range resolver requires at least one range", + }, + { + name: "empty shard ID", + topology: topology, + ranges: []TimeRange{ + {Start: start, End: end}, + }, + wantError: "xpg/shard/resolver: time range 0: shard ID must not be empty", }, { - name: "empty interval", - valueRange: TimeRange{Start: start, End: start, ShardID: "shard-a"}, - want: "must satisfy start < end", + name: "empty interval", + topology: topology, + ranges: []TimeRange{ + {Start: start, End: start, ShardID: "shard-a"}, + }, + wantError: "xpg/shard/resolver: time range 0 must satisfy start < end", }, { - name: "reversed interval", - valueRange: TimeRange{Start: end, End: start, ShardID: "shard-a"}, - want: "must satisfy start < end", + name: "reversed interval", + topology: topology, + ranges: []TimeRange{ + {Start: end, End: start, ShardID: "shard-a"}, + }, + wantError: "xpg/shard/resolver: time range 0 must satisfy start < end", }, { - name: "unknown shard", - valueRange: TimeRange{Start: start, End: end, ShardID: "missing"}, - want: `unknown shard "missing"`, + name: "unknown shard", + topology: topology, + ranges: []TimeRange{ + {Start: start, End: end, ShardID: "missing"}, + }, + wantError: `xpg/shard/resolver: time range 0: xpg/shard: unknown shard "missing"`, }, } for _, test := range tests { t.Run(test.name, func(t *testing.T) { - resolver, err := NewTimeRange(topology, []TimeRange{test.valueRange}) + t.Parallel() + + _, err := NewTimeRange( + test.topology, + test.ranges, + ) if err == nil { - _ = resolver t.Fatal("expected error") } - if !strings.Contains(err.Error(), test.want) { - t.Fatalf("error = %q, want substring %q", err, test.want) + if got := err.Error(); got != test.wantError { + t.Fatalf("error = %q, want %q", got, test.wantError) } }) } @@ -75,12 +93,11 @@ func TestNewTimeRangeRejectsOverlapUsingSourceIndexes(t *testing.T) { topology := newTestTopology(t, "shard-a", "shard-b") base := time.Date(2026, 1, 1, 0, 0, 0, 0, time.UTC) - resolver, err := NewTimeRange(topology, []TimeRange{ + _, err := NewTimeRange(topology, []TimeRange{ {Start: base.Add(2 * time.Hour), End: base.Add(4 * time.Hour), ShardID: "shard-b"}, {Start: base.Add(time.Hour), End: base.Add(3 * time.Hour), ShardID: "shard-a"}, }) if err == nil { - _ = resolver t.Fatal("expected overlap error") } @@ -99,7 +116,7 @@ func TestNewTimeRangeDoesNotModifyInput(t *testing.T) { {Start: base.Add(time.Hour), End: base.Add(2 * time.Hour), ShardID: "shard-b"}, {Start: base, End: base.Add(time.Hour), ShardID: "shard-a"}, } - want := append([]TimeRange(nil), ranges...) + want := slices.Clone(ranges) if _, err := NewTimeRange(topology, ranges); err != nil { t.Fatalf("NewTimeRange() error = %v", err) @@ -141,12 +158,14 @@ func TestTimeRangeResolverHalfOpenBoundariesAndGaps(t *testing.T) { if !errors.Is(err, test.wantErr) { t.Fatalf("Resolve(%v) error = %v, want %v", test.key, err, test.wantErr) } + continue } if err != nil { t.Fatalf("Resolve(%v) error = %v", test.key, err) } + if got := resolved.ID(); got != test.wantID { t.Fatalf("Resolve(%v).ID() = %q, want %q", test.key, got, test.wantID) } @@ -173,6 +192,7 @@ func TestTimeRangeResolverNormalizesToUTC(t *testing.T) { if err != nil { t.Fatalf("Resolve() error = %v", err) } + if got, want := resolved.ID(), shard.ID("shard-a"); got != want { t.Fatalf("Resolve().ID() = %q, want %q", got, want) } @@ -195,6 +215,7 @@ func TestTimeRangeResolverAllowsAdjacentRanges(t *testing.T) { if err != nil { t.Fatalf("Resolve() error = %v", err) } + if got, want := resolved.ID(), shard.ID("shard-b"); got != want { t.Fatalf("Resolve().ID() = %q, want %q", got, want) } diff --git a/shard/resolver/validation.go b/shard/resolver/validation.go index dce5f31..c70966f 100644 --- a/shard/resolver/validation.go +++ b/shard/resolver/validation.go @@ -8,9 +8,7 @@ import ( func requireTopology(topology *shard.Topology) error { if topology == nil || topology.Len() == 0 { - return errors.New( - "xpg/shard/resolver: topology is nil or empty", - ) + return errors.New("xpg/shard/resolver: topology is nil or empty") } return nil @@ -18,9 +16,7 @@ func requireTopology(topology *shard.Topology) error { func requireShardID(id shard.ID) error { if id == "" { - return errors.New( - "shard ID must not be empty", - ) + return errors.New("shard ID must not be empty") } return nil diff --git a/shard/shard.go b/shard/shard.go index 3394e5b..c3afb65 100644 --- a/shard/shard.go +++ b/shard/shard.go @@ -19,7 +19,7 @@ type Shard struct { cluster *cluster.Cluster } -// ID returns the stable logical shard ID. +// ID returns the logical shard ID. func (s Shard) ID() ID { if s.cluster == nil { return "" @@ -46,9 +46,10 @@ func (s Shard) Labels() map[string]string { return s.cluster.Labels() } -// Primary returns the shard primary pool, or nil when the shard cluster has no -// primary configured. The returned pool is borrowed and remains owned by the -// shard cluster. +// Primary returns the shard primary pool. +// +// Primary returns nil when no primary is configured. The returned pool is +// borrowed and must not be closed separately. func (s Shard) Primary() *xpg.Pool { if s.cluster == nil { return nil @@ -57,7 +58,7 @@ func (s Shard) Primary() *xpg.Pool { return s.cluster.Primary() } -// ReadPool returns a borrowed pool for a read operation according to policy. +// ReadPool returns a borrowed pool according to policy. func (s Shard) ReadPool( ctx context.Context, policy cluster.ReadPolicy, @@ -82,7 +83,8 @@ func (s Shard) InPrimaryTx( return s.cluster.InPrimaryTx(ctx, options, fn) } -// InReadTx executes fn in a read-only transaction resolved within this shard. +// InReadTx executes fn in a read-only transaction on a pool selected according +// to policy within the shard. func (s Shard) InReadTx( ctx context.Context, policy cluster.ReadPolicy, diff --git a/shard/shard_test.go b/shard/shard_test.go index 5355873..400822b 100644 --- a/shard/shard_test.go +++ b/shard/shard_test.go @@ -1,7 +1,6 @@ package shard import ( - "context" "errors" "testing" @@ -30,16 +29,16 @@ func TestShardZeroValue(t *testing.T) { t.Fatalf("Primary() = %p, want nil", primary) } - if _, err := shard.ReadPool(context.Background(), cluster.ReadPrimary); !errors.Is(err, ErrNoShard) { + if _, err := shard.ReadPool(t.Context(), cluster.ReadPrimary); !errors.Is(err, ErrNoShard) { t.Fatalf("ReadPool() error = %v, want ErrNoShard", err) } - if err := shard.InPrimaryTx(context.Background(), pgx.TxOptions{}, nil); !errors.Is(err, ErrNoShard) { + if err := shard.InPrimaryTx(t.Context(), pgx.TxOptions{}, nil); !errors.Is(err, ErrNoShard) { t.Fatalf("InPrimaryTx() error = %v, want ErrNoShard", err) } if err := shard.InReadTx( - context.Background(), + t.Context(), cluster.ReadPrimary, cluster.ReadTxOptions{}, nil, @@ -87,7 +86,7 @@ func TestShardDelegatesClusterMetadataAndRouting(t *testing.T) { t.Fatal("Primary() did not return cluster primary") } - pool, err := resolved.ReadPool(context.Background(), cluster.ReadPrimary) + pool, err := resolved.ReadPool(t.Context(), cluster.ReadPrimary) if err != nil { t.Fatalf("ReadPool() error = %v", err) } diff --git a/shard/topology.go b/shard/topology.go index 4c1647f..db7468a 100644 --- a/shard/topology.go +++ b/shard/topology.go @@ -9,10 +9,9 @@ import ( "github.com/mkbeh/xpg/cluster" ) -// Config registers one Cluster as a logical shard in a Topology. +// Config registers one Cluster as a logical shard. // -// The shard ID and labels are provided by Cluster. Config is intentionally -// retained as the extension point for future topology-specific options. +// The shard ID and labels are provided by Cluster. type Config struct { Cluster *cluster.Cluster } @@ -29,11 +28,15 @@ type Topology struct { closeOnce sync.Once } -// NewTopology validates and creates an immutable topology. Shards retain their -// registration order. Every cluster must have a non-empty and unique ID. +// NewTopology validates and creates an immutable topology. +// +// Shards retain their registration order. Every cluster must have a non-empty +// and unique ID. func NewTopology(configs []Config) (*Topology, error) { if len(configs) == 0 { - return nil, errors.New("xpg/shard: topology must contain at least one shard") + return nil, errors.New( + "xpg/shard: topology must contain at least one shard", + ) } shards := make([]Shard, len(configs)) @@ -41,16 +44,27 @@ func NewTopology(configs []Config) (*Topology, error) { for index, config := range configs { if config.Cluster == nil { - return nil, fmt.Errorf("xpg/shard: shard %d: cluster is nil", index) + return nil, fmt.Errorf( + "xpg/shard: shard %d: cluster is nil", + index, + ) } id := config.Cluster.ID() if id == "" { - return nil, fmt.Errorf("xpg/shard: shard %d: cluster ID must not be empty", index) + return nil, fmt.Errorf( + "xpg/shard: shard %d: cluster ID must not be empty", + index, + ) } - if _, exists := indexByID[id]; exists { - return nil, fmt.Errorf("xpg/shard: duplicate shard ID %q", id) + if previousIndex, exists := indexByID[id]; exists { + return nil, fmt.Errorf( + "xpg/shard: duplicate shard ID %q at indexes %d and %d", + id, + previousIndex, + index, + ) } shards[index] = Shard{ @@ -74,22 +88,14 @@ func (t *Topology) Len() int { return len(t.shards) } -// At returns the shard at index in registration order. At panics when index is -// outside the topology, matching ordinary slice indexing semantics. +// At returns the shard at index in registration order. +// +// At panics when t is nil or index is out of range. func (t *Topology) At(index int) Shard { return t.shards[index] } -// Shards returns a defensive copy of shards in registration order. -func (t *Topology) Shards() []Shard { - if t == nil { - return nil - } - - return slices.Clone(t.shards) -} - -// Shard returns one shard by stable ID. +// Shard returns the shard with id. func (t *Topology) Shard(id ID) (Shard, bool) { if t == nil { return Shard{}, false @@ -103,6 +109,15 @@ func (t *Topology) Shard(id ID) (Shard, bool) { return t.shards[index], true } +// Shards returns a defensive copy of shards in registration order. +func (t *Topology) Shards() []Shard { + if t == nil { + return nil + } + + return slices.Clone(t.shards) +} + // Close closes owned clusters in reverse registration order. Close is safe to // call multiple times. func (t *Topology) Close() { @@ -111,8 +126,8 @@ func (t *Topology) Close() { } t.closeOnce.Do(func() { - for _, v := range slices.Backward(t.shards) { - v.cluster.Close() + for _, shard := range slices.Backward(t.shards) { + shard.cluster.Close() } }) } diff --git a/shard/topology_test.go b/shard/topology_test.go index 4d50262..81712e5 100644 --- a/shard/topology_test.go +++ b/shard/topology_test.go @@ -1,7 +1,6 @@ package shard import ( - "strings" "testing" ) @@ -28,8 +27,8 @@ func TestNewTopologyRejectsNilCluster(t *testing.T) { t.Fatal("expected error") } - if !strings.Contains(err.Error(), "cluster is nil") { - t.Fatalf("error = %q, want cluster validation error", err) + if got, want := err.Error(), "xpg/shard: shard 0: cluster is nil"; got != want { + t.Fatalf("error = %q, want %q", got, want) } } @@ -43,8 +42,8 @@ func TestNewTopologyRejectsEmptyClusterID(t *testing.T) { t.Fatal("expected error") } - if !strings.Contains(err.Error(), "cluster ID must not be empty") { - t.Fatalf("error = %q, want cluster ID validation error", err) + if got, want := err.Error(), "xpg/shard: shard 0: cluster ID must not be empty"; got != want { + t.Fatalf("error = %q, want %q", got, want) } } @@ -62,8 +61,9 @@ func TestNewTopologyRejectsDuplicateIDs(t *testing.T) { t.Fatal("expected error") } - if !strings.Contains(err.Error(), `duplicate shard ID "shard-a"`) { - t.Fatalf("error = %q, want duplicate shard ID error", err) + if got, want := err.Error(), + `xpg/shard: duplicate shard ID "shard-a" at indexes 0 and 1`; got != want { + t.Fatalf("error = %q, want %q", got, want) } } diff --git a/stats.go b/stats.go index 353ff2b..8d1aafb 100644 --- a/stats.go +++ b/stats.go @@ -3,6 +3,8 @@ package xpg import "time" // PoolStats is a detached point-in-time snapshot of connection pool statistics. +// +// Counter fields are cumulative for the lifetime of the pool. type PoolStats struct { // Current state. @@ -45,15 +47,15 @@ type PoolStats struct { // Connection lifecycle. - // NewConnsCount is the cumulative number of connections opened by the pool. + // NewConnsCount is the cumulative number of connections created by the pool. NewConnsCount int64 - // MaxIdleDestroyCount is the cumulative number of connections closed after - // exceeding MaxConnIdleTime. + // MaxIdleDestroyCount is the cumulative number of connections closed because + // they exceeded MaxConnIdleTime. MaxIdleDestroyCount int64 // MaxLifetimeDestroyCount is the cumulative number of connections closed - // after exceeding MaxConnLifetime. + // because they exceeded MaxConnLifetime. MaxLifetimeDestroyCount int64 } diff --git a/tx.go b/tx.go index 178629d..b33ebfc 100644 --- a/tx.go +++ b/tx.go @@ -10,14 +10,14 @@ import ( // InTx executes fn in a transaction configured by txOptions. // -// The transaction is committed when fn returns nil and rolled back when fn -// returns an error. If fn panics, rollback is attempted before the panic is -// propagated. The callback must not call Commit or Rollback; InTx owns -// transaction finalization. +// If fn returns nil, the transaction is committed; otherwise it is rolled back. +// If fn panics, rollback is attempted before the panic is propagated. The +// callback must not call Commit or Rollback; InTx owns transaction +// finalization. // -// The callback receives the same context and an explicit pgx.Tx. Context -// cancellation does not automatically finalize the transaction while fn is -// running; fn should observe ctx and return promptly. +// The callback receives ctx unchanged. Context cancellation does not +// automatically finalize the transaction while fn is running; fn should +// observe ctx and return promptly. func (p *Pool) InTx( ctx context.Context, txOptions pgx.TxOptions, @@ -42,13 +42,14 @@ func (p *Pool) InTx( return nil } -// InSavepoint executes fn in a pseudo-nested transaction implemented with a -// PostgreSQL savepoint. +// InSavepoint executes fn within a PostgreSQL savepoint. // -// The savepoint is released when fn returns nil and rolled back when fn returns -// an error. If fn panics, rollback is attempted before the panic is propagated. -// The callback must not call Commit or Rollback; InSavepoint owns savepoint +// If fn returns nil, the savepoint is released; otherwise it is rolled back. +// If fn panics, rollback is attempted before the panic is propagated. The +// callback must not call Commit or Rollback; InSavepoint owns savepoint // finalization. +// +// The callback receives ctx unchanged and should observe its cancellation. func InSavepoint( ctx context.Context, tx pgx.Tx, diff --git a/tx_test.go b/tx_test.go index 0565735..f44d20e 100644 --- a/tx_test.go +++ b/tx_test.go @@ -2,20 +2,17 @@ package xpg import ( "context" + "errors" "testing" "github.com/jackc/pgx/v5" ) -type testTx struct { - pgx.Tx -} - func TestPoolInTxNilFunc(t *testing.T) { t.Parallel() err := (&Pool{}).InTx( - context.Background(), + t.Context(), pgx.TxOptions{}, nil, ) @@ -33,7 +30,7 @@ func TestInSavepointNilTx(t *testing.T) { called := false err := InSavepoint( - context.Background(), + t.Context(), nil, func(context.Context, pgx.Tx) error { called = true @@ -57,8 +54,8 @@ func TestInSavepointNilFunc(t *testing.T) { t.Parallel() err := InSavepoint( - context.Background(), - &testTx{}, + t.Context(), + &savepointParentTx{}, nil, ) @@ -69,6 +66,152 @@ func TestInSavepointNilFunc(t *testing.T) { ) } +func TestInSavepoint(t *testing.T) { + t.Parallel() + + savepoint := &savepointTestTx{} + parent := &savepointParentTx{ + savepoint: savepoint, + } + + called := false + + err := InSavepoint( + t.Context(), + parent, + func(_ context.Context, tx pgx.Tx) error { + called = true + + if tx != savepoint { + t.Fatal("unexpected savepoint transaction") + } + + return nil + }, + ) + if err != nil { + t.Fatalf("InSavepoint returned an error: %v", err) + } + + if !called { + t.Fatal("savepoint function was not called") + } + + if savepoint.commitCalls != 1 { + t.Fatalf( + "commit calls = %d, want 1", + savepoint.commitCalls, + ) + } +} + +func TestInSavepointPreservesCallbackError(t *testing.T) { + t.Parallel() + + expectedErr := errors.New("callback failed") + savepoint := &savepointTestTx{} + parent := &savepointParentTx{ + savepoint: savepoint, + } + + err := InSavepoint( + t.Context(), + parent, + func(context.Context, pgx.Tx) error { + return expectedErr + }, + ) + if !errors.Is(err, expectedErr) { + t.Fatalf("original error was not preserved: %v", err) + } + + if savepoint.commitCalls != 0 { + t.Fatalf( + "commit calls = %d, want 0", + savepoint.commitCalls, + ) + } + + if savepoint.rollbackCalls == 0 { + t.Fatal("savepoint was not rolled back") + } +} + +func TestInSavepointPreservesBeginError(t *testing.T) { + t.Parallel() + + expectedErr := errors.New("begin failed") + parent := &savepointParentTx{ + beginErr: expectedErr, + } + + called := false + + err := InSavepoint( + t.Context(), + parent, + func(context.Context, pgx.Tx) error { + called = true + + return nil + }, + ) + if !errors.Is(err, expectedErr) { + t.Fatalf("original error was not preserved: %v", err) + } + + if called { + t.Fatal("savepoint function was called after begin failure") + } +} + +type savepointParentTx struct { + pgx.Tx + + savepoint pgx.Tx + beginErr error +} + +func (tx *savepointParentTx) Begin(context.Context) (pgx.Tx, error) { + if tx.beginErr != nil { + return nil, tx.beginErr + } + + return tx.savepoint, nil +} + +type savepointTestTx struct { + pgx.Tx + + closed bool + commitCalls int + rollbackCalls int +} + +func (tx *savepointTestTx) Commit(context.Context) error { + tx.commitCalls++ + + if tx.closed { + return pgx.ErrTxClosed + } + + tx.closed = true + + return nil +} + +func (tx *savepointTestTx) Rollback(context.Context) error { + tx.rollbackCalls++ + + if tx.closed { + return pgx.ErrTxClosed + } + + tx.closed = true + + return nil +} + func assertErrorMessage( t *testing.T, err error, @@ -81,6 +224,10 @@ func assertErrorMessage( } if err.Error() != want { - t.Fatalf("unexpected error: got %q, want %q", err, want) + t.Fatalf( + "error = %q, want %q", + err, + want, + ) } } From a5e565046adcb676c921bc739e851412aac444f3 Mon Sep 17 00:00:00 2001 From: mkbeh Date: Wed, 26 Aug 2026 00:53:21 +0300 Subject: [PATCH 39/41] ci: fix lint --- .github/workflows/lint.yml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/.github/workflows/lint.yml b/.github/workflows/lint.yml index 4f7bdc8..9f2564f 100644 --- a/.github/workflows/lint.yml +++ b/.github/workflows/lint.yml @@ -74,7 +74,7 @@ jobs: env: GOWORK: ${{ github.workspace }}/go.work with: - version: v2.12.2 + version: v2.13.1 working-directory: ${{ matrix.module }} args: >- --config=${{ github.workspace }}/.golangci.yml From 90ff06c182c260809202d2bed445ba2ebf23dc9e Mon Sep 17 00:00:00 2001 From: mkbeh Date: Wed, 26 Aug 2026 00:59:32 +0300 Subject: [PATCH 40/41] ci: discover lint modules dynamically --- .github/workflows/lint.yml | 49 ++++++++++++++++++++++++++++++-------- 1 file changed, 39 insertions(+), 10 deletions(-) diff --git a/.github/workflows/lint.yml b/.github/workflows/lint.yml index 9f2564f..22648e6 100644 --- a/.github/workflows/lint.yml +++ b/.github/workflows/lint.yml @@ -20,24 +20,53 @@ env: GOTOOLCHAIN: local jobs: + modules: + name: Discover modules + runs-on: ubuntu-latest + + outputs: + modules: ${{ steps.modules.outputs.modules }} + + steps: + - name: Check out repository + uses: actions/checkout@v7 + + - name: Discover Go modules + id: modules + shell: bash + run: | + set -euo pipefail + + modules="$( + { + printf '.\n' + + find ./extra ./examples \ + -name go.mod \ + -type f \ + -print | + while IFS= read -r mod; do + dirname "${mod#./}" + done + } | + sort | + jq -R -s -c 'split("\n")[:-1]' + )" + + echo "modules=${modules}" >> "${GITHUB_OUTPUT}" + lint: name: Lint (${{ matrix.module }}) + needs: + - modules + runs-on: ubuntu-latest timeout-minutes: 15 strategy: fail-fast: false matrix: - module: - - . - - extra/otelxpg - - extra/slogxpg - - examples/advisory - - examples/basic - - examples/cluster - - examples/otel - - examples/shard - - examples/transactions + module: ${{ fromJSON(needs.modules.outputs.modules) }} steps: - name: Check out repository From 6605de3b351df0c494650a6020f4700b7312abfd Mon Sep 17 00:00:00 2001 From: mkbeh Date: Wed, 26 Aug 2026 01:57:45 +0300 Subject: [PATCH 41/41] fix: support single-line requires in workspace generator --- .github/scripts/create-go-workspace.sh | 10 +++++++++- README.md | 2 +- examples/shard/go.mod | 2 +- examples/shard_geo/go.mod | 2 +- 4 files changed, 12 insertions(+), 4 deletions(-) diff --git a/.github/scripts/create-go-workspace.sh b/.github/scripts/create-go-workspace.sh index acca8ac..950ceee 100755 --- a/.github/scripts/create-go-workspace.sh +++ b/.github/scripts/create-go-workspace.sh @@ -72,7 +72,15 @@ done < <( } | while IFS= read -r -d '' mod; do awk ' - $1 ~ /^github\.com\/mkbeh\/xpg(\/.*)?$/ && $2 ~ /^v[0-9]/ { + $1 == "require" && + $2 ~ /^github\.com\/mkbeh\/xpg(\/.*)?$/ && + $3 ~ /^v[0-9]/ { + print $2, $3 + next + } + + $1 ~ /^github\.com\/mkbeh\/xpg(\/.*)?$/ && + $2 ~ /^v[0-9]/ { print $1, $2 } ' "${mod}" diff --git a/README.md b/README.md index a2b0b18..3c83c29 100644 --- a/README.md +++ b/README.md @@ -29,7 +29,7 @@ connection management, routing, and common production workflows. * **Primary/Replica Routing:** Explicit read policies, replica selection, primary fallback, and read-only transactions across PostgreSQL nodes. * **Application-Level Sharding:** Hash, range, time-based, and custom routing with colocation checks, key grouping, and - bounded concurrent fan-out. + bounded parallel operations across shards. * **Observability:** Structured logging, tracing, pool statistics, and optional OpenTelemetry metrics. ## Installation diff --git a/examples/shard/go.mod b/examples/shard/go.mod index 4a05837..81abd8b 100644 --- a/examples/shard/go.mod +++ b/examples/shard/go.mod @@ -1,4 +1,4 @@ -module github.com/mkbeh/xpg/examples/shard +module shard go 1.27 diff --git a/examples/shard_geo/go.mod b/examples/shard_geo/go.mod index 320f239..7f51df5 100644 --- a/examples/shard_geo/go.mod +++ b/examples/shard_geo/go.mod @@ -1,4 +1,4 @@ -module github.com/mkbeh/xpg/examples/shard_geo +module shard_geo go 1.27