274 lines
7.3 KiB
Go
274 lines
7.3 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"
|
|
)
|
|
|
|
// 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 := s.cfg.TaskCenter.Archive.BatchSize
|
|
if batchSize <= 0 {
|
|
batchSize = 10000
|
|
}
|
|
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)
|
|
if err := tx.Clauses(clause.Locking{Strength: "UPDATE", Options: "SKIP LOCKED"}).
|
|
Where("status IN ? OR (status = ? AND update_time < ?)",
|
|
[]string{ArchiveStatusPending, ArchiveStatusFailed},
|
|
ArchiveStatusProcessing,
|
|
staleProcessingBefore,
|
|
).
|
|
Order("create_time asc, id asc").
|
|
Limit(batchSize).
|
|
Find(&rows).Error; err != nil {
|
|
return err
|
|
}
|
|
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 (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
|
|
}
|