NixMsg/internal/app/message/recover.go

package message

import (
	"context"
	"database/sql"
	"time"
)

// RecoverOnStart 启动恢复（DEVELOPMENT 7.8）：只做 SQL 修正，分发交给调度循环。
func (a *App) RecoverOnStart(ctx context.Context) error {
	nowMs := a.now().UnixMilli()
	graceMs := a.lim.GraceSeconds * 1000
	minExpire := nowMs + graceMs

	return a.db.Queue.Do(ctx, func(tx *sql.Tx) error {
		if _, err := tx.Exec(`
UPDATE deliveries SET
  pushed_conn = NULL,
  expire_at = CASE
    WHEN expire_at IS NULL OR expire_at < ? THEN ?
    ELSE expire_at
  END,
  updated_at = ?
WHERE state = 'pending'`, minExpire, minExpire, nowMs); err != nil {
			return err
		}
		_, err := tx.Exec(`
UPDATE endpoints SET offline_since = ?
WHERE online_since IS NOT NULL
  AND (offline_since IS NULL OR online_since > offline_since)`, nowMs)
		return err
	})
}

// CleanupOnce 兼容旧调用：先到期处理再做一次保留清理。
func (a *App) CleanupOnce(ctx context.Context, nowMs int64) error {
	if err := a.ExpireOnce(ctx, nowMs); err != nil {
		return err
	}
	return a.PurgeOnce(ctx, nowMs)
}

// ExpireOnce 处理未推送且到期的 pending，并给僵尸不保留投递补宽限。每写操作最多 expireBatch 条。
func (a *App) ExpireOnce(ctx context.Context, nowMs int64) error {
	deadline := time.Now().Add(expireBudget)
	for {
		if err := ctx.Err(); err != nil {
			return err
		}
		if time.Now().After(deadline) {
			return nil
		}
		n, err := a.expireOnceBatch(ctx, nowMs)
		if err != nil {
			return err
		}
		if n == 0 {
			break
		}
	}
	a.flushRevokes(ctx)
	return nil
}

func (a *App) expireOnceBatch(ctx context.Context, nowMs int64) (int, error) {
	var n int
	err := a.db.Queue.Do(ctx, func(tx *sql.Tx) error {
		if err := a.fillZombieExpireTx(tx, nowMs); err != nil {
			return err
		}
		rows, err := tx.Query(`
SELECT d.seq, d.endpoint_id, d.keep, m.sender_id, m.id
FROM deliveries d
JOIN messages m ON m.seq = d.seq
WHERE d.state = 'pending' AND d.pushed_conn IS NULL
  AND d.expire_at IS NOT NULL AND d.expire_at <= ?
LIMIT ?`, nowMs, expireBatch)
		if err != nil {
			return err
		}
		type item struct {
			seq, keep            int64
			endpointID, senderID string
			msgID                string
		}
		var list []item
		for rows.Next() {
			var it item
			if err := rows.Scan(&it.seq, &it.endpointID, &it.keep, &it.senderID, &it.msgID); err != nil {
				_ = rows.Close()
				return err
			}
			list = append(list, it)
		}
		_ = rows.Close()
		n = len(list)
		for _, it := range list {
			state := DeliveryDropped
			reason := ReasonOffline
			if it.keep != 0 {
				state = DeliveryExpired
				reason = ReasonTTL
			}
			if err := a.finishDeliveryTx(tx, it.seq, it.endpointID, it.senderID, it.msgID, state, reason, false, nowMs); err != nil {
				return err
			}
		}
		return finalizeStuckDispatchedTx(tx, nowMs, a.lim.RecordRetentionDays)
	})
	return n, err
}

func (a *App) fillZombieExpireTx(tx *sql.Tx, nowMs int64) error {
	graceMs := a.lim.GraceSeconds * 1000
	if graceMs < 0 {
		graceMs = 0
	}
	deadline := nowMs + graceMs
	rows, err := tx.Query(`
SELECT DISTINCT endpoint_id FROM deliveries
WHERE state = 'pending' AND keep = 0 AND pushed_conn IS NULL AND expire_at IS NULL
LIMIT 500`)
	if err != nil {
		return err
	}
	seen := map[string]struct{}{}
	var eps []string
	for rows.Next() {
		var ep string
		if err := rows.Scan(&ep); err != nil {
			_ = rows.Close()
			return err
		}
		if _, ok := seen[ep]; ok {
			continue
		}
		seen[ep] = struct{}{}
		if a.isReadyEndpoint(ep) {
			continue
		}
		eps = append(eps, ep)
	}
	_ = rows.Close()
	for _, ep := range eps {
		if _, err := tx.Exec(`
UPDATE deliveries SET expire_at = ?, updated_at = ?
WHERE endpoint_id = ? AND state = 'pending' AND keep = 0 AND pushed_conn IS NULL AND expire_at IS NULL`,
			deadline, nowMs, ep); err != nil {
			return err
		}
	}
	return nil
}

