344 lines
9.4 KiB
Go
344 lines
9.4 KiB
Go
package taskcenter
|
||
|
||
import (
|
||
"bytes"
|
||
"compress/gzip"
|
||
"context"
|
||
"fmt"
|
||
"log"
|
||
"net/http"
|
||
"net/url"
|
||
"path"
|
||
"sort"
|
||
"strings"
|
||
"time"
|
||
|
||
"chatapp3-golang/internal/model"
|
||
|
||
"github.com/tencentyun/cos-go-sdk-v5"
|
||
"gorm.io/gorm"
|
||
"gorm.io/gorm/clause"
|
||
)
|
||
|
||
const (
|
||
defaultTaskCenterArchiveBatchSize = 500
|
||
maxTaskCenterArchiveBatchSize = 1000
|
||
)
|
||
|
||
// StartArchiveWorker 启动任务中心事件 COS 归档 worker。
|
||
func (s *Service) StartArchiveWorker(ctx context.Context) error {
|
||
if !s.cfg.TaskCenter.Archive.Enabled {
|
||
log.Printf("task center archive worker disabled")
|
||
return nil
|
||
}
|
||
if strings.TrimSpace(s.cfg.TaskCenter.Archive.Bucket) == "" ||
|
||
strings.TrimSpace(s.cfg.TaskCenter.Archive.Region) == "" ||
|
||
strings.TrimSpace(s.cfg.TaskCenter.Archive.SecretID) == "" ||
|
||
strings.TrimSpace(s.cfg.TaskCenter.Archive.SecretKey) == "" {
|
||
log.Printf("task center archive worker skipped: cos bucket, region or credentials missing")
|
||
return nil
|
||
}
|
||
|
||
workers := s.cfg.TaskCenter.Archive.Workers
|
||
if workers <= 0 {
|
||
workers = 1
|
||
}
|
||
for index := 0; index < workers; index++ {
|
||
go s.runArchiveWorker(ctx, index)
|
||
}
|
||
return nil
|
||
}
|
||
|
||
func (s *Service) runArchiveWorker(ctx context.Context, workerIndex int) {
|
||
interval := time.Duration(s.cfg.TaskCenter.Archive.FlushSeconds) * time.Second
|
||
if interval <= 0 {
|
||
interval = time.Minute
|
||
}
|
||
ticker := time.NewTicker(interval)
|
||
defer ticker.Stop()
|
||
|
||
for {
|
||
if ctx.Err() != nil {
|
||
return
|
||
}
|
||
if err := s.archivePendingEvents(ctx); err != nil {
|
||
log.Printf("task center archive worker %d failed: %v", workerIndex, err)
|
||
}
|
||
select {
|
||
case <-ctx.Done():
|
||
return
|
||
case <-ticker.C:
|
||
}
|
||
}
|
||
}
|
||
|
||
func (s *Service) archivePendingEvents(ctx context.Context) error {
|
||
batchSize := normalizeTaskCenterArchiveBatchSize(s.cfg.TaskCenter.Archive.BatchSize)
|
||
rows, err := s.claimArchiveRows(ctx, batchSize)
|
||
if err != nil {
|
||
return err
|
||
}
|
||
if len(rows) == 0 {
|
||
return nil
|
||
}
|
||
|
||
client, err := s.newCOSClient()
|
||
if err != nil {
|
||
return err
|
||
}
|
||
for _, group := range groupArchiveRows(rows) {
|
||
if err := s.uploadArchiveGroup(ctx, client, group); err != nil {
|
||
_ = s.markArchiveGroupFailed(ctx, group, err)
|
||
return err
|
||
}
|
||
}
|
||
return nil
|
||
}
|
||
|
||
func (s *Service) claimArchiveRows(ctx context.Context, batchSize int) ([]model.TaskCenterEventArchiveOutbox, error) {
|
||
var rows []model.TaskCenterEventArchiveOutbox
|
||
err := s.db.WithContext(ctx).Transaction(func(tx *gorm.DB) error {
|
||
staleProcessingBefore := time.Now().Add(-10 * time.Minute)
|
||
staleLimit, _ := taskCenterArchiveClaimLimits(batchSize)
|
||
if staleLimit > 0 {
|
||
var staleRows []model.TaskCenterEventArchiveOutbox
|
||
// PROCESSING 超时行保留固定接管额度，避免 PENDING/FAILED 持续堆积时旧锁永久饥饿。
|
||
if err := tx.Clauses(clause.Locking{Strength: "UPDATE", Options: "SKIP LOCKED"}).
|
||
Where("status = ? AND update_time < ?", ArchiveStatusProcessing, staleProcessingBefore).
|
||
Order("update_time asc, id asc").
|
||
Limit(staleLimit).
|
||
Find(&staleRows).Error; err != nil {
|
||
return err
|
||
}
|
||
rows = append(rows, staleRows...)
|
||
}
|
||
if remaining := batchSize - len(rows); remaining > 0 {
|
||
pendingLimit := taskCenterArchivePendingClaimLimit(remaining)
|
||
for _, retryable := range []struct {
|
||
status string
|
||
limit int
|
||
}{
|
||
{status: ArchiveStatusPending, limit: pendingLimit},
|
||
{status: ArchiveStatusFailed, limit: remaining},
|
||
} {
|
||
if retryable.limit <= 0 || remaining <= 0 {
|
||
continue
|
||
}
|
||
if retryable.limit > remaining {
|
||
retryable.limit = remaining
|
||
}
|
||
var retryableRows []model.TaskCenterEventArchiveOutbox
|
||
// PENDING 和 FAILED 分开抢占，避免 IN 跨 status 排序触发 filesort，同时给 FAILED 保留重试额度。
|
||
if err := tx.Clauses(clause.Locking{Strength: "UPDATE", Options: "SKIP LOCKED"}).
|
||
Where("status = ?", retryable.status).
|
||
Order("create_time asc, id asc").
|
||
Limit(retryable.limit).
|
||
Find(&retryableRows).Error; err != nil {
|
||
return err
|
||
}
|
||
rows = append(rows, retryableRows...)
|
||
remaining -= len(retryableRows)
|
||
}
|
||
}
|
||
if len(rows) == 0 {
|
||
return nil
|
||
}
|
||
return tx.Model(&model.TaskCenterEventArchiveOutbox{}).
|
||
Where("id IN ?", archiveRowIDs(rows)).
|
||
Updates(map[string]any{
|
||
"status": ArchiveStatusProcessing,
|
||
"update_time": time.Now(),
|
||
}).Error
|
||
})
|
||
if err != nil {
|
||
return nil, err
|
||
}
|
||
return rows, nil
|
||
}
|
||
|
||
func normalizeTaskCenterArchiveBatchSize(batchSize int) int {
|
||
if batchSize <= 0 {
|
||
return defaultTaskCenterArchiveBatchSize
|
||
}
|
||
if batchSize > maxTaskCenterArchiveBatchSize {
|
||
return maxTaskCenterArchiveBatchSize
|
||
}
|
||
return batchSize
|
||
}
|
||
|
||
func taskCenterArchiveClaimLimits(batchSize int) (int, int) {
|
||
if batchSize <= 1 {
|
||
return batchSize, 0
|
||
}
|
||
staleLimit := batchSize / 10
|
||
if staleLimit < 1 {
|
||
staleLimit = 1
|
||
}
|
||
if staleLimit > 100 {
|
||
staleLimit = 100
|
||
}
|
||
return staleLimit, batchSize - staleLimit
|
||
}
|
||
|
||
func taskCenterArchivePendingClaimLimit(retryableLimit int) int {
|
||
if retryableLimit <= 1 {
|
||
return retryableLimit
|
||
}
|
||
failedReserve := retryableLimit / 10
|
||
if failedReserve < 1 {
|
||
failedReserve = 1
|
||
}
|
||
if failedReserve > 100 {
|
||
failedReserve = 100
|
||
}
|
||
return retryableLimit - failedReserve
|
||
}
|
||
|
||
func (s *Service) newCOSClient() (*cos.Client, error) {
|
||
bucketURL, err := url.Parse(fmt.Sprintf("https://%s.cos.%s.myqcloud.com",
|
||
strings.TrimSpace(s.cfg.TaskCenter.Archive.Bucket),
|
||
strings.TrimSpace(s.cfg.TaskCenter.Archive.Region),
|
||
))
|
||
if err != nil {
|
||
return nil, err
|
||
}
|
||
baseURL := &cos.BaseURL{BucketURL: bucketURL}
|
||
return cos.NewClient(baseURL, &http.Client{
|
||
Transport: &cos.AuthorizationTransport{
|
||
SecretID: strings.TrimSpace(s.cfg.TaskCenter.Archive.SecretID),
|
||
SecretKey: strings.TrimSpace(s.cfg.TaskCenter.Archive.SecretKey),
|
||
},
|
||
}), nil
|
||
}
|
||
|
||
func groupArchiveRows(rows []model.TaskCenterEventArchiveOutbox) [][]model.TaskCenterEventArchiveOutbox {
|
||
groupsByKey := make(map[string][]model.TaskCenterEventArchiveOutbox)
|
||
keys := make([]string, 0)
|
||
for _, row := range rows {
|
||
groupKey := row.SysOrigin + "\x00" + row.EventType
|
||
if _, exists := groupsByKey[groupKey]; !exists {
|
||
keys = append(keys, groupKey)
|
||
}
|
||
groupsByKey[groupKey] = append(groupsByKey[groupKey], row)
|
||
}
|
||
sort.Strings(keys)
|
||
groups := make([][]model.TaskCenterEventArchiveOutbox, 0, len(keys))
|
||
for _, key := range keys {
|
||
groups = append(groups, groupsByKey[key])
|
||
}
|
||
return groups
|
||
}
|
||
|
||
func (s *Service) uploadArchiveGroup(ctx context.Context, client *cos.Client, rows []model.TaskCenterEventArchiveOutbox) error {
|
||
if len(rows) == 0 {
|
||
return nil
|
||
}
|
||
cosKey := s.buildArchiveCOSKey(rows)
|
||
body, err := gzipArchiveRows(rows)
|
||
if err != nil {
|
||
return err
|
||
}
|
||
_, err = client.Object.Put(ctx, cosKey, bytes.NewReader(body), &cos.ObjectPutOptions{
|
||
ObjectPutHeaderOptions: &cos.ObjectPutHeaderOptions{
|
||
ContentType: "application/gzip",
|
||
ContentEncoding: "gzip",
|
||
},
|
||
})
|
||
if err != nil {
|
||
return err
|
||
}
|
||
return s.markArchiveGroupUploaded(ctx, rows, cosKey)
|
||
}
|
||
|
||
func gzipArchiveRows(rows []model.TaskCenterEventArchiveOutbox) ([]byte, error) {
|
||
var buffer bytes.Buffer
|
||
writer := gzip.NewWriter(&buffer)
|
||
for _, row := range rows {
|
||
line := strings.TrimSpace(row.RawJSON)
|
||
if line == "" {
|
||
line = "{}"
|
||
}
|
||
if _, err := writer.Write([]byte(line)); err != nil {
|
||
_ = writer.Close()
|
||
return nil, err
|
||
}
|
||
if _, err := writer.Write([]byte("\n")); err != nil {
|
||
_ = writer.Close()
|
||
return nil, err
|
||
}
|
||
}
|
||
if err := writer.Close(); err != nil {
|
||
return nil, err
|
||
}
|
||
return buffer.Bytes(), nil
|
||
}
|
||
|
||
func (s *Service) buildArchiveCOSKey(rows []model.TaskCenterEventArchiveOutbox) string {
|
||
first := rows[0]
|
||
occurredAt := first.OccurredAt
|
||
if occurredAt.IsZero() {
|
||
occurredAt = time.Now()
|
||
}
|
||
occurredAt = occurredAt.UTC()
|
||
prefix := strings.Trim(strings.TrimSpace(s.cfg.TaskCenter.Archive.Prefix), "/")
|
||
if prefix == "" {
|
||
prefix = "task-center-events/v1"
|
||
}
|
||
fileName := fmt.Sprintf("part-%s-%d-%d.jsonl.gz",
|
||
time.Now().UTC().Format("20060102T150405.000000000Z"),
|
||
first.ID,
|
||
len(rows),
|
||
)
|
||
return path.Join(
|
||
prefix,
|
||
"dt="+occurredAt.Format("2006-01-02"),
|
||
"hour="+occurredAt.Format("15"),
|
||
"sys_origin="+sanitizeArchivePartition(first.SysOrigin),
|
||
"event_type="+sanitizeArchivePartition(first.EventType),
|
||
fileName,
|
||
)
|
||
}
|
||
|
||
func sanitizeArchivePartition(value string) string {
|
||
value = strings.TrimSpace(value)
|
||
if value == "" {
|
||
return "UNKNOWN"
|
||
}
|
||
replacer := strings.NewReplacer("/", "_", "\\", "_", " ", "_", "=", "_")
|
||
return replacer.Replace(value)
|
||
}
|
||
|
||
func (s *Service) markArchiveGroupUploaded(ctx context.Context, rows []model.TaskCenterEventArchiveOutbox, cosKey string) error {
|
||
ids := archiveRowIDs(rows)
|
||
now := time.Now()
|
||
return s.db.WithContext(ctx).Model(&model.TaskCenterEventArchiveOutbox{}).
|
||
Where("id IN ?", ids).
|
||
Updates(map[string]any{
|
||
"status": ArchiveStatusUploaded,
|
||
"cos_key": cosKey,
|
||
"failure_reason": "",
|
||
"archived_at": now,
|
||
"update_time": now,
|
||
}).Error
|
||
}
|
||
|
||
func (s *Service) markArchiveGroupFailed(ctx context.Context, rows []model.TaskCenterEventArchiveOutbox, archiveErr error) error {
|
||
ids := archiveRowIDs(rows)
|
||
return s.db.WithContext(ctx).Model(&model.TaskCenterEventArchiveOutbox{}).
|
||
Where("id IN ?", ids).
|
||
Updates(map[string]any{
|
||
"status": ArchiveStatusFailed,
|
||
"retry_count": gorm.Expr("retry_count + 1"),
|
||
"failure_reason": truncateTaskCenterFailure(archiveErr.Error()),
|
||
"update_time": time.Now(),
|
||
}).Error
|
||
}
|
||
|
||
func archiveRowIDs(rows []model.TaskCenterEventArchiveOutbox) []int64 {
|
||
ids := make([]int64, 0, len(rows))
|
||
for _, row := range rows {
|
||
ids = append(ids, row.ID)
|
||
}
|
||
return ids
|
||
}
|