compaction.go

package storage

import (
	"fmt"
	"os"
	"sort"
	"sync"
	"time"
)

type CompactionConfig struct {
	// MinDirtyRatio is the minimum ratio of "dirty" (duplicate key) records
	// before compaction is triggered (0.0-1.0). Default: 0.5
	MinDirtyRatio float64

	// MinCompactionInterval is the minimum time between compactions.
	MinCompactionInterval time.Duration

	// Enabled controls whether compaction runs.
	Enabled bool
}

// DefaultCompactionConfig returns sensible defaults.
func DefaultCompactionConfig() CompactionConfig {
	return CompactionConfig{
		MinDirtyRatio:         0.5,
		MinCompactionInterval: 5 * time.Minute,
		Enabled:               false, // Opt-in
	}
}

// KeyIndex tracks the latest offset for each key.
// Used during compaction to determine which records to keep.
type KeyIndex struct {
	mu      sync.RWMutex
	keys    map[string]uint64 // key -> latest offset
	offsets map[uint64]string // offset -> key (for reverse lookup)
}

func NewKeyIndex() *KeyIndex {
	return &KeyIndex{
		keys:    make(map[string]uint64),
		offsets: make(map[uint64]string),
	}
}

// Update records that a key was seen at the given offset.
// Only updates if this is a newer offset for the key.
func (ki *KeyIndex) Update(key []byte, offset uint64) {
	ki.mu.Lock()
	defer ki.mu.Unlock()

	keyStr := string(key)
	if existing, ok := ki.keys[keyStr]; ok {
		if offset > existing {
			// Remove old offset mapping
			delete(ki.offsets, existing)
			ki.keys[keyStr] = offset
			ki.offsets[offset] = keyStr
		}
	} else {
		ki.keys[keyStr] = offset
		ki.offsets[offset] = keyStr
	}
}

// IsLatest returns true if the given offset is the latest occurrence of the key.
func (ki *KeyIndex) IsLatest(key []byte, offset uint64) bool {
	ki.mu.RLock()
	defer ki.mu.RUnlock()

	latest, exists := ki.keys[string(key)]
	if !exists {
		return true // Unknown key, keep it
	}
	return latest == offset
}

func (ki *KeyIndex) GetLatest(key []byte) (uint64, bool) {
	ki.mu.RLock()
	defer ki.mu.RUnlock()
	offset, ok := ki.keys[string(key)]
	return offset, ok
}

// Size returns the number of unique keys.
func (ki *KeyIndex) Size() int {
	ki.mu.RLock()
	defer ki.mu.RUnlock()
	return len(ki.keys)
}

// UniqueKeys returns all unique keys.
func (ki *KeyIndex) UniqueKeys() []string {
	ki.mu.RLock()
	defer ki.mu.RUnlock()

	keys := make([]string, 0, len(ki.keys))
	for k := range ki.keys {
		keys = append(keys, k)
	}
	return keys
}

type CompactionStats struct {
	RecordsScanned  int
	RecordsRetained int
	RecordsRemoved  int
	BytesBefore     uint64
	BytesAfter      uint64
	Duration        time.Duration
}

// Compactor handles log compaction for a single log.
type Compactor struct {
	mu          sync.Mutex
	log         *Log
	config      CompactionConfig
	lastCompact time.Time
	stats       CompactionStats
}

func NewCompactor(log *Log, config CompactionConfig) *Compactor {
	return &Compactor{
		log:    log,
		config: config,
	}
}

// ShouldCompact returns true if compaction should run.
func (c *Compactor) ShouldCompact() bool {
	c.mu.Lock()
	defer c.mu.Unlock()

	if !c.config.Enabled {
		return false
	}

	if time.Since(c.lastCompact) < c.config.MinCompactionInterval {
		return false
	}

	// This is a lightweight check that just counts keys
	keyIndex := NewKeyIndex()
	duplicates := 0
	total := 0

	c.log.Scan(func(offset uint64, data []byte) error {
		total++
		record, err := UnmarshalRecord(data)
		if err != nil || len(record.Key) == 0 {
			return nil // Skip keyless records
		}

		if _, exists := keyIndex.GetLatest(record.Key); exists {
			duplicates++
		}
		keyIndex.Update(record.Key, offset)
		return nil
	})

	if total == 0 {
		return false
	}

	dirtyRatio := float64(duplicates) / float64(total)
	return dirtyRatio >= c.config.MinDirtyRatio
}

