geo/server/internal/scheduler/monitoring_result_recovery_worker.go

package scheduler

import (
	"context"
	"fmt"
	"time"

	"github.com/jackc/pgx/v5"
	"github.com/jackc/pgx/v5/pgxpool"
	"go.uber.org/zap"

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

const (
	defaultMonitoringResultRecoveryInterval = 1 * time.Minute
	defaultMonitoringResultRecoveryTimeout  = 30 * time.Second
	defaultMonitoringResultRecoveryLimit    = 100
	defaultMonitoringResultRecoveryClaimTTL = 5 * time.Minute
)

type MonitoringResultRecoveryWorker struct {
	monitoringPool *pgxpool.Pool
	service        *tenantapp.MonitoringCallbackService
	logger         *zap.Logger
	interval       time.Duration
	timeout        time.Duration
	limit          int
}

type monitoringRecoverableTaskResult struct {
	ID                 int64
	CallbackReceivedAt time.Time
	RequestPayloadJSON []byte
}

func NewMonitoringResultRecoveryWorker(
	monitoringPool *pgxpool.Pool,
	service *tenantapp.MonitoringCallbackService,
	logger *zap.Logger,
) *MonitoringResultRecoveryWorker {
	return &MonitoringResultRecoveryWorker{
		monitoringPool: monitoringPool,
		service:        service,
		logger:         logger,
		interval:       defaultMonitoringResultRecoveryInterval,
		timeout:        defaultMonitoringResultRecoveryTimeout,
		limit:          defaultMonitoringResultRecoveryLimit,
	}
}

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

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

func (w *MonitoringResultRecoveryWorker) 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 *MonitoringResultRecoveryWorker) runOnce(parent context.Context) {
	ctx, cancel := context.WithTimeout(parent, w.timeout)
	defer cancel()

	tx, err := w.monitoringPool.Begin(ctx)
	if err != nil {
		if w.logger != nil {
			w.logger.Warn("monitoring result recovery begin failed", zap.Error(err))
		}
		return
	}
	defer tx.Rollback(ctx)

	items, err := loadMonitoringRecoverableTaskResults(ctx, tx, w.limit, defaultMonitoringResultRecoveryClaimTTL)
	if err != nil {
		if w.logger != nil {
			w.logger.Warn("monitoring result recovery failed", tenantapp.MonitoringWorkerErrorFields(err)...)
		}
		return
	}

	if err := tx.Commit(ctx); err != nil {
		if w.logger != nil {
			w.logger.Warn("monitoring result recovery commit failed", zap.Error(err))
		}
		return
	}

	if len(items) == 0 {
		return
	}

	republishedCount := 0
	for _, item := range items {
		if w.tryRepublish(parent, item) {
			republishedCount++
		}
	}

	if republishedCount > 0 && w.logger != nil {
		w.logger.Info("monitoring result recovery republished tasks", zap.Int("task_count", republishedCount))
	}
}

func (w *MonitoringResultRecoveryWorker) RunOnce(parent context.Context, run JobRunContext) (map[string]any, error) {
	if w == nil || w.monitoringPool == 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()

	tx, err := w.monitoringPool.Begin(ctx)
	if err != nil {
		return map[string]any{"stage": "begin"}, err
	}
	defer tx.Rollback(ctx)

	claimTTL := AsDuration(run.Config, "claim_ttl", defaultMonitoringResultRecoveryClaimTTL)
	items, err := loadMonitoringRecoverableTaskResults(ctx, tx, limit, claimTTL)
	if err != nil {
		return map[string]any{"stage": "load"}, err
	}
	if err := tx.Commit(ctx); err != nil {
		return map[string]any{"stage": "commit"}, err
	}

	republishedCount := 0
	for _, item := range items {
		if w.tryRepublish(parent, item) {
			republishedCount++
		}
	}
	return map[string]any{
		"recoverable_count": len(items),
		"republished_count": republishedCount,
		"duration_ms":       time.Since(startedAt).Milliseconds(),
		"batch_size":        limit,
		"claim_ttl_seconds": int64(claimTTL.Seconds()),
	}, nil
}

