42 lines
903 B
Go
42 lines
903 B
Go
package postgres
|
|
|
|
import (
|
|
"context"
|
|
"hash/crc32"
|
|
"log/slog"
|
|
"sync"
|
|
|
|
"git.optclblast.xyz/draincloud/draincloud-core/internal/logger"
|
|
"github.com/google/uuid"
|
|
"github.com/jackc/pgx/v5"
|
|
)
|
|
|
|
type ShardMap = map[uint32]*pgx.ConnConfig
|
|
|
|
type ShardCluster struct {
|
|
m sync.Mutex
|
|
shards []*pgx.Conn
|
|
}
|
|
|
|
func NewShardCluster(ctx context.Context, shardMap ShardMap) *ShardCluster {
|
|
shards := make([]*pgx.Conn, len(shardMap))
|
|
for n, cfg := range shardMap {
|
|
conn, err := pgx.ConnectConfig(ctx, cfg)
|
|
if err != nil {
|
|
logger.Fatal(ctx, "failed to connect to shard", slog.Uint64("num", uint64(n)), logger.Err(err))
|
|
}
|
|
shards[n] = conn
|
|
}
|
|
return &ShardCluster{shards: shards}
|
|
}
|
|
|
|
func (c *ShardCluster) PickShard(n uint32) *pgx.Conn {
|
|
c.m.Lock()
|
|
defer c.m.Unlock()
|
|
return c.shards[n]
|
|
}
|
|
|
|
func UUIDShardFn(id uuid.UUID, numShards uint32) uint32 {
|
|
return crc32.ChecksumIEEE(id[:]) % numShards
|
|
}
|