126 lines
2.8 KiB
Go
126 lines
2.8 KiB
Go
|
package cassandra
|
|||
|
|
|
|||
|
|
import (
|
|||
|
|
"context"
|
|||
|
|
"fmt"
|
|||
|
|
"github.com/scylladb/gocqlx/v3/qb"
|
|||
|
|
"time"
|
|||
|
|
)
|
|||
|
|
|
|||
|
|
const (
|
|||
|
|
defaultTTLSec = 30
|
|||
|
|
defaultRetry = 3
|
|||
|
|
baseDelay = 100 * time.Millisecond
|
|||
|
|
)
|
|||
|
|
|
|||
|
|
// LockOption 用來設定 TryLock 的 TTL 行為
|
|||
|
|
type LockOption func(*lockOptions)
|
|||
|
|
|
|||
|
|
type lockOptions struct {
|
|||
|
|
ttlSeconds int // TTL，單位秒；<=0 代表不 expire
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
func WithLockTTL(d time.Duration) LockOption {
|
|||
|
|
return func(o *lockOptions) {
|
|||
|
|
o.ttlSeconds = int(d.Seconds())
|
|||
|
|
}
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
// WithNoLockExpire 永不自動解鎖
|
|||
|
|
func WithNoLockExpire() LockOption {
|
|||
|
|
return func(o *lockOptions) {
|
|||
|
|
o.ttlSeconds = 0
|
|||
|
|
}
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
// TryLock 嘗試在表上插入一筆唯一鍵（IF NOT EXISTS）作為鎖
|
|||
|
|
// 預設 30 秒 TTL，可透過 option 調整或取消 TTL
|
|||
|
|
func (db *DB) TryLock(
|
|||
|
|
ctx context.Context,
|
|||
|
|
document any,
|
|||
|
|
keyspace string,
|
|||
|
|
opts ...LockOption,
|
|||
|
|
) error {
|
|||
|
|
// 1. 解析 metadata
|
|||
|
|
metadata, err := GenerateTableMetadata(document, keyspace)
|
|||
|
|
if err != nil {
|
|||
|
|
return err
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
// 2. 組合 option
|
|||
|
|
options := &lockOptions{ttlSeconds: defaultTTLSec}
|
|||
|
|
for _, opt := range opts {
|
|||
|
|
opt(options)
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
// 3. 建 TTL 子句
|
|||
|
|
builder := qb.Insert(metadata.Name).
|
|||
|
|
Unique(). // IF NOT EXISTS
|
|||
|
|
Columns(metadata.Columns...)
|
|||
|
|
|
|||
|
|
if options.ttlSeconds > 0 {
|
|||
|
|
ttl := time.Duration(options.ttlSeconds) * time.Second
|
|||
|
|
builder = builder.TTL(ttl)
|
|||
|
|
}
|
|||
|
|
stmt, names := builder.ToCql()
|
|||
|
|
|
|||
|
|
// 4. 執行 CAS
|
|||
|
|
q := db.GetSession().Query(stmt, names).BindStruct(document).
|
|||
|
|
WithContext(ctx).
|
|||
|
|
WithTimestamp(time.Now().UnixNano() / 1e3)
|
|||
|
|
|
|||
|
|
applied, err := q.ExecCASRelease()
|
|||
|
|
if err != nil {
|
|||
|
|
return err
|
|||
|
|
}
|
|||
|
|
if !applied {
|
|||
|
|
return fmt.Errorf("failed to acquire lock")
|
|||
|
|
}
|
|||
|
|
return nil
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
// UnLock 砍掉鎖，其實就是 Delete
|
|||
|
|
func (db *DB) UnLock(ctx context.Context, filter any, keyspace string) error {
|
|||
|
|
if filter == nil {
|
|||
|
|
return fmt.Errorf("unlock failed: nil filter")
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
metadata, err := GenerateTableMetadata(filter, keyspace)
|
|||
|
|
if err != nil {
|
|||
|
|
return fmt.Errorf("unlock: generate metadata failed: %w", err)
|
|||
|
|
}
|
|||
|
|
if len(metadata.Columns) == 0 {
|
|||
|
|
return fmt.Errorf("unlock failed: missing primary key in struct")
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
var lastErr error
|
|||
|
|
|
|||
|
|
for i := 0; i < defaultRetry; i++ {
|
|||
|
|
builder := qb.Delete(metadata.Name).Existing()
|
|||
|
|
|
|||
|
|
// 動態添加 WHERE 條件
|
|||
|
|
for _, key := range metadata.PartKey {
|
|||
|
|
builder = builder.Where(qb.Eq(key))
|
|||
|
|
}
|
|||
|
|
stmt, names := builder.ToCql()
|
|||
|
|
q := db.GetSession().Query(stmt, names).BindStruct(filter).
|
|||
|
|
WithContext(ctx).
|
|||
|
|
WithTimestamp(time.Now().UnixNano() / 1e3)
|
|||
|
|
|
|||
|
|
applied, err := q.ExecCASRelease()
|
|||
|
|
if err == nil && applied {
|
|||
|
|
return nil
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
if err != nil {
|
|||
|
|
lastErr = fmt.Errorf("unlock error: %w", err)
|
|||
|
|
} else if !applied {
|
|||
|
|
lastErr = fmt.Errorf("unlock not applied: row not found or not visible yet")
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
time.Sleep(baseDelay * time.Duration(1<<i)) // 100ms → 200ms → 400ms
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
return fmt.Errorf("unlock failed after retries: %w", lastErr)
|
|||
|
|
}
|