// Compact performs log compaction, retaining only the latest value per key.
// Records without keys are always retained.
func (c *Compactor) Compact() (*CompactionStats, error) {
	c.mu.Lock()
	defer c.mu.Unlock()

	startTime := time.Now()

	// Phase 1: Build key index by scanning all records
	keyIndex := NewKeyIndex()
	var bytesBeforeTotal uint64

	err := c.log.Scan(func(offset uint64, data []byte) error {
		bytesBeforeTotal += uint64(len(data))
		record, err := UnmarshalRecord(data)
		if err != nil || len(record.Key) == 0 {
			return nil // Records without keys indexed at their offset
		}
		keyIndex.Update(record.Key, offset)
		return nil
	})
	if err != nil {
		return nil, fmt.Errorf("scan for key index: %w", err)
	}

	// Phase 2: Collect records to keep
	type keepRecord struct {
		offset uint64
		data   []byte
	}
	var recordsToKeep []keepRecord
	recordsScanned := 0
	var bytesAfter uint64

	err = c.log.Scan(func(offset uint64, data []byte) error {
		recordsScanned++
		record, err := UnmarshalRecord(data)
		if err != nil || len(record.Key) == 0 {
			// Keep keyless records
			recordsToKeep = append(recordsToKeep, keepRecord{offset, data})
			bytesAfter += uint64(len(data))
			return nil
		}

		// Keep only if this is the latest for the key
		if keyIndex.IsLatest(record.Key, offset) {
			recordsToKeep = append(recordsToKeep, keepRecord{offset, data})
			bytesAfter += uint64(len(data))
		}
		return nil
	})
	if err != nil {
		return nil, fmt.Errorf("scan for records to keep: %w", err)
	}

	sort.Slice(recordsToKeep, func(i, j int) bool {
		return recordsToKeep[i].offset < recordsToKeep[j].offset
	})

	// Phase 3: Create new compacted log
	compactedDir := c.log.Dir() + ".compacting"
	if err := os.MkdirAll(compactedDir, 0755); err != nil {
		return nil, fmt.Errorf("create compacted dir: %w", err)
	}

	compactedLog, err := NewLog(Config{
		Dir:             compactedDir,
		MaxSegmentBytes: c.log.config.MaxSegmentBytes,
		SyncWrites:      c.log.config.SyncWrites,
	})
	if err != nil {
		os.RemoveAll(compactedDir)
		return nil, fmt.Errorf("create compacted log: %w", err)
	}

	// Records keep their original offsets so committed consumer offsets stay
	// valid; gaps left by removed records start a new segment.
	for _, rec := range recordsToKeep {
		if _, err := compactedLog.AppendAt(rec.offset, rec.data); err != nil {
			compactedLog.Close()
			os.RemoveAll(compactedDir)
			return nil, fmt.Errorf("write to compacted log: %w", err)
		}
	}

	if err := compactedLog.Close(); err != nil {
		os.RemoveAll(compactedDir)
		return nil, fmt.Errorf("close compacted log: %w", err)
	}

	// Phase 4: Atomic swap
	originalDir := c.log.Dir()
	backupDir := originalDir + ".backup"

	if err := c.log.Close(); err != nil {
		os.RemoveAll(compactedDir)
		return nil, fmt.Errorf("close original log: %w", err)
	}

	// Rename original to backup
	if err := os.Rename(originalDir, backupDir); err != nil {
		os.RemoveAll(compactedDir)
		return nil, fmt.Errorf("backup original: %w", err)
	}

	// Rename compacted to original
	if err := os.Rename(compactedDir, originalDir); err != nil {
		// Try to restore backup
		os.Rename(backupDir, originalDir)
		return nil, fmt.Errorf("swap compacted: %w", err)
	}

	os.RemoveAll(backupDir)

	// Reopen log
	newLog, err := NewLog(c.log.config)
	if err != nil {
		return nil, fmt.Errorf("reopen log: %w", err)
	}
	c.log = newLog

	c.lastCompact = time.Now()

	stats := &CompactionStats{
		RecordsScanned:  recordsScanned,
		RecordsRetained: len(recordsToKeep),
		RecordsRemoved:  recordsScanned - len(recordsToKeep),
		BytesBefore:     bytesBeforeTotal,
		BytesAfter:      bytesAfter,
		Duration:        time.Since(startTime),
	}
	c.stats = *stats

	return stats, nil
}

func (c *Compactor) GetStats() CompactionStats {
	c.mu.Lock()
	defer c.mu.Unlock()
	return c.stats
}

// SetEnabled enables or disables compaction.
func (c *Compactor) SetEnabled(enabled bool) {
	c.mu.Lock()
	defer c.mu.Unlock()
	c.config.Enabled = enabled
}

func (c *Compactor) GetLog() *Log {
	c.mu.Lock()
	defer c.mu.Unlock()
	return c.log
}