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

package app

import (
	"context"
	"database/sql"
	"errors"
	"fmt"
	"os"
	"strings"
	"time"

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

const (
	knowledgeDeletedCleanupNotifyChannel = "knowledge_deleted_cleanup_events"
	defaultKnowledgeCleanupBatchSize     = 50
	defaultKnowledgeCleanupLockTTL       = 10 * time.Minute
	defaultKnowledgeCleanupMaxWait       = time.Hour
	defaultKnowledgeCleanupRetryBase     = 5 * time.Second
	defaultKnowledgeCleanupRetryMax      = time.Hour
	defaultKnowledgeCleanupListenRetry   = 5 * time.Second
	defaultKnowledgeCleanupProcessTTL    = 2 * time.Minute
)

type knowledgeDeletedCleanupEvent struct {
	ID           int64
	TenantID     int64
	ItemID       int64
	AttemptCount int
}

type KnowledgeDeletedCleanupEventWorker struct {
	service        *KnowledgeService
	logger         *zap.Logger
	workerID       string
	batchSize      int
	lockTTL        time.Duration
	maxWait        time.Duration
	processTimeout time.Duration
}

func NewKnowledgeDeletedCleanupEventWorker(service *KnowledgeService, logger *zap.Logger) *KnowledgeDeletedCleanupEventWorker {
	hostname, _ := os.Hostname()
	if hostname == "" {
		hostname = "unknown"
	}
	return &KnowledgeDeletedCleanupEventWorker{
		service:        service,
		logger:         logger,
		workerID:       fmt.Sprintf("%s:%d:%d", hostname, os.Getpid(), time.Now().UnixNano()),
		batchSize:      defaultKnowledgeCleanupBatchSize,
		lockTTL:        defaultKnowledgeCleanupLockTTL,
		maxWait:        defaultKnowledgeCleanupMaxWait,
		processTimeout: defaultKnowledgeCleanupProcessTTL,
	}
}

func (w *KnowledgeDeletedCleanupEventWorker) Start(ctx context.Context) {
	if w == nil {
		return
	}
	go w.run(ctx)
}

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

	wake := make(chan struct{}, 1)
	go w.listen(ctx, wake)
	signalKnowledgeCleanup(wake)

	for {
		processed, err := w.processDue(ctx)
		if err != nil && ctx.Err() == nil && w.logger != nil {
			w.logger.Warn("knowledge deleted cleanup worker failed", zap.Error(err))
		}
		if processed > 0 {
			continue
		}

		waitFor := w.nextWait(ctx)
		timer := time.NewTimer(waitFor)
		select {
		case <-ctx.Done():
			timer.Stop()
			return
		case <-wake:
			timer.Stop()
		case <-timer.C:
		}
	}
}

func (w *KnowledgeDeletedCleanupEventWorker) listen(ctx context.Context, wake chan<- struct{}) {
	for {
		if err := w.listenOnce(ctx, wake); err != nil && ctx.Err() == nil && w.logger != nil {
			w.logger.Warn("knowledge deleted cleanup listener stopped", zap.Error(err))
		}
		signalKnowledgeCleanup(wake)

		select {
		case <-ctx.Done():
			return
		case <-time.After(defaultKnowledgeCleanupListenRetry):
			signalKnowledgeCleanup(wake)
		}
	}
}

func (w *KnowledgeDeletedCleanupEventWorker) listenOnce(ctx context.Context, wake chan<- struct{}) error {
	conn, err := w.service.pool.Acquire(ctx)
	if err != nil {
		return err
	}
	defer conn.Release()

	if _, err := conn.Exec(ctx, fmt.Sprintf("LISTEN %s", knowledgeDeletedCleanupNotifyChannel)); err != nil {
		return err
	}
	defer func() {
		unlistenCtx, cancel := context.WithTimeout(context.Background(), time.Second)
		defer cancel()
		_, _ = conn.Exec(unlistenCtx, fmt.Sprintf("UNLISTEN %s", knowledgeDeletedCleanupNotifyChannel))
	}()

	for {
		if _, err := conn.Conn().WaitForNotification(ctx); err != nil {
			if errors.Is(err, context.Canceled) || errors.Is(err, context.DeadlineExceeded) {
				return nil
			}
			return err
		}
		signalKnowledgeCleanup(wake)
	}
}

