Skip to content
Open
Show file tree
Hide file tree
Changes from 2 commits
Commits
File filter

Filter by extension

Filter by extension


Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
7 changes: 5 additions & 2 deletions Makefile
Original file line number Diff line number Diff line change
Expand Up @@ -43,11 +43,14 @@ test:
test-integration:
go test -tags=integration -p 1 ./... -race $(GOTESTFLAGS)

COMPOSE_PROFILE ?=
compose_profile_args = $(foreach p,$(COMPOSE_PROFILE),--profile $(p))

compose-up:
$(DOCKER_COMPOSE) up -d --wait
$(DOCKER_COMPOSE) $(compose_profile_args) up -d --wait

compose-down:
$(DOCKER_COMPOSE) down
$(DOCKER_COMPOSE) $(compose_profile_args) down

lint:
@command -v golangci-lint >/dev/null 2>&1 || { \
Expand Down
227 changes: 227 additions & 0 deletions cmd/workload/main.go
Original file line number Diff line number Diff line change
@@ -0,0 +1,227 @@
package main

import (
"context"
"fmt"
"log"
"math/rand/v2"
"os"
"os/signal"
"sync"
"syscall"
"time"

"github.com/jackc/pgx/v5/pgxpool"
)

const (
// defaultDSN points at the local compose database; password is the dev one.
defaultDSN = "postgres://postgres:pass@postgres:5432/dev" //nolint:gosec // dev-only default DSN
poolSize = 20
seedRows = 1000
normalWorkers = 8
extraWorkers = 4 // nPlusOne, longQueries, blocker, contender
hotRowID = 1 // the single row the blocker and contenders fight over
)

func main() {
if err := run(); err != nil {
log.Fatalf("workload: %v", err)
}
}

func run() error {
ctx, stop := signal.NotifyContext(context.Background(), os.Interrupt, syscall.SIGTERM)
defer stop()

dsn := os.Getenv("WORKLOAD_DSN")
if dsn == "" {
dsn = defaultDSN
}

pool, err := connect(ctx, dsn)
if err != nil {
return err
}
defer pool.Close()

if err := setup(ctx, pool); err != nil {
return err
}

log.Println("workload: generating traffic (ctrl-c to stop)")

workers := make([]func(context.Context, *pgxpool.Pool), 0, normalWorkers+extraWorkers)
for range normalWorkers {
workers = append(workers, normalLoad)
}

workers = append(workers, nPlusOne, longQueries, blocker, contender)

var wg sync.WaitGroup

for _, w := range workers {
wg.Go(func() {
w(ctx, pool)
})
}
Comment thread
mickamy marked this conversation as resolved.
Comment thread
mickamy marked this conversation as resolved.
Comment thread
mickamy marked this conversation as resolved.

<-ctx.Done()
wg.Wait()
log.Println("workload: stopped")

return nil
}

func connect(ctx context.Context, dsn string) (*pgxpool.Pool, error) {
cfg, err := pgxpool.ParseConfig(dsn)
if err != nil {
return nil, fmt.Errorf("parse dsn: %w", err)
}

cfg.MaxConns = poolSize

for {
pool, err := pgxpool.NewWithConfig(ctx, cfg)
if err == nil {
if pingErr := pool.Ping(ctx); pingErr == nil {
return pool, nil
} else {
pool.Close()
err = pingErr
}
}

log.Printf("workload: waiting for db: %v", err)

if !wait(ctx, time.Second) {
return nil, fmt.Errorf("canceled while connecting: %w", ctx.Err())
}
}
}

func setup(ctx context.Context, pool *pgxpool.Pool) error {
// pg_stat_statements needs shared_preload_libraries; tolerate its absence.
if _, err := pool.Exec(ctx, `CREATE EXTENSION IF NOT EXISTS pg_stat_statements`); err != nil {
log.Printf("workload: pg_stat_statements unavailable: %v", err)
}

schema := []string{
`CREATE TABLE IF NOT EXISTS items (
id int PRIMARY KEY,
name text NOT NULL,
value int NOT NULL
)`,
fmt.Sprintf(`INSERT INTO items
SELECT g, 'item-' || g, g
FROM generate_series(1, %d) g
ON CONFLICT (id) DO NOTHING`, seedRows),
}

for _, stmt := range schema {
if _, err := pool.Exec(ctx, stmt); err != nil {
return fmt.Errorf("setup: %w", err)
}
}

return nil
}

// normalLoad does steady point reads and writes across random rows.
func normalLoad(ctx context.Context, pool *pgxpool.Pool) {
for wait(ctx, jitter(5*time.Millisecond, 45*time.Millisecond)) {
id := rand.IntN(seedRows) + 1 //nolint:gosec // weak RNG is fine for a load generator

var (
name string
value int
)

_ = pool.QueryRow(ctx, `SELECT name, value FROM items WHERE id = $1`, id).Scan(&name, &value)
_, _ = pool.Exec(ctx, `UPDATE items SET value = value + 1 WHERE id = $1`, id)
}
}

// nPlusOne fetches a batch of ids then queries each one separately, the classic
// pattern that is cheap per call but dominates pg_stat_statements by count.
func nPlusOne(ctx context.Context, pool *pgxpool.Pool) {
for wait(ctx, 500*time.Millisecond) {
rows, err := pool.Query(ctx, `SELECT id FROM items ORDER BY random() LIMIT 50`)
if err != nil {
continue
}

var ids []int

for rows.Next() {
var id int
if err := rows.Scan(&id); err == nil {
ids = append(ids, id)
}
}

rows.Close()

for _, id := range ids {
var name string

_ = pool.QueryRow(ctx, `SELECT name FROM items WHERE id = $1`, id).Scan(&name)
}
}
}

// longQueries runs an occasional multi-second query so the Activity screen has
// a long DURATION to surface.
func longQueries(ctx context.Context, pool *pgxpool.Pool) {
for wait(ctx, 2*time.Second) {
seconds := rand.IntN(8) + 3 //nolint:gosec // weak RNG is fine for a load generator

_, _ = pool.Exec(ctx, `SELECT pg_sleep($1)`, seconds)
}
}

// blocker holds a row lock inside an open transaction, then sits idle, creating
// an idle-in-transaction backend that blocks the contender.
func blocker(ctx context.Context, pool *pgxpool.Pool) {
for wait(ctx, 3*time.Second) {
holdLock(ctx, pool)
}
}

func holdLock(ctx context.Context, pool *pgxpool.Pool) {
tx, err := pool.Begin(ctx)
if err != nil {
return
}
defer func() { _ = tx.Rollback(ctx) }()

if _, err := tx.Exec(ctx, `UPDATE items SET value = value WHERE id = $1`, hotRowID); err != nil {
return
}

wait(ctx, 5*time.Second)

_ = tx.Commit(ctx)
}

// contender repeatedly updates the hot row, so it blocks whenever the blocker
// is holding the lock.
func contender(ctx context.Context, pool *pgxpool.Pool) {
for wait(ctx, time.Second) {
_, _ = pool.Exec(ctx, `UPDATE items SET value = value + 1 WHERE id = $1`, hotRowID)
}
}

// wait sleeps for d, returning false if the context is canceled first.
func wait(ctx context.Context, d time.Duration) bool {
select {
case <-ctx.Done():
return false
case <-time.After(d):
return true
}
}

func jitter(minimum, maximum time.Duration) time.Duration {
return minimum + rand.N(maximum-minimum) //nolint:gosec // weak RNG is fine for a load generator
}
19 changes: 19 additions & 0 deletions compose.yaml
Original file line number Diff line number Diff line change
@@ -1,11 +1,14 @@
services:
postgres:
image: postgres:16
command: ["postgres", "-c", "shared_preload_libraries=pg_stat_statements"]
environment:
POSTGRES_PASSWORD: pass
POSTGRES_DB: dev
ports:
- "5432:5432"
volumes:
- ./compose/initdb:/docker-entrypoint-initdb.d:ro
healthcheck:
test: ["CMD-SHELL", "pg_isready -U postgres"]
interval: 2s
Expand All @@ -24,3 +27,19 @@ services:
interval: 2s
timeout: 3s
retries: 10

# Synthetic workload for developing dbtop. Opt-in via the "load" profile:
# make compose-up COMPOSE_PROFILE=load
workload:
image: golang:1.26-alpine
profiles: ["load"]
network_mode: "service:postgres"
working_dir: /src
command: ["go", "run", "./cmd/workload"]
environment:
WORKLOAD_DSN: postgres://postgres:pass@localhost:5432/dev
volumes:
- .:/src:ro
depends_on:
postgres:
condition: service_healthy
4 changes: 4 additions & 0 deletions compose/initdb/init.sql
Original file line number Diff line number Diff line change
@@ -0,0 +1,4 @@
-- Runs once on a fresh data directory (after postgres starts with
-- shared_preload_libraries=pg_stat_statements), so the Statements screen works
-- out of the box in local and CI databases.
CREATE EXTENSION IF NOT EXISTS pg_stat_statements;
26 changes: 25 additions & 1 deletion go.mod
Original file line number Diff line number Diff line change
Expand Up @@ -2,12 +2,36 @@ module github.com/mickamy/dbtop

go 1.26.3

require github.com/jackc/pgx/v5 v5.9.2
require (
github.com/charmbracelet/bubbles v1.0.0
github.com/charmbracelet/bubbletea v1.3.10
github.com/charmbracelet/lipgloss v1.1.0
github.com/jackc/pgx/v5 v5.9.2
)

require (
github.com/aymanbagabas/go-osc52/v2 v2.0.1 // indirect
github.com/charmbracelet/colorprofile v0.4.1 // indirect
github.com/charmbracelet/x/ansi v0.11.6 // indirect
github.com/charmbracelet/x/cellbuf v0.0.15 // indirect
github.com/charmbracelet/x/term v0.2.2 // indirect
github.com/clipperhouse/displaywidth v0.9.0 // indirect
github.com/clipperhouse/stringish v0.1.1 // indirect
github.com/clipperhouse/uax29/v2 v2.5.0 // indirect
github.com/erikgeiser/coninput v0.0.0-20211004153227-1c3628e74d0f // 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/lucasb-eyer/go-colorful v1.3.0 // indirect
github.com/mattn/go-isatty v0.0.20 // indirect
github.com/mattn/go-localereader v0.0.1 // indirect
github.com/mattn/go-runewidth v0.0.19 // indirect
github.com/muesli/ansi v0.0.0-20230316100256-276c6243b2f6 // indirect
github.com/muesli/cancelreader v0.2.2 // indirect
github.com/muesli/termenv v0.16.0 // indirect
github.com/rivo/uniseg v0.4.7 // indirect
github.com/xo/terminfo v0.0.0-20220910002029-abceb7e1c41e // indirect
golang.org/x/sync v0.17.0 // indirect
golang.org/x/sys v0.38.0 // indirect
golang.org/x/text v0.29.0 // indirect
)
48 changes: 48 additions & 0 deletions go.sum
Original file line number Diff line number Diff line change
@@ -1,6 +1,30 @@
github.com/aymanbagabas/go-osc52/v2 v2.0.1 h1:HwpRHbFMcZLEVr42D4p7XBqjyuxQH5SMiErDT4WkJ2k=
github.com/aymanbagabas/go-osc52/v2 v2.0.1/go.mod h1:uYgXzlJ7ZpABp8OJ+exZzJJhRNQ2ASbcXHWsFqH8hp8=
github.com/charmbracelet/bubbles v1.0.0 h1:12J8/ak/uCZEMQ6KU7pcfwceyjLlWsDLAxB5fXonfvc=
github.com/charmbracelet/bubbles v1.0.0/go.mod h1:9d/Zd5GdnauMI5ivUIVisuEm3ave1XwXtD1ckyV6r3E=
github.com/charmbracelet/bubbletea v1.3.10 h1:otUDHWMMzQSB0Pkc87rm691KZ3SWa4KUlvF9nRvCICw=
github.com/charmbracelet/bubbletea v1.3.10/go.mod h1:ORQfo0fk8U+po9VaNvnV95UPWA1BitP1E0N6xJPlHr4=
github.com/charmbracelet/colorprofile v0.4.1 h1:a1lO03qTrSIRaK8c3JRxJDZOvhvIeSco3ej+ngLk1kk=
github.com/charmbracelet/colorprofile v0.4.1/go.mod h1:U1d9Dljmdf9DLegaJ0nGZNJvoXAhayhmidOdcBwAvKk=
github.com/charmbracelet/lipgloss v1.1.0 h1:vYXsiLHVkK7fp74RkV7b2kq9+zDLoEU4MZoFqR/noCY=
github.com/charmbracelet/lipgloss v1.1.0/go.mod h1:/6Q8FR2o+kj8rz4Dq0zQc3vYf7X+B0binUUBwA0aL30=
github.com/charmbracelet/x/ansi v0.11.6 h1:GhV21SiDz/45W9AnV2R61xZMRri5NlLnl6CVF7ihZW8=
github.com/charmbracelet/x/ansi v0.11.6/go.mod h1:2JNYLgQUsyqaiLovhU2Rv/pb8r6ydXKS3NIttu3VGZQ=
github.com/charmbracelet/x/cellbuf v0.0.15 h1:ur3pZy0o6z/R7EylET877CBxaiE1Sp1GMxoFPAIztPI=
github.com/charmbracelet/x/cellbuf v0.0.15/go.mod h1:J1YVbR7MUuEGIFPCaaZ96KDl5NoS0DAWkskup+mOY+Q=
github.com/charmbracelet/x/term v0.2.2 h1:xVRT/S2ZcKdhhOuSP4t5cLi5o+JxklsoEObBSgfgZRk=
github.com/charmbracelet/x/term v0.2.2/go.mod h1:kF8CY5RddLWrsgVwpw4kAa6TESp6EB5y3uxGLeCqzAI=
github.com/clipperhouse/displaywidth v0.9.0 h1:Qb4KOhYwRiN3viMv1v/3cTBlz3AcAZX3+y9OLhMtAtA=
github.com/clipperhouse/displaywidth v0.9.0/go.mod h1:aCAAqTlh4GIVkhQnJpbL0T/WfcrJXHcj8C0yjYcjOZA=
github.com/clipperhouse/stringish v0.1.1 h1:+NSqMOr3GR6k1FdRhhnXrLfztGzuG+VuFDfatpWHKCs=
github.com/clipperhouse/stringish v0.1.1/go.mod h1:v/WhFtE1q0ovMta2+m+UbpZ+2/HEXNWYXQgCt4hdOzA=
github.com/clipperhouse/uax29/v2 v2.5.0 h1:x7T0T4eTHDONxFJsL94uKNKPHrclyFI0lm7+w94cO8U=
github.com/clipperhouse/uax29/v2 v2.5.0/go.mod h1:Wn1g7MK6OoeDT0vL+Q0SQLDz/KpfsVRgg6W7ihQeh4g=
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/erikgeiser/coninput v0.0.0-20211004153227-1c3628e74d0f h1:Y/CXytFA4m6baUTXGLOoWe4PQhGxaX0KpnayAqC48p4=
github.com/erikgeiser/coninput v0.0.0-20211004153227-1c3628e74d0f/go.mod h1:vw97MGsxSvLiUE2X8qFplwetxpGLQrlU1Q9AUEIzCaM=
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=
Expand All @@ -9,15 +33,39 @@ 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/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/lucasb-eyer/go-colorful v1.3.0 h1:2/yBRLdWBZKrf7gB40FoiKfAWYQ0lqNcbuQwVHXptag=
github.com/lucasb-eyer/go-colorful v1.3.0/go.mod h1:R4dSotOR9KMtayYi1e77YzuveK+i7ruzyGqttikkLy0=
github.com/mattn/go-isatty v0.0.20 h1:xfD0iDuEKnDkl03q4limB+vH+GxLEtL/jb4xVJSWWEY=
github.com/mattn/go-isatty v0.0.20/go.mod h1:W+V8PltTTMOvKvAeJH7IuucS94S2C6jfK/D7dTCTo3Y=
github.com/mattn/go-localereader v0.0.1 h1:ygSAOl7ZXTx4RdPYinUpg6W99U8jWvWi9Ye2JC/oIi4=
github.com/mattn/go-localereader v0.0.1/go.mod h1:8fBrzywKY7BI3czFoHkuzRoWE9C+EiG4R1k4Cjx5p88=
github.com/mattn/go-runewidth v0.0.19 h1:v++JhqYnZuu5jSKrk9RbgF5v4CGUjqRfBm05byFGLdw=
github.com/mattn/go-runewidth v0.0.19/go.mod h1:XBkDxAl56ILZc9knddidhrOlY5R/pDhgLpndooCuJAs=
github.com/muesli/ansi v0.0.0-20230316100256-276c6243b2f6 h1:ZK8zHtRHOkbHy6Mmr5D264iyp3TiX5OmNcI5cIARiQI=
github.com/muesli/ansi v0.0.0-20230316100256-276c6243b2f6/go.mod h1:CJlz5H+gyd6CUWT45Oy4q24RdLyn7Md9Vj2/ldJBSIo=
github.com/muesli/cancelreader v0.2.2 h1:3I4Kt4BQjOR54NavqnDogx/MIoWBFa0StPA8ELUXHmA=
github.com/muesli/cancelreader v0.2.2/go.mod h1:3XuTXfFS2VjM+HTLZY9Ak0l6eUKfijIfMUZ4EgX0QYo=
github.com/muesli/termenv v0.16.0 h1:S5AlUN9dENB57rsbnkPyfdGuWIlkmzJjbFf0Tf5FWUc=
github.com/muesli/termenv v0.16.0/go.mod h1:ZRfOIKPFDYQoDFF4Olj7/QJbW60Ol/kL1pU3VfY/Cnk=
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/rivo/uniseg v0.4.7 h1:WUdvkW8uEhrYfLC4ZzdpI2ztxP1I582+49Oc5Mq64VQ=
github.com/rivo/uniseg v0.4.7/go.mod h1:FN3SvrM+Zdj16jyLfmOkMNblXMcoc8DfTHruCPUcx88=
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/xo/terminfo v0.0.0-20220910002029-abceb7e1c41e h1:JVG44RsyaB9T2KIHavMF/ppJZNG9ZpyihvCd0w101no=
github.com/xo/terminfo v0.0.0-20220910002029-abceb7e1c41e/go.mod h1:RbqR21r5mrJuqunuUZ/Dhy/avygyECGrLceyNeo4LiM=
golang.org/x/exp v0.0.0-20231006140011-7918f672742d h1:jtJma62tbqLibJ5sFQz8bKtEM8rJBtfilJ2qTU199MI=
golang.org/x/exp v0.0.0-20231006140011-7918f672742d/go.mod h1:ldy0pHrwJyGW56pPQzzkH36rKxoZW1tw7ZJpeKx+hdo=
golang.org/x/sync v0.17.0 h1:l60nONMj9l5drqw6jlhIELNv9I0A4OFgRsG9k2oT9Ug=
golang.org/x/sync v0.17.0/go.mod h1:9KTHXmSnoGruLpwFjVSX0lNNA75CykiMECbovNTZqGI=
golang.org/x/sys v0.0.0-20210809222454-d867a43fc93e/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
golang.org/x/sys v0.6.0/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
golang.org/x/sys v0.38.0 h1:3yZWxaJjBmCWXqhN1qh02AkOnCQ1poK6oF+a7xWL6Gc=
golang.org/x/sys v0.38.0/go.mod h1:OgkHotnGiDImocRcuBABYBEXf8A9a87e/uXjp9XT3ks=
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=
Expand Down
Loading
Loading