geo/server/internal/scheduler/publish_lease_recovery_worker.go

package scheduler

import (
	"context"
	"time"

	"go.uber.org/zap"

	tenantapp "github.com/geo-platform/tenant-api/internal/tenant/app"
)

const (
	defaultPublishLeaseRecoveryInterval = 3 * time.Minute
	defaultPublishLeaseRecoveryTimeout  = 30 * time.Second
	defaultPublishLeaseRecoveryLimit    = 1000
	defaultPublishQueuedTaskTimeout     = 3 * 24 * time.Hour
)

type PublishLeaseRecoveryWorker struct {
	service           *tenantapp.DesktopTaskService
	logger            *zap.Logger
	interval          time.Duration
	timeout           time.Duration
	limit             int
	queuedTaskTimeout time.Duration
}

func NewPublishLeaseRecoveryWorker(service *tenantapp.DesktopTaskService, logger *zap.Logger) *PublishLeaseRecoveryWorker {
	return &PublishLeaseRecoveryWorker{
		service:           service,
		logger:            logger,
		interval:          defaultPublishLeaseRecoveryInterval,
		timeout:           defaultPublishLeaseRecoveryTimeout,
		limit:             defaultPublishLeaseRecoveryLimit,
		queuedTaskTimeout: defaultPublishQueuedTaskTimeout,
	}
}

func (w *PublishLeaseRecoveryWorker) WithQueuedTaskTimeout(timeout time.Duration) *PublishLeaseRecoveryWorker {
	if w != nil && timeout > 0 {
		w.queuedTaskTimeout = timeout
	}
	return w
}

func (w *PublishLeaseRecoveryWorker) Start(ctx context.Context) {
	go w.Run(ctx)
}

func (w *PublishLeaseRecoveryWorker) Run(ctx context.Context) {
	if w == nil || w.service == nil {
		return
	}
	w.run(ctx)
}

func (w *PublishLeaseRecoveryWorker) run(ctx context.Context) {
	w.runOnce(context.Background())

	ticker := time.NewTicker(w.interval)
	defer ticker.Stop()

	for {
		select {
		case <-ctx.Done():
			return
		case <-ticker.C:
			w.runOnce(context.Background())
		}
	}
}

func (w *PublishLeaseRecoveryWorker) runOnce(parent context.Context) {
	ctx, cancel := context.WithTimeout(parent, w.timeout)
	defer cancel()

	result, err := w.service.RecoverExpiredPublishTasks(ctx, w.limit)
	if err != nil {
		if w.logger != nil {
			w.logger.Warn("publish lease recovery failed", zap.Error(err))
		}
		return
	}

	queuedTaskTimeout := w.queuedTaskTimeout
	if queuedTaskTimeout <= 0 {
		queuedTaskTimeout = defaultPublishQueuedTaskTimeout
	}
	cancelledQueued, err := w.service.CancelStaleQueuedPublishTasks(ctx, queuedTaskTimeout, w.limit)
	if err != nil {
		if w.logger != nil {
			w.logger.Warn("publish queued task timeout cleanup failed", zap.Error(err))
		}
		return
	}
	result.CancelledQueued = cancelledQueued

	if (result.Requeued > 0 || result.Failed > 0 || result.Uncertain > 0 || result.CancelledQueued > 0) && w.logger != nil {
		w.logger.Info("publish lease recovery completed",
			zap.Int("requeued_task_count", result.Requeued),
			zap.Int("failed_task_count", result.Failed),
			zap.Int("uncertain_task_count", result.Uncertain),
			zap.Int("cancelled_queued_task_count", result.CancelledQueued),
		)
	}
}

func publishQueuedTaskTimeoutFromRun(run JobRunContext) time.Duration {
	timeout := AsDuration(run.Config, "queued_task_timeout", defaultPublishQueuedTaskTimeout)
	if seconds := AsInt(run.Config, "queued_task_timeout_seconds", 0); seconds > 0 {
		timeout = time.Duration(seconds) * time.Second
	}
	if timeout <= 0 {
		return defaultPublishQueuedTaskTimeout
	}
	return timeout
}

func (w *PublishLeaseRecoveryWorker) RunOnce(parent context.Context, run JobRunContext) (map[string]any, error) {
	if w == nil || w.service == nil {
		return map[string]any{"skipped": true, "reason": "worker_not_ready"}, nil
	}
	startedAt := time.Now()
	timeout := w.timeout
	if run.Job != nil && run.Job.TimeoutSeconds > 0 {
		timeout = time.Duration(run.Job.TimeoutSeconds) * time.Second
	}
	limit := w.limit
	if run.Job != nil && run.Job.BatchSize != nil && *run.Job.BatchSize > 0 {
		limit = *run.Job.BatchSize
	}
	if value := AsInt(run.Config, "batch_size", 0); value > 0 && (run.Job == nil || run.Job.BatchSize == nil) {
		limit = value
	}
	ctx, cancel := context.WithTimeout(parent, timeout)
	defer cancel()

	result, err := w.service.RecoverExpiredPublishTasks(ctx, limit)
	if err != nil {
		return map[string]any{"stage": "recover"}, err
	}
	queuedTimeout := publishQueuedTaskTimeoutFromRun(run)
	cancelledQueued, err := w.service.CancelStaleQueuedPublishTasks(ctx, queuedTimeout, limit)
	if err != nil {
		return map[string]any{"stage": "cancel_stale_queued"}, err
	}
	result.CancelledQueued = cancelledQueued
	return map[string]any{
		"requeued_task_count":         result.Requeued,
		"failed_task_count":           result.Failed,
		"uncertain_task_count":        result.Uncertain,
		"cancelled_queued_task_count": result.CancelledQueued,
		"queued_task_timeout_seconds": int64(queuedTimeout.Seconds()),
		"duration_ms":                 time.Since(startedAt).Milliseconds(),
		"batch_size":                  limit,
	}, nil
}