draincloud-core/internal/storage/postgres/sharded_cluster.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
}