func (w *KnowledgeDeletedCleanupEventWorker) processDue(parent context.Context) (int, error) {
	events, err := w.service.claimDeletedCleanupEvents(parent, w.workerID, w.batchSize, w.lockTTL)
	if err != nil || len(events) == 0 {
		return len(events), err
	}

	for _, event := range events {
		ctx, cancel := context.WithTimeout(parent, w.processTimeout)
		err := w.service.ProcessDeletedKnowledgeCleanup(ctx, event.TenantID, event.ItemID)
		cancel()

		if err != nil {
			if markErr := w.service.failDeletedCleanupEvent(parent, event.ID, err.Error(), retryDelayForKnowledgeCleanup(event.AttemptCount)); markErr != nil && w.logger != nil {
				w.logger.Warn("knowledge deleted cleanup event fail mark failed",
					zap.Error(markErr),
					zap.Int64("event_id", event.ID),
				)
			}
			if w.logger != nil {
				w.logger.Warn("knowledge deleted cleanup event failed",
					zap.Error(err),
					zap.Int64("event_id", event.ID),
					zap.Int64("tenant_id", event.TenantID),
					zap.Int64("knowledge_item_id", event.ItemID),
				)
			}
			continue
		}

		if err := w.service.completeDeletedCleanupEvent(parent, event.ID); err != nil && w.logger != nil {
			w.logger.Warn("knowledge deleted cleanup event complete mark failed",
				zap.Error(err),
				zap.Int64("event_id", event.ID),
			)
		}
	}
	return len(events), nil
}

func (w *KnowledgeDeletedCleanupEventWorker) nextWait(ctx context.Context) time.Duration {
	waitFor := w.maxWait
	if waitFor <= 0 {
		waitFor = defaultKnowledgeCleanupMaxWait
	}

	nextAt, err := w.service.nextDeletedCleanupDueAt(ctx, w.lockTTL)
	if err != nil {
		if w.logger != nil {
			w.logger.Warn("knowledge deleted cleanup next due query failed", zap.Error(err))
		}
		return minDuration(waitFor, time.Minute)
	}
	if nextAt == nil {
		return waitFor
	}

	until := time.Until(*nextAt)
	if until <= 0 {
		return time.Millisecond
	}
	return minDuration(until, waitFor)
}

func (s *KnowledgeService) enqueueDeletedKnowledgeCleanupEvents(ctx context.Context, tx pgx.Tx, tenantID int64, itemIDs []int64) error {
	if s == nil || tx == nil || tenantID <= 0 || len(itemIDs) == 0 {
		return nil
	}

	ids := normalizeInt64IDs(itemIDs)
	if len(ids) == 0 {
		return nil
	}

	if _, err := tx.Exec(ctx, `
		INSERT INTO knowledge_deleted_cleanup_events (tenant_id, knowledge_item_id)
		SELECT $1, item_id
		FROM unnest($2::bigint[]) AS item_ids(item_id)
		WHERE item_id > 0
		ON CONFLICT (tenant_id, knowledge_item_id) WHERE completed_at IS NULL
		DO UPDATE SET status = 'pending',
		              available_at = NOW(),
		              locked_by = NULL,
		              locked_at = NULL,
		              last_error = NULL,
		              updated_at = NOW()
	`, tenantID, ids); err != nil {
		return err
	}

	_, err := tx.Exec(ctx, `SELECT pg_notify($1, '')`, knowledgeDeletedCleanupNotifyChannel)
	return err
}

func (s *KnowledgeService) ProcessDeletedKnowledgeCleanup(ctx context.Context, tenantID, itemID int64) error {
	if s == nil {
		return nil
	}
	if !s.startCleanupJob(tenantID, itemID) {
		return nil
	}
	defer s.finishCleanupJob(tenantID, itemID)
	return s.cleanupDeletedKnowledgeItem(ctx, tenantID, itemID)
}

