From 2697ba7ac09ae03501cc827ae4bc6bfe40f8f392 Mon Sep 17 00:00:00 2001 From: smail Date: Tue, 6 Oct 2026 03:12:46 +0800 Subject: [PATCH] =?UTF-8?q?=E4=BF=AE=E5=A4=8D=E9=AA=8C=E8=AF=81=E7=A0=81?= =?UTF-8?q?=E6=9C=8D=E5=8A=A1=E5=9C=A8=20PostgreSQL=20=E4=B8=8B=E7=9A=84?= =?UTF-8?q?=E4=BA=8B=E5=8A=A1=E8=A1=8C=E9=94=81=E6=9F=A5=E8=AF=A2?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- backend/services/verification/event_store.go | 8 +++--- backend/services/verification/rate_limit.go | 14 +++------- backend/services/verification/row_lock.go | 27 +++++++++++++++++++ .../services/verification/row_lock_test.go | 25 +++++++++++++++++ 4 files changed, 59 insertions(+), 15 deletions(-) create mode 100644 backend/services/verification/row_lock.go create mode 100644 backend/services/verification/row_lock_test.go diff --git a/backend/services/verification/event_store.go b/backend/services/verification/event_store.go index 1b406eb5..3d6ce8a9 100644 --- a/backend/services/verification/event_store.go +++ b/backend/services/verification/event_store.go @@ -186,6 +186,7 @@ func (s *EventStore) TryAcquireCleanupLease(ctx context.Context, now time.Time) now = s.now() } session := s.engine.NewSession().Context(ctx) + defer session.Close() if err := session.Begin(); err != nil { return cleanupLease{}, false, fmt.Errorf("begin cleanup lease transaction: %w", err) } @@ -196,12 +197,9 @@ func (s *EventStore) TryAcquireCleanupLease(ctx context.Context, now time.Time) if err := insertCleanupState(session, now); err != nil { return rollback(fmt.Errorf("initialize cleanup state: %w", err)) } - query := session.Where("state_key = ?", cleanupStateKey) - if s.engine.DriverName() != "sqlite3" && s.engine.DriverName() != "sqlite" { - query = query.ForUpdate() - } var state models.VerificationCleanupState - has, err := query.Get(&state) + has, err := getVerificationRowForUpdate(session, &state, + "SELECT * FROM verification_cleanup_state WHERE state_key = ?", cleanupStateKey) if err != nil { return rollback(fmt.Errorf("load cleanup state: %w", err)) } diff --git a/backend/services/verification/rate_limit.go b/backend/services/verification/rate_limit.go index 571be7eb..615fd6b0 100644 --- a/backend/services/verification/rate_limit.go +++ b/backend/services/verification/rate_limit.go @@ -102,11 +102,8 @@ func (r *PersistentRateLimiter) allowWithRelease(ctx context.Context, request Ra return rollback(fmt.Errorf("initialize verification rate limit: %w", err)) } var row models.VerificationRateLimit - query := session.Where("dimension_key = ?", key) - if r.engine.DriverName() != "sqlite3" && r.engine.DriverName() != "sqlite" { - query = query.ForUpdate() - } - has, err := query.Get(&row) + has, err := getVerificationRowForUpdate(session, &row, + "SELECT * FROM verification_rate_limit WHERE dimension_key = ?", key) if err != nil { return rollback(fmt.Errorf("load verification rate limit: %w", err)) } @@ -159,11 +156,8 @@ func (r *PersistentRateLimiter) releaseCooldown(rows []models.VerificationRateLi for _, row := range rows { // 同一维度已有新发送时，不清理新请求的冷却，也不重置小时和每日计数。 var current models.VerificationRateLimit - query := session.ID(row.ID) - if r.engine.DriverName() != "sqlite3" && r.engine.DriverName() != "sqlite" { - query = query.ForUpdate() - } - has, err := query.Get(¤t) + has, err := getVerificationRowForUpdate(session, ¤t, + "SELECT * FROM verification_rate_limit WHERE id = ?", row.ID) if err != nil { return err } diff --git a/backend/services/verification/row_lock.go b/backend/services/verification/row_lock.go new file mode 100644 index 00000000..4b61ce0f --- /dev/null +++ b/backend/services/verification/row_lock.go @@ -0,0 +1,27 @@ +package verification + +import ( + "fmt" + + "xorm.io/xorm" +) + +// 使用参数化查询保留事务行锁，避免当前 Xorm 仅支持 MySQL 的行锁构造器。 +func getVerificationRowForUpdate(session *xorm.Session, bean any, query string, args ...any) (bool, error) { + statement, err := verificationLockQuery(session.Engine().DriverName(), query) + if err != nil { + return false, err + } + return session.SQL(statement, args...).Get(bean) +} + +func verificationLockQuery(driver, query string) (string, error) { + switch driver { + case "postgres", "mysql": + return query + " FOR UPDATE", nil + case "sqlite", "sqlite3": + return query, nil + default: + return "", fmt.Errorf("unsupported verification lock driver: %s", driver) + } +} diff --git a/backend/services/verification/row_lock_test.go b/backend/services/verification/row_lock_test.go new file mode 100644 index 00000000..4a9891df --- /dev/null +++ b/backend/services/verification/row_lock_test.go @@ -0,0 +1,25 @@ +package verification + +import "testing" + +func TestVerificationLockQuery(t *testing.T) { + const query = "SELECT * FROM verification_rate_limit WHERE dimension_key = ?" + for _, driver := range []string{"postgres", "mysql", "sqlite", "sqlite3"} { + t.Run(driver, func(t *testing.T) { + got, err := verificationLockQuery(driver, query) + if err != nil { + t.Fatal(err) + } + want := query + if driver == "postgres" || driver == "mysql" { + want += " FOR UPDATE" + } + if got != want { + t.Fatalf("query = %q, want %q", got, want) + } + }) + } + if got, err := verificationLockQuery("unknown", query); err == nil || got != "" { + t.Fatalf("unsupported driver did not fail closed: %q, %v", got, err) + } +}