329 lines
11 KiB
Go
329 lines
11 KiB
Go
package region
|
||
|
||
import (
|
||
"context"
|
||
"database/sql"
|
||
"strings"
|
||
"time"
|
||
|
||
"hyapp/pkg/appcode"
|
||
"hyapp/pkg/xerr"
|
||
userdomain "hyapp/services/user-service/internal/domain/user"
|
||
"hyapp/services/user-service/internal/storage/mysql/shared"
|
||
)
|
||
|
||
func (r *Repository) ProcessNextRegionRebuildTask(ctx context.Context, workerID string, nowMs int64, lockTTL time.Duration, batchSize int) (userdomain.RegionRebuildProcessResult, error) {
|
||
if lockTTL <= 0 {
|
||
lockTTL = 30 * time.Second
|
||
}
|
||
if batchSize <= 0 {
|
||
batchSize = 500
|
||
}
|
||
|
||
// claim 必须独立提交：running/locked_until_ms/attempt_count 要对其他 worker 和运维可见。
|
||
task, ok, err := claimRegionRebuildTask(ctx, r.db, appcode.FromContext(ctx), workerID, nowMs, lockTTL)
|
||
if err != nil {
|
||
return userdomain.RegionRebuildProcessResult{}, err
|
||
}
|
||
if !ok {
|
||
return userdomain.RegionRebuildProcessResult{NoTask: true}, nil
|
||
}
|
||
|
||
return r.processClaimedRegionRebuildTask(ctx, task.TaskID, workerID, nowMs, batchSize)
|
||
}
|
||
|
||
func (r *Repository) processClaimedRegionRebuildTask(ctx context.Context, taskID int64, workerID string, nowMs int64, batchSize int) (userdomain.RegionRebuildProcessResult, error) {
|
||
tx, err := r.db.BeginTx(ctx, nil)
|
||
if err != nil {
|
||
return userdomain.RegionRebuildProcessResult{}, err
|
||
}
|
||
defer tx.Rollback()
|
||
|
||
task, ok, err := lockClaimedRegionRebuildTask(ctx, tx, taskID, workerID)
|
||
if err != nil {
|
||
return userdomain.RegionRebuildProcessResult{}, err
|
||
}
|
||
if !ok {
|
||
// task 在 claim 提交后被接管或人工改动时，本轮不应继续写 users。
|
||
return userdomain.RegionRebuildProcessResult{TaskID: taskID}, xerr.New(xerr.Conflict, "region rebuild task lock is lost")
|
||
}
|
||
result := userdomain.RegionRebuildProcessResult{
|
||
AppCode: task.AppCode,
|
||
TaskID: task.TaskID,
|
||
CountryCode: task.CountryCode,
|
||
TargetRegionID: task.TargetRegionID,
|
||
Status: userdomain.RegionRebuildTaskStatusRunning,
|
||
CursorUserID: task.CursorUserID,
|
||
}
|
||
|
||
stale, err := newerRegionRebuildTaskExists(ctx, tx, task.AppCode, task.CountryCode, task.MappingRevision, task.TaskID)
|
||
if err != nil {
|
||
return result, err
|
||
}
|
||
if stale {
|
||
// 旧任务不能覆盖新映射；直接标记 skipped，等待较新任务刷新历史用户。
|
||
if err := finishRegionRebuildTask(ctx, tx, task.TaskID, workerID, userdomain.RegionRebuildTaskStatusSkipped, task.CursorUserID, "stale mapping revision", nowMs); err != nil {
|
||
return result, err
|
||
}
|
||
if err := tx.Commit(); err != nil {
|
||
return result, err
|
||
}
|
||
result.Status = userdomain.RegionRebuildTaskStatusSkipped
|
||
return result, nil
|
||
}
|
||
|
||
userIDs, err := selectRegionRebuildUserIDs(ctx, tx, task.AppCode, task.CountryCode, task.CursorUserID, batchSize)
|
||
if err != nil {
|
||
return result, err
|
||
}
|
||
if len(userIDs) == 0 {
|
||
if err := finishRegionRebuildTask(ctx, tx, task.TaskID, workerID, userdomain.RegionRebuildTaskStatusCompleted, task.CursorUserID, "", nowMs); err != nil {
|
||
return result, err
|
||
}
|
||
if err := tx.Commit(); err != nil {
|
||
return result, err
|
||
}
|
||
result.Status = userdomain.RegionRebuildTaskStatusCompleted
|
||
return result, nil
|
||
}
|
||
|
||
if err := updateUsersRegionByIDs(ctx, tx, task.AppCode, userIDs, task.TargetRegionID, nowMs); err != nil {
|
||
return result, err
|
||
}
|
||
newCursor := userIDs[len(userIDs)-1]
|
||
status := userdomain.RegionRebuildTaskStatusPending
|
||
if len(userIDs) < batchSize {
|
||
status = userdomain.RegionRebuildTaskStatusCompleted
|
||
}
|
||
if err := finishRegionRebuildTask(ctx, tx, task.TaskID, workerID, status, newCursor, "", nowMs); err != nil {
|
||
return result, err
|
||
}
|
||
if err := tx.Commit(); err != nil {
|
||
return result, err
|
||
}
|
||
|
||
result.Status = status
|
||
result.ProcessedUsers = len(userIDs)
|
||
result.CursorUserID = newCursor
|
||
return result, nil
|
||
}
|
||
|
||
func insertRebuildTask(ctx context.Context, tx *sql.Tx, appCode string, regionID int64, countryCode string, targetRegionID int64, revision int64) error {
|
||
_, err := tx.ExecContext(ctx, `
|
||
INSERT INTO user_region_rebuild_tasks (app_code, region_id, country_code, target_region_id, mapping_revision, status, cursor_user_id, created_at_ms, updated_at_ms)
|
||
VALUES (?, ?, ?, ?, ?, ?, 0, ?, ?)
|
||
`, appcode.Normalize(appCode), shared.NullableRegionID(regionID), countryCode, shared.NullableRegionID(targetRegionID), revision, userdomain.RegionRebuildTaskStatusPending, revision, revision)
|
||
return err
|
||
}
|
||
|
||
func claimRegionRebuildTask(ctx context.Context, db *sql.DB, appCode string, workerID string, nowMs int64, lockTTL time.Duration) (userdomain.UserRegionRebuildTask, bool, error) {
|
||
tx, err := db.BeginTx(ctx, nil)
|
||
if err != nil {
|
||
return userdomain.UserRegionRebuildTask{}, false, err
|
||
}
|
||
defer tx.Rollback()
|
||
|
||
task, ok, err := selectClaimableRegionRebuildTask(ctx, tx, appCode, nowMs)
|
||
if err != nil {
|
||
return userdomain.UserRegionRebuildTask{}, false, err
|
||
}
|
||
if !ok {
|
||
return userdomain.UserRegionRebuildTask{}, false, nil
|
||
}
|
||
|
||
lockedUntilMs := nowMs + lockTTL.Milliseconds()
|
||
result, err := tx.ExecContext(ctx, `
|
||
UPDATE user_region_rebuild_tasks
|
||
SET status = ?, locked_by = ?, locked_until_ms = ?, attempt_count = attempt_count + 1, error_message = NULL, updated_at_ms = ?
|
||
WHERE task_id = ?
|
||
`, userdomain.RegionRebuildTaskStatusRunning, workerID, lockedUntilMs, nowMs, task.TaskID)
|
||
if err != nil {
|
||
return userdomain.UserRegionRebuildTask{}, false, err
|
||
}
|
||
if affected, err := result.RowsAffected(); err != nil {
|
||
return userdomain.UserRegionRebuildTask{}, false, err
|
||
} else if affected != 1 {
|
||
return userdomain.UserRegionRebuildTask{}, false, xerr.New(xerr.Conflict, "region rebuild task claim conflict")
|
||
}
|
||
if err := tx.Commit(); err != nil {
|
||
return userdomain.UserRegionRebuildTask{}, false, err
|
||
}
|
||
task.Status = userdomain.RegionRebuildTaskStatusRunning
|
||
task.LockedBy = workerID
|
||
task.LockedUntilMs = lockedUntilMs
|
||
task.AttemptCount++
|
||
task.UpdatedAtMs = nowMs
|
||
|
||
return task, true, nil
|
||
}
|
||
|
||
func selectClaimableRegionRebuildTask(ctx context.Context, tx *sql.Tx, appCode string, nowMs int64) (userdomain.UserRegionRebuildTask, bool, error) {
|
||
row := tx.QueryRowContext(ctx, `
|
||
SELECT app_code, task_id, region_id, country_code, target_region_id, mapping_revision, status, cursor_user_id,
|
||
locked_by, locked_until_ms, attempt_count, error_message, created_at_ms, updated_at_ms
|
||
FROM user_region_rebuild_tasks
|
||
WHERE app_code = ?
|
||
AND (status = ?
|
||
OR (status = ? AND COALESCE(locked_until_ms, 0) <= ?))
|
||
ORDER BY updated_at_ms ASC, task_id ASC
|
||
LIMIT 1
|
||
FOR UPDATE SKIP LOCKED
|
||
`, appcode.Normalize(appCode), userdomain.RegionRebuildTaskStatusPending, userdomain.RegionRebuildTaskStatusRunning, nowMs)
|
||
task, err := scanRegionRebuildTask(row)
|
||
if err == sql.ErrNoRows {
|
||
return userdomain.UserRegionRebuildTask{}, false, nil
|
||
}
|
||
if err != nil {
|
||
return userdomain.UserRegionRebuildTask{}, false, err
|
||
}
|
||
|
||
return task, true, nil
|
||
}
|
||
|
||
func lockClaimedRegionRebuildTask(ctx context.Context, tx *sql.Tx, taskID int64, workerID string) (userdomain.UserRegionRebuildTask, bool, error) {
|
||
row := tx.QueryRowContext(ctx, `
|
||
SELECT app_code, task_id, region_id, country_code, target_region_id, mapping_revision, status, cursor_user_id,
|
||
locked_by, locked_until_ms, attempt_count, error_message, created_at_ms, updated_at_ms
|
||
FROM user_region_rebuild_tasks
|
||
WHERE task_id = ? AND status = ? AND locked_by = ?
|
||
FOR UPDATE
|
||
`, taskID, userdomain.RegionRebuildTaskStatusRunning, workerID)
|
||
task, err := scanRegionRebuildTask(row)
|
||
if err == sql.ErrNoRows {
|
||
return userdomain.UserRegionRebuildTask{}, false, nil
|
||
}
|
||
if err != nil {
|
||
return userdomain.UserRegionRebuildTask{}, false, err
|
||
}
|
||
|
||
return task, true, nil
|
||
}
|
||
|
||
func scanRegionRebuildTask(scanner interface{ Scan(dest ...any) error }) (userdomain.UserRegionRebuildTask, error) {
|
||
var task userdomain.UserRegionRebuildTask
|
||
var regionID sql.Null[int64]
|
||
var targetRegionID sql.Null[int64]
|
||
var lockedBy sql.Null[string]
|
||
var lockedUntil sql.Null[int64]
|
||
var errorMessage sql.Null[string]
|
||
err := scanner.Scan(
|
||
&task.AppCode,
|
||
&task.TaskID,
|
||
®ionID,
|
||
&task.CountryCode,
|
||
&targetRegionID,
|
||
&task.MappingRevision,
|
||
&task.Status,
|
||
&task.CursorUserID,
|
||
&lockedBy,
|
||
&lockedUntil,
|
||
&task.AttemptCount,
|
||
&errorMessage,
|
||
&task.CreatedAtMs,
|
||
&task.UpdatedAtMs,
|
||
)
|
||
if err != nil {
|
||
return userdomain.UserRegionRebuildTask{}, err
|
||
}
|
||
if regionID.Valid {
|
||
task.RegionID = regionID.V
|
||
}
|
||
if targetRegionID.Valid {
|
||
task.TargetRegionID = targetRegionID.V
|
||
}
|
||
if lockedBy.Valid {
|
||
task.LockedBy = lockedBy.V
|
||
}
|
||
if lockedUntil.Valid {
|
||
task.LockedUntilMs = lockedUntil.V
|
||
}
|
||
if errorMessage.Valid {
|
||
task.ErrorMessage = errorMessage.V
|
||
}
|
||
|
||
return task, nil
|
||
}
|
||
|
||
func newerRegionRebuildTaskExists(ctx context.Context, q queryer, appCode string, countryCode string, revision int64, taskID int64) (bool, error) {
|
||
var one int
|
||
err := q.QueryRowContext(ctx, `
|
||
SELECT 1
|
||
FROM user_region_rebuild_tasks
|
||
WHERE app_code = ?
|
||
AND country_code = ?
|
||
AND (mapping_revision > ? OR (mapping_revision = ? AND task_id > ?))
|
||
LIMIT 1
|
||
`, appcode.Normalize(appCode), countryCode, revision, revision, taskID).Scan(&one)
|
||
if err == sql.ErrNoRows {
|
||
return false, nil
|
||
}
|
||
if err != nil {
|
||
return false, err
|
||
}
|
||
|
||
return true, nil
|
||
}
|
||
|
||
func selectRegionRebuildUserIDs(ctx context.Context, q queryer, appCode string, countryCode string, cursorUserID int64, batchSize int) ([]int64, error) {
|
||
rows, err := queryRows(ctx, q, `
|
||
SELECT user_id
|
||
FROM users
|
||
WHERE app_code = ? AND country = ? AND user_id > ?
|
||
ORDER BY user_id ASC
|
||
LIMIT ?
|
||
`, appcode.Normalize(appCode), countryCode, cursorUserID, batchSize)
|
||
if err != nil {
|
||
return nil, err
|
||
}
|
||
defer rows.Close()
|
||
|
||
userIDs := make([]int64, 0, batchSize)
|
||
for rows.Next() {
|
||
var userID int64
|
||
if err := rows.Scan(&userID); err != nil {
|
||
return nil, err
|
||
}
|
||
userIDs = append(userIDs, userID)
|
||
}
|
||
|
||
return userIDs, rows.Err()
|
||
}
|
||
|
||
func updateUsersRegionByIDs(ctx context.Context, tx *sql.Tx, appCode string, userIDs []int64, targetRegionID int64, nowMs int64) error {
|
||
if len(userIDs) == 0 {
|
||
return nil
|
||
}
|
||
placeholders := make([]string, 0, len(userIDs))
|
||
args := make([]any, 0, len(userIDs)+3)
|
||
args = append(args, shared.NullableRegionID(targetRegionID), nowMs, appcode.Normalize(appCode))
|
||
for _, userID := range userIDs {
|
||
placeholders = append(placeholders, "?")
|
||
args = append(args, userID)
|
||
}
|
||
_, err := tx.ExecContext(ctx, `
|
||
UPDATE users
|
||
SET region_id = ?, updated_at_ms = ?
|
||
WHERE app_code = ? AND user_id IN (`+strings.Join(placeholders, ",")+`)
|
||
`, args...)
|
||
return err
|
||
}
|
||
|
||
func finishRegionRebuildTask(ctx context.Context, tx *sql.Tx, taskID int64, workerID string, status string, cursorUserID int64, errorMessage string, nowMs int64) error {
|
||
result, err := tx.ExecContext(ctx, `
|
||
UPDATE user_region_rebuild_tasks
|
||
SET status = ?, cursor_user_id = ?, locked_by = NULL, locked_until_ms = NULL, error_message = ?, updated_at_ms = ?
|
||
WHERE task_id = ? AND status = ? AND locked_by = ?
|
||
`, status, cursorUserID, shared.NullableString(errorMessage), nowMs, taskID, userdomain.RegionRebuildTaskStatusRunning, workerID)
|
||
if err != nil {
|
||
return err
|
||
}
|
||
if affected, err := result.RowsAffected(); err != nil {
|
||
return err
|
||
} else if affected != 1 {
|
||
// finish 和 users 更新在同一个事务里；丢失 claim 时让事务回滚，避免旧 worker 覆盖新进度。
|
||
return xerr.New(xerr.Conflict, "region rebuild task lock is lost")
|
||
}
|
||
|
||
return nil
|
||
}
|