hyapp-server/services/cron-service/internal/scheduler/scheduler.go

// Package scheduler runs configured cron tasks and delegates business work to handlers.
package scheduler

import (
	"context"
	"log/slog"
	"sort"
	"strings"
	"sync"
	"time"

	"hyapp/pkg/idgen"
	"hyapp/pkg/logx"
	"hyapp/services/cron-service/internal/config"
	mysqlstorage "hyapp/services/cron-service/internal/storage/mysql"
)

const (
	runStatusRunning   = "running"
	runStatusSucceeded = "succeeded"
	runStatusFailed    = "failed"
)

// Store is the cron metadata boundary used by scheduler loops.
type Store interface {
	TryAcquireTaskLease(ctx context.Context, taskName string, appCode string, ownerNodeID string, nowMs int64, leaseTTL time.Duration) (bool, error)
	StartTaskRun(ctx context.Context, run mysqlstorage.TaskRun) error
	FinishTaskRun(ctx context.Context, run mysqlstorage.TaskRun) error
}

// BatchRequest is the stable input passed to task handlers.
type BatchRequest struct {
	TaskName                string
	AppCode                 string
	RunID                   string
	WorkerID                string
	BatchSize               int
	LockTTL                 time.Duration
	PendingPublishMaxAge    time.Duration
	PublishingSessionMaxAge time.Duration
}

// BatchResult summarizes one owner-service batch response.
type BatchResult struct {
	ClaimedCount   int
	ProcessedCount int
	SuccessCount   int
	FailureCount   int
	HasMore        bool
}

// Handler executes one batch through an owner service client.
type Handler func(ctx context.Context, request BatchRequest) (BatchResult, error)

// Scheduler owns task loops. It never implements domain logic directly.
type Scheduler struct {
	nodeID   string
	tasks    map[string]config.TaskConfig
	store    Store
	handlers map[string]Handler
}

// New creates a scheduler with explicit handlers for tasks that are already migrated.
func New(nodeID string, tasks map[string]config.TaskConfig, store Store, handlers map[string]Handler) *Scheduler {
	return &Scheduler{
		nodeID:   strings.TrimSpace(nodeID),
		tasks:    tasks,
		store:    store,
		handlers: handlers,
	}
}

// Run starts one loop per enabled task and app_code until ctx is cancelled.
func (s *Scheduler) Run(ctx context.Context) {
	if s == nil {
		return
	}
	taskNames := make([]string, 0, len(s.tasks))
	for name := range s.tasks {
		taskNames = append(taskNames, name)
	}
	sort.Strings(taskNames)

	var wg sync.WaitGroup
	for _, taskName := range taskNames {
		task := s.tasks[taskName]
		if !task.Enabled {
			continue
		}
		handler := s.handlers[taskName]
		if handler == nil {
			// 配置可以先落地，但没有 owner-service RPC handler 时不能启动空跑循环。
			logx.Warn(ctx, "cron_task_handler_missing", slog.String("task_name", taskName))
			continue
		}
		for _, appCode := range task.AppCodes {
			wg.Add(1)
			go func(taskName string, appCode string, task config.TaskConfig, handler Handler) {
				defer wg.Done()
				s.runTaskLoop(ctx, taskName, appCode, task, handler)
			}(taskName, appCode, task, handler)
		}
	}
	wg.Wait()
}

func (s *Scheduler) runTaskLoop(ctx context.Context, taskName string, appCode string, task config.TaskConfig, handler Handler) {
	interval := parseDuration(task.Interval, 5*time.Second)
	timer := time.NewTimer(0)
	defer timer.Stop()

	for {
		select {
		case <-ctx.Done():
			return
		case <-timer.C:
			hasMore := s.runOnce(ctx, taskName, appCode, task, handler)
			if hasMore {
				timer.Reset(0)
				continue
			}
			timer.Reset(interval)
		}
	}
}

func (s *Scheduler) runOnce(ctx context.Context, taskName string, appCode string, task config.TaskConfig, handler Handler) bool {
	if s.store == nil {
		logx.Error(ctx, "cron_store_missing", nil, slog.String("task_name", taskName), slog.String("app_code", appCode))
		return false
	}
	now := time.Now()
	lockTTL := parseDuration(task.LockTTL, 30*time.Second)
	acquired, err := s.store.TryAcquireTaskLease(ctx, taskName, appCode, s.nodeID, now.UnixMilli(), lockTTL)
	if err != nil {
		logx.Error(ctx, "cron_task_lease_failed", err, slog.String("task_name", taskName), slog.String("app_code", appCode), slog.String("worker_id", s.nodeID))
		return false
	}
	if !acquired {
		return false
	}

	runID := idgen.New("cronrun")
	run := mysqlstorage.TaskRun{
		RunID:       runID,
		TaskName:    taskName,
		AppCode:     appCode,
		OwnerNodeID: s.nodeID,
		Status:      runStatusRunning,
		StartedAtMs: now.UnixMilli(),
	}
	if err := s.store.StartTaskRun(ctx, run); err != nil {
		logx.Error(ctx, "cron_task_run_start_failed", err, slog.String("task_name", taskName), slog.String("app_code", appCode), slog.String("run_id", runID))
		return false
	}

	timeout := parseDuration(task.Timeout, 5*time.Second)
	taskCtx, cancel := context.WithTimeout(ctx, timeout)
	result, handlerErr := handler(taskCtx, BatchRequest{
		TaskName:                taskName,
		AppCode:                 appCode,
		RunID:                   runID,
		WorkerID:                s.nodeID,
		BatchSize:               task.BatchSize,
		LockTTL:                 lockTTL,
		PendingPublishMaxAge:    parseDuration(task.PendingPublishMaxAge, 0),
		PublishingSessionMaxAge: parseDuration(task.PublishingSessionMaxAge, 0),
	})
	cancel()

	finished := time.Now()
	run.FinishedAtMs = finished.UnixMilli()
	run.DurationMs = finished.Sub(now).Milliseconds()
	run.ClaimedCount = result.ClaimedCount
	run.ProcessedCount = result.ProcessedCount
	run.SuccessCount = result.SuccessCount
	run.FailureCount = result.FailureCount
	run.Status = runStatusSucceeded
	if handlerErr != nil {
		run.Status = runStatusFailed
		run.ErrorMessage = handlerErr.Error()
	}
	if err := s.store.FinishTaskRun(context.Background(), run); err != nil {
		logx.Error(ctx, "cron_task_run_finish_failed", err, slog.String("task_name", taskName), slog.String("app_code", appCode), slog.String("run_id", runID))
		return false
	}
	if handlerErr != nil {
		logx.Error(ctx, "cron_task_run_failed", handlerErr, slog.String("task_name", taskName), slog.String("app_code", appCode), slog.String("run_id", runID), slog.Int("processed_count", run.ProcessedCount), slog.Int64("duration_ms", run.DurationMs))
		return false
	}
	logx.Info(ctx, "cron_task_run_succeeded", slog.String("task_name", taskName), slog.String("app_code", appCode), slog.String("run_id", runID), slog.Int("processed_count", run.ProcessedCount), slog.Int("claimed_count", run.ClaimedCount), slog.Bool("has_more", result.HasMore), slog.Int64("duration_ms", run.DurationMs))
	return result.HasMore
}

func parseDuration(raw string, fallback time.Duration) time.Duration {
	duration, err := time.ParseDuration(strings.TrimSpace(raw))
	if err != nil || duration <= 0 {
		return fallback
	}
	return duration
}