rvbox/internal/server/store/recovery.go

package store

import (
	"context"
	"database/sql"
	"errors"
	"fmt"
	"os"
	"path/filepath"
)

var (
	ErrIntegrityCheck       = errors.New("SQLite integrity check failed")
	ErrQuotaCounterMismatch = errors.New("quota counter mismatch")
)

type RecoveryReport struct {
	SegmentsChecked    uint64
	RecordsChecked     uint64
	TailBytesTruncated uint64
	OrphansRemoved     uint64
}

func (store *Store) QuickCheck(ctx context.Context) error {
	store.writeMu.Lock()
	defer store.writeMu.Unlock()
	database, err := store.openDatabase()
	if err != nil {
		return err
	}
	return store.quickCheckDatabase(ctx, database)
}

func (store *Store) RecoverCommandSegments(ctx context.Context) (RecoveryReport, error) {
	var report RecoveryReport
	store.writeMu.Lock()
	defer store.writeMu.Unlock()
	database, err := store.openDatabase()
	if err != nil {
		return report, err
	}
	if err := store.quickCheckDatabase(ctx, database); err != nil {
		return report, err
	}
	rows, err := database.QueryContext(ctx, `SELECT issue_uuid, ordinal, path, committed_end_offset,
min_event_seq, max_event_seq, stream_mix, compressed_bytes, raw_bytes
FROM output_segments ORDER BY issue_uuid, ordinal`)
	if err != nil {
		return report, err
	}
	type segmentMetadata struct {
		owner                                                                 [16]byte
		ordinal                                                               uint32
		name                                                                  string
		committed, minSequence, maxSequence, streamMix, storedBytes, rawBytes uint64
	}
	var metadata []segmentMetadata
	for rows.Next() {
		var current segmentMetadata
		var ownerBytes []byte
		if err := rows.Scan(&ownerBytes, &current.ordinal, &current.name, &current.committed, &current.minSequence, &current.maxSequence, &current.streamMix, &current.storedBytes, &current.rawBytes); err != nil {
			_ = rows.Close()
			return report, err
		}
		if len(ownerBytes) != 16 {
			_ = rows.Close()
			return report, fmt.Errorf("%w: invalid segment owner", ErrCommittedRangeMissing)
		}
		copy(current.owner[:], ownerBytes)
		metadata = append(metadata, current)
	}
	if err := rows.Close(); err != nil {
		return report, err
	}
	referenced := map[string]bool{}
	for _, current := range metadata {
		nameOwner, nameOrdinal, err := parseSegmentName(current.name)
		if err != nil || nameOwner != current.owner || nameOrdinal != current.ordinal {
			return report, ErrUnsafeSegmentReference
		}
		recovery, err := RecoverSegment(store.segmentDirectory(), current.name, current.committed, DefaultSegmentLimits())
		if err != nil {
			return report, err
		}
		if err := validateRecoveredMetadata(ctx, database, current.owner, current.ordinal, recovery, current.minSequence, current.maxSequence, current.streamMix, current.storedBytes, current.rawBytes); err != nil {
			return report, err
		}
		referenced[current.name] = true
		report.SegmentsChecked++
		report.RecordsChecked += uint64(len(recovery.Records))
		report.TailBytesTruncated += recovery.TruncatedTail
	}
	entries, err := os.ReadDir(store.segmentDirectory())
	if err != nil {
		return report, err
	}
	directoryChanged := false
	for _, entry := range entries {
		name := entry.Name()
		if referenced[name] {
			continue
		}
		if _, _, err := parseSegmentName(name); err != nil {
			return report, ErrUnsafeSegmentReference
		}
		path := filepath.Join(store.segmentDirectory(), name)
		info, err := os.Lstat(path)
		if err != nil {
			return report, err
		}
		if !info.Mode().IsRegular() || info.Mode().Perm()&0o077 != 0 {
			return report, ErrUnsafeSegmentReference
		}
		if err := os.Remove(path); err != nil {
			return report, err
		}
		directoryChanged = true
		report.OrphansRemoved++
	}
	if directoryChanged {
		if err := syncDirectory(store.segmentDirectory()); err != nil {
			return report, err
		}
	}
	if err := checkQuotaCounters(ctx, database); err != nil {
		return report, err
	}
	if _, err := database.ExecContext(ctx, `UPDATE output_segments SET sealed = 1 WHERE sealed = 0`); err != nil {
		return report, err
	}
	return report, nil
}

