hyapp-server/server/admin/internal/job/runner.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 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
	maxAttempts int
	store       Store
	workerID    string

	closeOnce sync.Once
	done      chan struct{}
	stop      chan struct{}
}

func NewRunner(store Store, locker Locker, cfg config.Config) *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"
	}
	return &Runner{
		artifactDir: artifactDir,
		interval:    interval,
		leaseTTL:    leaseTTL,
		locker:      locker,
		maxAttempts: maxAttempts,
		store:       store,
		workerID:    workerID,
		done:        make(chan struct{}),
		stop:        make(chan struct{}),
	}
}

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:
		return fmt.Errorf("unsupported job type: %s", job.Type)
	}
}

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 "未开启"
}