func (s *KnowledgeService) claimDeletedCleanupEvents(
	ctx context.Context,
	workerID string,
	batchSize int,
	lockTTL time.Duration,
) ([]knowledgeDeletedCleanupEvent, error) {
	if s == nil || s.pool == nil {
		return nil, nil
	}
	if batchSize <= 0 {
		batchSize = defaultKnowledgeCleanupBatchSize
	}
	if lockTTL <= 0 {
		lockTTL = defaultKnowledgeCleanupLockTTL
	}

	staleBefore := time.Now().UTC().Add(-lockTTL)
	rows, err := s.pool.Query(ctx, `
		WITH next_events AS (
			SELECT id
			FROM knowledge_deleted_cleanup_events
			WHERE completed_at IS NULL
			  AND (
			      (status IN ('pending', 'failed') AND available_at <= NOW())
			      OR (status = 'processing' AND locked_at < $2)
			  )
			ORDER BY available_at ASC, id ASC
			FOR UPDATE SKIP LOCKED
			LIMIT $1
		)
		UPDATE knowledge_deleted_cleanup_events e
		SET status = 'processing',
		    attempt_count = attempt_count + 1,
		    locked_by = $3,
		    locked_at = NOW(),
		    updated_at = NOW()
		FROM next_events
		WHERE e.id = next_events.id
		RETURNING e.id, e.tenant_id, e.knowledge_item_id, e.attempt_count
	`, batchSize, staleBefore, strings.TrimSpace(workerID))
	if err != nil {
		return nil, err
	}
	defer rows.Close()

	events := make([]knowledgeDeletedCleanupEvent, 0, batchSize)
	for rows.Next() {
		var event knowledgeDeletedCleanupEvent
		if err := rows.Scan(&event.ID, &event.TenantID, &event.ItemID, &event.AttemptCount); err != nil {
			return nil, err
		}
		events = append(events, event)
	}
	return events, rows.Err()
}

func (s *KnowledgeService) completeDeletedCleanupEvent(ctx context.Context, eventID int64) error {
	if s == nil || s.pool == nil || eventID <= 0 {
		return nil
	}
	_, err := s.pool.Exec(ctx, `
		UPDATE knowledge_deleted_cleanup_events
		SET status = 'completed',
		    completed_at = NOW(),
		    locked_by = NULL,
		    locked_at = NULL,
		    last_error = NULL,
		    updated_at = NOW()
		WHERE id = $1
		  AND completed_at IS NULL
	`, eventID)
	return err
}

func (s *KnowledgeService) failDeletedCleanupEvent(ctx context.Context, eventID int64, message string, delay time.Duration) error {
	if s == nil || s.pool == nil || eventID <= 0 {
		return nil
	}
	if delay <= 0 {
		delay = defaultKnowledgeCleanupRetryBase
	}
	nextAt := time.Now().UTC().Add(delay)
	_, err := s.pool.Exec(ctx, `
		UPDATE knowledge_deleted_cleanup_events
		SET status = 'failed',
		    available_at = $2,
		    locked_by = NULL,
		    locked_at = NULL,
		    last_error = $3,
		    updated_at = NOW()
		WHERE id = $1
		  AND completed_at IS NULL
	`, eventID, nextAt, strings.TrimSpace(message))
	return err
}

func (s *KnowledgeService) nextDeletedCleanupDueAt(ctx context.Context, lockTTL time.Duration) (*time.Time, error) {
	if s == nil || s.pool == nil {
		return nil, nil
	}
	if lockTTL <= 0 {
		lockTTL = defaultKnowledgeCleanupLockTTL
	}

	var due sql.NullTime
	err := s.pool.QueryRow(ctx, `
		SELECT MIN(due_at)
		FROM (
			SELECT available_at AS due_at
			FROM knowledge_deleted_cleanup_events
			WHERE completed_at IS NULL
			  AND status IN ('pending', 'failed')
			UNION ALL
			SELECT locked_at + ($1::double precision * INTERVAL '1 second') AS due_at
			FROM knowledge_deleted_cleanup_events
			WHERE completed_at IS NULL
			  AND status = 'processing'
			  AND locked_at IS NOT NULL
		) due_events
	`, lockTTL.Seconds()).Scan(&due)
	if err != nil {
		return nil, err
	}
	if !due.Valid {
		return nil, nil
	}
	return &due.Time, nil
}

func retryDelayForKnowledgeCleanup(attempt int) time.Duration {
	if attempt <= 1 {
		return defaultKnowledgeCleanupRetryBase
	}
	shift := attempt - 1
	if shift > 8 {
		shift = 8
	}
	delay := defaultKnowledgeCleanupRetryBase * time.Duration(1<<shift)
	return minDuration(delay, defaultKnowledgeCleanupRetryMax)
}

func signalKnowledgeCleanup(ch chan<- struct{}) {
	select {
	case ch <- struct{}{}:
	default:
	}
}

func minDuration(a, b time.Duration) time.Duration {
	if a < b {
		return a
	}
	return b
}