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) type Option func(*Runner) type Store interface { LeaseNextJob(workerID string, lease time.Duration, maxAttempts int) (*model.AdminJob, error) CompleteJobWithArtifact(id uint, resultJSON string, artifactPath string) error FailJobAttempt(id uint, message string) error ReleaseJobLease(id uint) 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) } type Runner struct { artifactDir string interval time.Duration leaseTTL time.Duration locker Locker handlers map[string]HandlerFunc 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 } } 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), 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) return } if !locked { _ = r.store.ReleaseJobLease(job.ID) return } err = r.handle(ctx, job) if unlock != nil { _ = unlock(ctx) } if err != nil { _ = r.store.FailJobAttempt(job.ID, err.Error()) } } } func (r *Runner) handle(ctx context.Context, job *model.AdminJob) error { switch job.Type { case userExportJobType: return r.exportAdminUsers(ctx, job) default: 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, 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, 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 "未开启" }