// PurgeOnce 分批删除过期记录、回执和防重行，然后 wal_checkpoint + optimize。
func (a *App) PurgeOnce(ctx context.Context, nowMs int64) error {
	a.purgeMu.Lock()
	a.lastPurge = a.now()
	a.purgeMu.Unlock()

	if err := a.purgeCompletedMessages(ctx, nowMs); err != nil {
		return err
	}
	if err := a.purgeReceipts(ctx, nowMs); err != nil {
		return err
	}
	if err := a.purgeSendKeys(ctx, nowMs); err != nil {
		return err
	}
	if a.db != nil && a.db.Queue != nil {
		if err := a.db.Queue.Checkpoint(ctx); err != nil {
			return err
		}
		_ = a.db.Queue.Optimize(ctx)
	}
	return nil
}

func (a *App) purgeCompletedMessages(ctx context.Context, nowMs int64) error {
	days := a.lim.RecordRetentionDays
	if days < 0 {
		return nil
	}
	for {
		if err := ctx.Err(); err != nil {
			return err
		}
		var seqs []int64
		err := a.db.Queue.Do(ctx, func(tx *sql.Tx) error {
			q := `
SELECT m.seq FROM messages m
WHERE m.state = 'completed'`
			args := []any{}
			if days > 0 {
				cutoff := nowMs - int64(days)*24*3600*1000
				q += ` AND COALESCE(
  (SELECT MAX(d.updated_at) FROM deliveries d WHERE d.seq = m.seq),
  m.send_at
) < ?`
				args = append(args, cutoff)
			}
			q += ` LIMIT ?`
			args = append(args, purgeMsgBatch)
			rows, err := tx.Query(q, args...)
			if err != nil {
				return err
			}
			for rows.Next() {
				var seq int64
				if err := rows.Scan(&seq); err != nil {
					_ = rows.Close()
					return err
				}
				seqs = append(seqs, seq)
			}
			_ = rows.Close()
			if len(seqs) == 0 {
				return nil
			}
			for _, seq := range seqs {
				if _, err := tx.Exec(`DELETE FROM deliveries WHERE seq = ?`, seq); err != nil {
					return err
				}
				if _, err := tx.Exec(`DELETE FROM messages WHERE seq = ?`, seq); err != nil {
					return err
				}
			}
			return nil
		})
		if err != nil {
			return err
		}
		if len(seqs) == 0 {
			return nil
		}
	}
}

func (a *App) purgeReceipts(ctx context.Context, nowMs int64) error {
	if a.lim.ReceiptRetentionDays <= 0 {
		return nil
	}
	cutoff := nowMs - int64(a.lim.ReceiptRetentionDays)*24*3600*1000
	for {
		var n int64
		err := a.db.Queue.Do(ctx, func(tx *sql.Tx) error {
			res, err := tx.Exec(`
DELETE FROM receipts WHERE receipt_id IN (
  SELECT receipt_id FROM receipts WHERE created_at < ? LIMIT ?
)`, cutoff, purgeRowBatch)
			if err != nil {
				return err
			}
			n, _ = res.RowsAffected()
			return nil
		})
		if err != nil {
			return err
		}
		if n == 0 {
			return nil
		}
	}
}

func (a *App) purgeSendKeys(ctx context.Context, nowMs int64) error {
	if a.lim.IdempotencyHours <= 0 {
		return nil
	}
	cutoff := nowMs - int64(a.lim.IdempotencyHours)*3600*1000
	for {
		var n int64
		err := a.db.Queue.Do(ctx, func(tx *sql.Tx) error {
			res, err := tx.Exec(`
DELETE FROM send_keys WHERE rowid IN (
  SELECT sk.rowid FROM send_keys sk
  WHERE sk.created_at < ?
    AND NOT EXISTS (
      SELECT 1 FROM messages m WHERE m.sender_id = sk.sender_id AND m.id = sk.msg_id
    )
  LIMIT ?
)`, cutoff, purgeRowBatch)
			if err != nil {
				return err
			}
			n, _ = res.RowsAffected()
			return nil
		})
		if err != nil {
			return err
		}
		if n == 0 {
			return nil
		}
	}
}

// finalizeStuckDispatchedTx 收尾「dispatched 且已无 pending 投递」的消息（C-04 兜底）。
func finalizeStuckDispatchedTx(tx *sql.Tx, nowMs int64, recordDays int) error {
	rows, err := tx.Query(`
SELECT seq FROM messages
WHERE state = ?
  AND NOT EXISTS (
    SELECT 1 FROM deliveries d WHERE d.seq = messages.seq AND d.state = 'pending'
  )
LIMIT 500`, StateDispatched)
	if err != nil {
		return err
	}
	var seqs []int64
	for rows.Next() {
		var seq int64
		if err := rows.Scan(&seq); err != nil {
			_ = rows.Close()
			return err
		}
		seqs = append(seqs, seq)
	}
	_ = rows.Close()
	if err := rows.Err(); err != nil {
		return err
	}
	for _, seq := range seqs {
		if err := TryFinalizeTx(tx, seq, nowMs, recordDays); err != nil {
			return err
		}
	}
	return nil
}