func checkQuotaCounters(ctx context.Context, database *sql.DB) error {
	var mismatchedClients int
	if err := database.QueryRowContext(ctx, `SELECT count(*) FROM clients WHERE charged_bytes !=
COALESCE((SELECT sum(commands.charged_bytes) FROM commands WHERE commands.client_id = clients.client_id), 0)`).Scan(&mismatchedClients); err != nil {
		return err
	}
	var storedServer, computedServer uint64
	if err := database.QueryRowContext(ctx, `SELECT command_charged_bytes,
COALESCE((SELECT sum(charged_bytes) FROM commands), 0) FROM storage_counters WHERE singleton = 1`).Scan(&storedServer, &computedServer); err != nil {
		return err
	}
	var invalidCommands int
	if err := database.QueryRowContext(ctx, `SELECT count(*) FROM commands WHERE output_charged_bytes > charged_bytes OR charge_version != ?`, ChargeFormulaVersion).Scan(&invalidCommands); err != nil {
		return err
	}
	if mismatchedClients != 0 || storedServer != computedServer || invalidCommands != 0 {
		return ErrQuotaCounterMismatch
	}
	return nil
}

func (store *Store) quickCheckDatabase(ctx context.Context, database *sql.DB) error {
	rows, err := database.QueryContext(ctx, `PRAGMA quick_check(100)`)
	if err != nil {
		return fmt.Errorf("%w: query failed", ErrIntegrityCheck)
	}
	defer rows.Close()
	count := 0
	for rows.Next() {
		var result string
		if err := rows.Scan(&result); err != nil {
			return fmt.Errorf("%w: result unreadable", ErrIntegrityCheck)
		}
		count++
		if result != "ok" {
			return fmt.Errorf("%w: database reported an inconsistency", ErrIntegrityCheck)
		}
	}
	if err := rows.Err(); err != nil || count != 1 {
		return fmt.Errorf("%w: incomplete result", ErrIntegrityCheck)
	}
	return nil
}

func validateRecoveredMetadata(ctx context.Context, database *sql.DB, owner [16]byte, ordinal uint32, recovery SegmentRecovery, minSequence, maxSequence, streamMix, storedBytes, rawBytes uint64) error {
	rows, err := database.QueryContext(ctx, `SELECT event_seq, segment_record_offset, segment_record_length,
compression, raw_bytes, stored_bytes FROM command_events
WHERE issue_uuid = ? AND segment_ordinal = ? ORDER BY segment_record_offset`, owner[:], ordinal)
	if err != nil {
		return err
	}
	defer rows.Close()
	var index int
	var computedStored, computedRaw, computedStreams uint64
	for rows.Next() {
		if index >= len(recovery.Records) {
			return ErrCommittedRangeMissing
		}
		var sequence, offset, length, compression, raw, stored uint64
		if err := rows.Scan(&sequence, &offset, &length, &compression, &raw, &stored); err != nil {
			return err
		}
		recovered := recovery.Records[index]
		if recovered.Offset != offset || recovered.Length != length || recovered.Record.OwnerUUID != owner || recovered.Record.Sequence != sequence || uint64(recovered.Record.Compression) != compression || recovered.Record.RawLength != raw || uint64(len(recovered.Record.Payload)) != stored {
			return ErrCommittedRangeMissing
		}
		if recovered.Record.Stream > 0 {
			computedStreams |= uint64(1) << (recovered.Record.Stream - 1)
		}
		computedStored += stored
		computedRaw += raw
		index++
	}
	if err := rows.Err(); err != nil {
		return err
	}
	if index != len(recovery.Records) || index == 0 || recovery.Records[0].Record.Sequence != minSequence || recovery.Records[index-1].Record.Sequence != maxSequence || computedStreams != streamMix || computedStored != storedBytes || computedRaw != rawBytes {
		return ErrCommittedRangeMissing
	}
	return nil
}