350 lines
8.9 KiB
Go
350 lines
8.9 KiB
Go
package job
|
||
|
||
import (
|
||
"bytes"
|
||
"context"
|
||
"encoding/csv"
|
||
"encoding/json"
|
||
"errors"
|
||
"fmt"
|
||
"log/slog"
|
||
"os"
|
||
"path/filepath"
|
||
"strings"
|
||
"sync"
|
||
"time"
|
||
|
||
"hyapp-admin-server/internal/cache"
|
||
"hyapp-admin-server/internal/config"
|
||
"hyapp-admin-server/internal/model"
|
||
"hyapp-admin-server/internal/repository"
|
||
|
||
"gorm.io/gorm"
|
||
)
|
||
|
||
const userExportJobType = "user-export"
|
||
|
||
type HandlerFunc func(ctx context.Context, job *model.AdminJob) (string, error)
|
||
|
||
// ArtifactHandlerFunc 供需要生成可下载文件的业务任务使用；artifactPath 必须位于 runner 配置的导出目录内。
|
||
type ArtifactHandlerFunc func(ctx context.Context, job *model.AdminJob) (resultJSON string, artifactPath string, err error)
|
||
|
||
type Option func(*Runner)
|
||
|
||
type Store interface {
|
||
LeaseNextJob(workerID string, lease time.Duration, maxAttempts int) (*model.AdminJob, error)
|
||
RenewJobLease(id uint, workerID string, lease time.Duration) error
|
||
CompleteJobWithArtifact(id uint, workerID string, resultJSON string, artifactPath string) error
|
||
FailJobAttempt(id uint, workerID string, message string) error
|
||
ReleaseJobLease(id uint, workerID string) error
|
||
ExportUsers(options repository.ListOptions) ([]model.User, error)
|
||
}
|
||
|
||
type Locker interface {
|
||
TryLock(ctx context.Context, key string, owner string, ttl time.Duration) (bool, cache.UnlockFunc, error)
|
||
RefreshLock(ctx context.Context, key string, owner string, ttl time.Duration) (bool, error)
|
||
}
|
||
|
||
type Runner struct {
|
||
artifactDir string
|
||
interval time.Duration
|
||
leaseTTL time.Duration
|
||
locker Locker
|
||
handlers map[string]HandlerFunc
|
||
artifactHandlers map[string]ArtifactHandlerFunc
|
||
maxAttempts int
|
||
store Store
|
||
workerID string
|
||
|
||
closeOnce sync.Once
|
||
done chan struct{}
|
||
stop chan struct{}
|
||
}
|
||
|
||
func WithHandler(jobType string, handler HandlerFunc) Option {
|
||
return func(r *Runner) {
|
||
jobType = strings.TrimSpace(jobType)
|
||
if jobType == "" || handler == nil {
|
||
return
|
||
}
|
||
r.handlers[jobType] = handler
|
||
}
|
||
}
|
||
|
||
// WithArtifactHandler 注册带文件产物的任务，避免业务模块绕过统一 job lease、重试和下载校验。
|
||
func WithArtifactHandler(jobType string, handler ArtifactHandlerFunc) Option {
|
||
return func(r *Runner) {
|
||
jobType = strings.TrimSpace(jobType)
|
||
if jobType == "" || handler == nil {
|
||
return
|
||
}
|
||
r.artifactHandlers[jobType] = handler
|
||
}
|
||
}
|
||
|
||
func NewRunner(store Store, locker Locker, cfg config.Config, opts ...Option) *Runner {
|
||
interval := cfg.Jobs.WorkerInterval
|
||
if interval <= 0 {
|
||
interval = 5 * time.Second
|
||
}
|
||
leaseTTL := cfg.Jobs.LeaseTTL
|
||
if leaseTTL <= 0 {
|
||
leaseTTL = 2 * time.Minute
|
||
}
|
||
maxAttempts := cfg.Jobs.MaxAttempts
|
||
if maxAttempts <= 0 {
|
||
maxAttempts = 3
|
||
}
|
||
artifactDir := strings.TrimSpace(cfg.Jobs.ArtifactDir)
|
||
if artifactDir == "" {
|
||
artifactDir = "storage/exports"
|
||
}
|
||
workerID := strings.TrimSpace(cfg.NodeID)
|
||
if workerID == "" {
|
||
workerID = "admin-job-worker"
|
||
}
|
||
runner := &Runner{
|
||
artifactDir: artifactDir,
|
||
interval: interval,
|
||
leaseTTL: leaseTTL,
|
||
locker: locker,
|
||
handlers: make(map[string]HandlerFunc),
|
||
artifactHandlers: make(map[string]ArtifactHandlerFunc),
|
||
maxAttempts: maxAttempts,
|
||
store: store,
|
||
workerID: workerID,
|
||
done: make(chan struct{}),
|
||
stop: make(chan struct{}),
|
||
}
|
||
for _, opt := range opts {
|
||
if opt != nil {
|
||
opt(runner)
|
||
}
|
||
}
|
||
return runner
|
||
}
|
||
|
||
func (r *Runner) Start() {
|
||
go r.loop()
|
||
}
|
||
|
||
func (r *Runner) Close() {
|
||
r.closeOnce.Do(func() {
|
||
close(r.stop)
|
||
<-r.done
|
||
})
|
||
}
|
||
|
||
func (r *Runner) Status() map[string]any {
|
||
return map[string]any{
|
||
"artifactDir": r.artifactDir,
|
||
"interval": r.interval.String(),
|
||
"leaseTTL": r.leaseTTL.String(),
|
||
"maxAttempts": r.maxAttempts,
|
||
"workerID": r.workerID,
|
||
}
|
||
}
|
||
|
||
func (r *Runner) loop() {
|
||
defer close(r.done)
|
||
ticker := time.NewTicker(r.interval)
|
||
defer ticker.Stop()
|
||
|
||
r.runOnce(context.Background())
|
||
for {
|
||
select {
|
||
case <-ticker.C:
|
||
r.runOnce(context.Background())
|
||
case <-r.stop:
|
||
return
|
||
}
|
||
}
|
||
}
|
||
|
||
func (r *Runner) runOnce(ctx context.Context) {
|
||
for {
|
||
job, err := r.store.LeaseNextJob(r.workerID, r.leaseTTL, r.maxAttempts)
|
||
if err != nil {
|
||
if errors.Is(err, gorm.ErrRecordNotFound) {
|
||
return
|
||
}
|
||
slog.Error("job_find_pending_failed", "worker_id", r.workerID, "error", err.Error())
|
||
return
|
||
}
|
||
locked, unlock, err := r.lock(ctx, job)
|
||
if err != nil {
|
||
slog.Error("job_lock_failed", "worker_id", r.workerID, "job_id", job.ID, "error", err.Error())
|
||
_ = r.store.ReleaseJobLease(job.ID, r.workerID)
|
||
return
|
||
}
|
||
if !locked {
|
||
_ = r.store.ReleaseJobLease(job.ID, r.workerID)
|
||
return
|
||
}
|
||
err = r.handleWithHeartbeat(ctx, job)
|
||
if unlock != nil {
|
||
_ = unlock(ctx)
|
||
}
|
||
if err != nil {
|
||
_ = r.store.FailJobAttempt(job.ID, r.workerID, err.Error())
|
||
}
|
||
}
|
||
}
|
||
|
||
func (r *Runner) handleWithHeartbeat(ctx context.Context, job *model.AdminJob) error {
|
||
jobCtx, cancel := context.WithCancel(ctx)
|
||
defer cancel()
|
||
heartbeatResult := make(chan error, 1)
|
||
go r.heartbeat(jobCtx, cancel, job, heartbeatResult)
|
||
|
||
handleErr := r.handle(jobCtx, job)
|
||
cancel()
|
||
heartbeatErr := <-heartbeatResult
|
||
if handleErr != nil {
|
||
return handleErr
|
||
}
|
||
return heartbeatErr
|
||
}
|
||
|
||
func (r *Runner) heartbeat(ctx context.Context, cancel context.CancelFunc, job *model.AdminJob, result chan<- error) {
|
||
interval := r.leaseTTL / 3
|
||
if interval < 100*time.Millisecond {
|
||
interval = 100 * time.Millisecond
|
||
}
|
||
ticker := time.NewTicker(interval)
|
||
defer ticker.Stop()
|
||
for {
|
||
select {
|
||
case <-ctx.Done():
|
||
result <- nil
|
||
return
|
||
case <-ticker.C:
|
||
if err := r.store.RenewJobLease(job.ID, r.workerID, r.leaseTTL); err != nil {
|
||
cancel()
|
||
result <- fmt.Errorf("renew job lease: %w", err)
|
||
return
|
||
}
|
||
if r.locker != nil {
|
||
refreshed, err := r.locker.RefreshLock(ctx, fmt.Sprintf("jobs:%d", job.ID), r.workerID, r.leaseTTL)
|
||
if err != nil {
|
||
cancel()
|
||
result <- fmt.Errorf("refresh job lock: %w", err)
|
||
return
|
||
}
|
||
if !refreshed {
|
||
cancel()
|
||
result <- errors.New("job distributed lock is no longer owned by worker")
|
||
return
|
||
}
|
||
}
|
||
}
|
||
}
|
||
}
|
||
|
||
func (r *Runner) handle(ctx context.Context, job *model.AdminJob) error {
|
||
switch job.Type {
|
||
case userExportJobType:
|
||
return r.exportAdminUsers(ctx, job)
|
||
default:
|
||
if handler := r.artifactHandlers[strings.TrimSpace(job.Type)]; handler != nil {
|
||
resultJSON, artifactPath, err := handler(ctx, job)
|
||
if err != nil {
|
||
return err
|
||
}
|
||
return r.store.CompleteJobWithArtifact(job.ID, r.workerID, resultJSON, artifactPath)
|
||
}
|
||
handler := r.handlers[strings.TrimSpace(job.Type)]
|
||
if handler == nil {
|
||
return fmt.Errorf("unsupported job type: %s", job.Type)
|
||
}
|
||
resultJSON, err := handler(ctx, job)
|
||
if err != nil {
|
||
return err
|
||
}
|
||
return r.store.CompleteJobWithArtifact(job.ID, r.workerID, resultJSON, "")
|
||
}
|
||
}
|
||
|
||
func (r *Runner) exportAdminUsers(_ context.Context, job *model.AdminJob) error {
|
||
var options repository.ListOptions
|
||
if strings.TrimSpace(job.PayloadJSON) != "" {
|
||
if err := json.Unmarshal([]byte(job.PayloadJSON), &options); err != nil {
|
||
return err
|
||
}
|
||
}
|
||
users, err := r.store.ExportUsers(options)
|
||
if err != nil {
|
||
return err
|
||
}
|
||
artifactPath, err := r.writeUserExport(job, users)
|
||
if err != nil {
|
||
return err
|
||
}
|
||
result, err := json.Marshal(map[string]any{
|
||
"count": len(users),
|
||
"downloadPath": fmt.Sprintf("/api/v1/jobs/%d/artifact", job.ID),
|
||
})
|
||
if err != nil {
|
||
return err
|
||
}
|
||
return r.store.CompleteJobWithArtifact(job.ID, r.workerID, string(result), artifactPath)
|
||
}
|
||
|
||
func (r *Runner) lock(ctx context.Context, job *model.AdminJob) (bool, cache.UnlockFunc, error) {
|
||
if r.locker == nil {
|
||
return true, func(context.Context) error { return nil }, nil
|
||
}
|
||
return r.locker.TryLock(ctx, fmt.Sprintf("jobs:%d", job.ID), r.workerID, r.leaseTTL)
|
||
}
|
||
|
||
func (r *Runner) writeUserExport(job *model.AdminJob, users []model.User) (string, error) {
|
||
if err := os.MkdirAll(r.artifactDir, 0o755); err != nil {
|
||
return "", err
|
||
}
|
||
fileName := fmt.Sprintf("user-export-%d.csv", job.ID)
|
||
artifactPath := filepath.Join(r.artifactDir, fileName)
|
||
|
||
var buf bytes.Buffer
|
||
writer := csv.NewWriter(&buf)
|
||
_ = writer.Write([]string{"ID", "账号", "姓名", "角色", "团队", "状态", "MFA", "最后登录"})
|
||
for _, user := range users {
|
||
lastLogin := ""
|
||
if user.LastLoginAtMS != nil {
|
||
lastLogin = time.UnixMilli(*user.LastLoginAtMS).UTC().Format("2006-01-02 15:04:05")
|
||
}
|
||
_ = writer.Write([]string{
|
||
fmt.Sprintf("%d", user.ID),
|
||
user.Username,
|
||
user.Name,
|
||
strings.Join(roleNames(user.Roles), "、"),
|
||
user.Team,
|
||
user.Status,
|
||
boolText(user.MFAEnabled),
|
||
lastLogin,
|
||
})
|
||
}
|
||
writer.Flush()
|
||
if err := writer.Error(); err != nil {
|
||
return "", err
|
||
}
|
||
if err := os.WriteFile(artifactPath, buf.Bytes(), 0o644); err != nil {
|
||
return "", err
|
||
}
|
||
return fileName, nil
|
||
}
|
||
|
||
func roleNames(roles []model.Role) []string {
|
||
names := make([]string, 0, len(roles))
|
||
for _, role := range roles {
|
||
names = append(names, role.Name)
|
||
}
|
||
return names
|
||
}
|
||
|
||
func boolText(value bool) string {
|
||
if value {
|
||
return "已开启"
|
||
}
|
||
return "未开启"
|
||
}
|