freedakgmail
45 lines
1.3 KiB
Go
45 lines
1.3 KiB
Go
package cccode
|
|
|
|
import (
|
|
"database/sql"
|
|
"fmt"
|
|
)
|
|
|
|
// PostgresStore 是 AllocationStore 的 PostgreSQL 实现。
|
|
// 通过行级原子 UPSERT + 返回值保证多实例下序列分配原子、持久、不重号。
|
|
// 解决 MemoryStore 重启丢号、多实例重号的问题（生产用）。
|
|
type PostgresStore struct {
|
|
db *sql.DB
|
|
}
|
|
|
|
// NewPostgresStore 创建基于 *sql.DB 的分配存储。
|
|
func NewPostgresStore(db *sql.DB) *PostgresStore {
|
|
return &PostgresStore{db: db}
|
|
}
|
|
|
|
// Next 原子获取下一个序列。
|
|
// 使用 INSERT ... ON CONFLICT DO UPDATE 的原子自增语义：
|
|
// - 首次：cursor = start
|
|
// - 之后：cursor = cursor + 1
|
|
//
|
|
// 单条 SQL 在行锁内完成读改写，并发安全；超过 end 返回耗尽错误。
|
|
func (s *PostgresStore) Next(segmentKey string, start, end uint64) (uint64, error) {
|
|
const q = `
|
|
INSERT INTO cccode_cursor (segment_key, cursor)
|
|
VALUES ($1, $2)
|
|
ON CONFLICT (segment_key)
|
|
DO UPDATE SET cursor = cccode_cursor.cursor + 1, updated_at = NOW()
|
|
RETURNING cursor;`
|
|
|
|
var next uint64
|
|
if err := s.db.QueryRow(q, segmentKey, start).Scan(&next); err != nil {
|
|
return 0, fmt.Errorf("cccode: pg next: %w", err)
|
|
}
|
|
if next > end {
|
|
return 0, ErrSegmentExhausted
|
|
}
|
|
return next, nil
|
|
}
|
|
|
|
var _ AllocationStore = (*PostgresStore)(nil)
|