func (w *MonitoringResultRecoveryWorker) tryRepublish(parent context.Context, item monitoringRecoverableTaskResult) bool {
	if err := tenantapp.ValidateMonitoringTaskResultEnvelope(item.RequestPayloadJSON); err != nil {
		w.service.PatchTaskResultQueueMeta(item.ID, item.CallbackReceivedAt, tenantapp.MonitoringTaskResultQueueMetaPatch{
			QueueStatus:           "invalid_payload",
			QueueClaimedAt:        nil,
			QueueEnqueuedAt:       nil,
			LastQueuePublishError: stringPtr(err.Error()),
		})
		if w.logger != nil {
			w.logger.Warn("monitoring result recovery skipped invalid payload", zap.Int64("task_id", item.ID), zap.Error(err))
		}
		return false
	}

	ctx, cancel := context.WithTimeout(parent, 5*time.Second)
	defer cancel()

	if err := w.service.PublishTaskResultEnvelope(ctx, item.RequestPayloadJSON); err != nil {
		w.service.PatchTaskResultQueueMeta(item.ID, item.CallbackReceivedAt, tenantapp.MonitoringTaskResultQueueMetaPatch{
			QueueStatus:           "publish_failed",
			QueueClaimedAt:        nil,
			QueueEnqueuedAt:       nil,
			LastQueuePublishError: stringPtr(err.Error()),
		})
		if w.logger != nil {
			w.logger.Warn("monitoring result recovery republish failed", zap.Int64("task_id", item.ID), zap.Error(err))
		}
		return false
	}

	enqueuedAt := time.Now().UTC().Round(0)
	w.service.PatchTaskResultQueueMeta(item.ID, item.CallbackReceivedAt, tenantapp.MonitoringTaskResultQueueMetaPatch{
		QueueStatus:           "queued",
		QueueClaimedAt:        nil,
		QueueEnqueuedAt:       &enqueuedAt,
		LastQueuePublishError: nil,
	})

	return true
}

func loadMonitoringRecoverableTaskResults(ctx context.Context, tx pgx.Tx, limit int, claimTTL time.Duration) ([]monitoringRecoverableTaskResult, error) {
	if limit <= 0 {
		limit = defaultMonitoringResultRecoveryLimit
	}
	if claimTTL <= 0 {
		claimTTL = defaultMonitoringResultRecoveryClaimTTL
	}

	rows, err := tx.Query(ctx, `
		WITH candidates AS (
			SELECT
				id,
				callback_received_at,
				request_payload_json
			FROM monitoring_collect_tasks
			WHERE collector_type = $1
			  AND (
				status = 'received'
				OR (
					COALESCE(execution_owner, 'legacy') = 'desktop_tasks'
					AND status = 'pending'
					AND callback_received_at IS NOT NULL
				)
			  )
			  AND callback_received_at IS NOT NULL
			  AND request_payload_json IS NOT NULL
			  AND jsonb_typeof(request_payload_json) = 'object'
			  AND (
				COALESCE(NULLIF(request_payload_json ->> 'queue_status', ''), 'pending') IN ('pending', 'publish_failed')
				OR (
					COALESCE(NULLIF(request_payload_json ->> 'queue_status', ''), 'pending') = 'claiming'
					AND COALESCE(NULLIF(request_payload_json ->> 'queue_claimed_at', ''), '') <> ''
					AND NULLIF(request_payload_json ->> 'queue_claimed_at', '')::timestamptz < NOW() - $3::interval
				)
			  )
			ORDER BY callback_received_at ASC, id ASC
			FOR UPDATE SKIP LOCKED
			LIMIT $2
		),
		claimed AS (
			UPDATE monitoring_collect_tasks t
			SET request_payload_json = (
					CASE
						WHEN t.request_payload_json IS NULL OR jsonb_typeof(t.request_payload_json) <> 'object'
							THEN '{}'::jsonb
						ELSE t.request_payload_json
					END
				) || jsonb_build_object(
					'queue_status', 'claiming',
					'queue_claimed_at', NOW()::text
				),
				updated_at = NOW()
			FROM candidates c
			WHERE t.id = c.id
			  AND t.callback_received_at = c.callback_received_at
			RETURNING c.id, c.callback_received_at, c.request_payload_json
		)
		SELECT
			id,
			callback_received_at,
			request_payload_json
		FROM claimed
		ORDER BY callback_received_at ASC, id ASC
	`, "desktop", limit, fmt.Sprintf("%d milliseconds", claimTTL.Milliseconds()))
	if err != nil {
		return nil, tenantapp.InternalMonitoringError(50041, "result_recovery_query_failed", "failed to load recoverable monitoring results", err)
	}
	defer rows.Close()

	result := make([]monitoringRecoverableTaskResult, 0, limit)
	for rows.Next() {
		var item monitoringRecoverableTaskResult
		if scanErr := rows.Scan(&item.ID, &item.CallbackReceivedAt, &item.RequestPayloadJSON); scanErr != nil {
			return nil, tenantapp.InternalMonitoringError(50041, "result_recovery_scan_failed", "failed to parse recoverable monitoring results", scanErr)
		}
		result = append(result, item)
	}
	if err := rows.Err(); err != nil {
		return nil, tenantapp.InternalMonitoringError(50041, "result_recovery_scan_failed", "failed to iterate recoverable monitoring results", err)
	}

	return result, nil
}

func stringPtr(value string) *string {
	if value == "" {
		return nil
	}
	result := value
	return &result
}