my example was not good tbh
Bel LaPointe
parent
0db7008aac
commit
faa128400f
|
|
@ -0,0 +1 @@
|
||||||
|
https://topicpartition.io/blog/postgres-pubsub-queue-benchmarks
|
||||||
|
|
@ -0,0 +1,162 @@
|
||||||
|
package main
|
||||||
|
|
||||||
|
import (
|
||||||
|
"context"
|
||||||
|
"database/sql"
|
||||||
|
"encoding/json"
|
||||||
|
"flag"
|
||||||
|
"fmt"
|
||||||
|
"log"
|
||||||
|
"math/rand"
|
||||||
|
"os"
|
||||||
|
"strconv"
|
||||||
|
"time"
|
||||||
|
|
||||||
|
"pg/src/with"
|
||||||
|
|
||||||
|
_ "github.com/lib/pq"
|
||||||
|
)
|
||||||
|
|
||||||
|
func main() {
|
||||||
|
if err := with.Context(run); err != nil {
|
||||||
|
panic(err)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func run(ctx context.Context) error {
|
||||||
|
fs := flag.NewFlagSet(os.Args[0], flag.ContinueOnError)
|
||||||
|
c := fs.String("c", "postgresql://pulsegres:pulsegres@localhost:15432", "conn string")
|
||||||
|
if err := fs.Parse(os.Args[1:]); err != nil {
|
||||||
|
panic(err)
|
||||||
|
}
|
||||||
|
|
||||||
|
return with.PSQL(ctx, *c, func(pg *sql.DB) error {
|
||||||
|
log.Println("creating tables...")
|
||||||
|
for _, sql := range []string{
|
||||||
|
`CREATE SCHEMA IF NOT EXISTS "pubsub"`,
|
||||||
|
|
||||||
|
`CREATE TABLE IF NOT EXISTS "pubsub"."topic_a" (
|
||||||
|
partition INT PRIMARY KEY,
|
||||||
|
partition_offset BIGINT NOT NULL DEFAULT 0
|
||||||
|
)`,
|
||||||
|
`CREATE TABLE IF NOT EXISTS "pubsub"."topic_a_partition_0" (
|
||||||
|
id UUID NOT NULL DEFAULT GEN_RANDOM_UUID() PRIMARY KEY,
|
||||||
|
partition_offset BIGINT UNIQUE NOT NULL,
|
||||||
|
payload BYTEA NOT NULL,
|
||||||
|
created_at TIMESTAMP NOT NULL DEFAULT NOW()
|
||||||
|
)`,
|
||||||
|
`CREATE TABLE IF NOT EXISTS "pubsub"."topic_a_partition_1" (
|
||||||
|
id UUID NOT NULL DEFAULT GEN_RANDOM_UUID() PRIMARY KEY,
|
||||||
|
partition_offset BIGINT UNIQUE NOT NULL,
|
||||||
|
payload BYTEA NOT NULL,
|
||||||
|
created_at TIMESTAMP NOT NULL DEFAULT NOW()
|
||||||
|
)`,
|
||||||
|
|
||||||
|
`INSERT INTO "pubsub"."topic_a" (partition, partition_offset) VALUES (0, 0) ON CONFLICT DO NOTHING`,
|
||||||
|
`INSERT INTO "pubsub"."topic_a" (partition, partition_offset) VALUES (1, 0) ON CONFLICT DO NOTHING`,
|
||||||
|
|
||||||
|
`CREATE TABLE IF NOT EXISTS "pubsub"."consumer_partition_offset" (
|
||||||
|
consumer_group TEXT NOT NULL,
|
||||||
|
partition INT REFERENCES "pubsub"."topic_a"("partition"),
|
||||||
|
partition_offset BIGINT NOT NULL DEFAULT 0,
|
||||||
|
PRIMARY KEY (consumer_group, partition)
|
||||||
|
)`,
|
||||||
|
} {
|
||||||
|
if _, err := pg.ExecContext(ctx, sql); err != nil {
|
||||||
|
return fmt.Errorf("failed setup sql %s: %w", sql, err)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
pub := func(a any) error {
|
||||||
|
b, _ := json.Marshal(a)
|
||||||
|
partition := rand.Int() % 2
|
||||||
|
partitionS := strconv.Itoa(partition)
|
||||||
|
_, err := pg.ExecContext(ctx, `
|
||||||
|
WITH reserve AS (
|
||||||
|
UPDATE "pubsub"."topic_a"
|
||||||
|
SET partition_offset=partition_offset+1
|
||||||
|
WHERE partition=`+partitionS+`
|
||||||
|
RETURNING (partition_offset-1) AS offset
|
||||||
|
)
|
||||||
|
INSERT INTO "pubsub"."topic_a_partition_`+partitionS+`" (partition_offset, payload)
|
||||||
|
SELECT reserve.offset, $1 FROM reserve
|
||||||
|
`, b)
|
||||||
|
return err
|
||||||
|
}
|
||||||
|
sub := func(group string) (any, error) {
|
||||||
|
partition := rand.Int() % 2
|
||||||
|
partitionS := strconv.Itoa(partition)
|
||||||
|
|
||||||
|
if _, err := pg.ExecContext(ctx, `
|
||||||
|
INSERT INTO
|
||||||
|
"pubsub"."consumer_partition_offset"
|
||||||
|
(consumer_group, partition)
|
||||||
|
VALUES
|
||||||
|
($1, $2)
|
||||||
|
ON CONFLICT DO NOTHING
|
||||||
|
`, group, partition); err != nil {
|
||||||
|
return nil, err
|
||||||
|
}
|
||||||
|
|
||||||
|
row := pg.QueryRowContext(ctx, `
|
||||||
|
WITH tip AS (
|
||||||
|
SELECT partition_offset-1
|
||||||
|
FROM "pubsub"."topic_a"
|
||||||
|
WHERE partition=`+partitionS+`
|
||||||
|
)
|
||||||
|
, to_claim AS (
|
||||||
|
SELECT partition_offset
|
||||||
|
FROM "pubsub"."consumer_partition_offset"
|
||||||
|
WHERE consumer_group=$1 AND partition=$2
|
||||||
|
FOR UPDATE SKIP LOCKED
|
||||||
|
)
|
||||||
|
, claimed AS (
|
||||||
|
UPDATE "pubsub"."consumer_partition_offset" consumer_partition_offset
|
||||||
|
SET consumer_partition_offset.partition_offset=to_claim.partition_offset+1
|
||||||
|
FROM to_claim
|
||||||
|
WHERE consumer_partition_offset.consumer_group=$1 AND consumer_partition_offset.partition=$2
|
||||||
|
RETURNING consumer_partition_offset.partition_offset
|
||||||
|
)
|
||||||
|
SELECT partition_offset FROM claimed
|
||||||
|
`, group, partition)
|
||||||
|
if err := row.Err(); err != nil {
|
||||||
|
return nil, err
|
||||||
|
}
|
||||||
|
|
||||||
|
var result any
|
||||||
|
if err := row.Scan(&result); err != nil {
|
||||||
|
return nil, err
|
||||||
|
}
|
||||||
|
|
||||||
|
return result, row.Err()
|
||||||
|
}
|
||||||
|
|
||||||
|
go with.Every(ctx, time.Second, func() {
|
||||||
|
for i := 0; i < 2; i++ {
|
||||||
|
if err := pub(1); err != nil {
|
||||||
|
log.Printf("failed pub: %v", err)
|
||||||
|
} else {
|
||||||
|
log.Printf("pubbed")
|
||||||
|
}
|
||||||
|
}
|
||||||
|
})
|
||||||
|
|
||||||
|
with.Every(ctx, time.Second, func() {
|
||||||
|
if err := func() error {
|
||||||
|
for i := 0; i < 2; i++ {
|
||||||
|
group := fmt.Sprintf("group_%d", i)
|
||||||
|
if v, err := sub(group); err != nil {
|
||||||
|
return err
|
||||||
|
} else {
|
||||||
|
log.Printf("%s => %v", group, v)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
return nil
|
||||||
|
}(); err != nil {
|
||||||
|
log.Printf("failed subs: %v", err)
|
||||||
|
}
|
||||||
|
})
|
||||||
|
|
||||||
|
return ctx.Err()
|
||||||
|
})
|
||||||
|
}
|
||||||
Loading…
Reference in New Issue