geo/server/internal/tenant/app/monitoring_lease_recovery.go

package app

import (
	"context"
	"time"

	"github.com/jackc/pgx/v5"
)

const (
	monitoringLeaseMaxRetryCount    = 1
	monitoringLeaseExpiredErrPrefix = "monitoring lease expired before result callback"
	monitoringStaleTaskDropReason   = "stale_business_date"
	monitoringStaleTaskDropError    = "monitoring task dropped after business day rollover"
)

func expireMonitoringLeasedTasks(ctx context.Context, tx pgx.Tx, tenantID *int64, now time.Time, limits ...int) (int64, error) {
	limit := 0
	if len(limits) > 0 {
		limit = limits[0]
	}
	if limit <= 0 {
		limit = 1000
	}
	tag, err := tx.Exec(ctx, `
		WITH candidates AS (
			SELECT id
			FROM monitoring_collect_tasks
			WHERE collector_type = $2
			  AND status = 'leased'
			  AND lease_expires_at IS NOT NULL
			  AND lease_expires_at < $3
			  AND ($4::bigint IS NULL OR tenant_id = $4)
			ORDER BY lease_expires_at ASC, id ASC
			FOR UPDATE SKIP LOCKED
			LIMIT $5
		)
			UPDATE monitoring_collect_tasks
			SET status = 'expired',
			    lease_token_hash = NULL,
			    leased_to_executor = NULL,
			    leased_at = NULL,
			    lease_expires_at = NULL,
			    error_message = COALESCE(NULLIF(error_message, ''), $1),
			    request_payload_json = (
				CASE
					WHEN request_payload_json IS NULL OR jsonb_typeof(request_payload_json) <> 'object'
						THEN '{}'::jsonb
					ELSE request_payload_json
				END
				) || jsonb_build_object(
					'lease_retry_count',
					(
						CASE
							WHEN COALESCE(request_payload_json ->> 'lease_retry_count', '') ~ '^[0-9]+$'
								THEN (request_payload_json ->> 'lease_retry_count')::int
							ELSE 0
						END
					) + 1
				),
			    updated_at = NOW()
		FROM candidates c
		WHERE monitoring_collect_tasks.id = c.id
	`, monitoringLeaseExpiredErrPrefix, monitoringCollectorType, now.UTC(), nullableInt64(tenantID), limit)
	if err != nil {
		return 0, monitoringInternalError(50041, "lease_expire_failed", "failed to expire overdue monitoring lease tasks", err)
	}
	return tag.RowsAffected(), nil
}

func requeueMonitoringExpiredTasks(ctx context.Context, tx pgx.Tx, tenantID int64, businessDate time.Time) (int64, error) {
	tag, err := tx.Exec(ctx, `
		UPDATE monitoring_collect_tasks
		SET status = 'pending',
		    planned_at = NOW(),
		    callback_received_at = NULL,
		    completed_at = NULL,
		    skip_reason = NULL,
		    error_message = NULL,
		    updated_at = NOW()
		WHERE tenant_id = $1
		  AND collector_type = $2
		  AND business_date = $3::date
		  AND status = 'expired'
		  AND (
			CASE
				WHEN COALESCE(request_payload_json ->> 'lease_retry_count', '') ~ '^[0-9]+$'
					THEN (request_payload_json ->> 'lease_retry_count')::int
				ELSE 0
			END
		  ) <= $4
	`, tenantID, monitoringCollectorType, businessDate.Format("2006-01-02"), monitoringLeaseMaxRetryCount)
	if err != nil {
		return 0, monitoringInternalError(50041, "lease_requeue_failed", "failed to recycle expired monitoring lease tasks", err)
	}
	return tag.RowsAffected(), nil
}

func skipMonitoringTasksBeforeBusinessDate(
	ctx context.Context,
	tx pgx.Tx,
	tenantID int64,
	businessDate time.Time,
	executionOwner string,
) (int64, error) {
	if tx == nil {
		return 0, nil
	}
	if executionOwner == "" {
		executionOwner = "legacy"
	}

	tag, err := tx.Exec(ctx, `
		UPDATE monitoring_collect_tasks
		SET status = 'skipped',
		    lease_token_hash = NULL,
		    leased_to_executor = NULL,
		    leased_at = NULL,
		    lease_expires_at = NULL,
		    completed_at = COALESCE(completed_at, NOW()),
		    skip_reason = $4,
		    error_message = COALESCE(NULLIF(error_message, ''), $5),
		    updated_at = NOW()
		WHERE tenant_id = $1
		  AND collector_type = $2
		  AND business_date < $3::date
		  AND COALESCE(execution_owner, 'legacy') = $6
		  AND callback_received_at IS NULL
		  AND status IN ('pending', 'leased', 'expired')
	`, tenantID, monitoringCollectorType, businessDate.Format("2006-01-02"), monitoringStaleTaskDropReason, monitoringStaleTaskDropError, executionOwner)
	if err != nil {
		return 0, monitoringInternalError(50041, "task_skip_failed", "failed to drop stale monitoring tasks", err)
	}
	return tag.RowsAffected(), nil
}

func ExpireMonitoringLeasedTasks(ctx context.Context, tx pgx.Tx, tenantID *int64, now time.Time, limit int) (int64, error) {
	return expireMonitoringLeasedTasks(ctx, tx, tenantID, now, limit)
}