monlet/agent/internal/spool/spool.go

package spool

import (
	"encoding/json"
	"fmt"
	"os"
	"path/filepath"
	"sort"
	"strconv"
	"strings"
	"sync"

	"github.com/monlet/agent/internal/scheduler"
)

const (
	DefaultMaxEvents = 10_000
	DefaultMaxBytes  = 50 * 1024 * 1024
	fileExt          = ".json"
)

type Spool struct {
	dir       string
	maxEvents int
	maxBytes  int64

	mu      sync.Mutex
	nextSeq uint64
	files   []entry // sorted ascending by seq
	bytes   int64
}

type entry struct {
	seq  uint64
	size int64
	name string
}

type Item struct {
	Seq   uint64
	Event scheduler.Event
}

func Open(dir string, maxEvents int, maxBytes int64) (*Spool, error) {
	if err := os.MkdirAll(dir, 0o755); err != nil {
		return nil, err
	}
	if maxEvents <= 0 {
		maxEvents = DefaultMaxEvents
	}
	if maxBytes <= 0 {
		maxBytes = DefaultMaxBytes
	}
	s := &Spool{dir: dir, maxEvents: maxEvents, maxBytes: maxBytes}
	if err := s.rescan(); err != nil {
		return nil, err
	}
	s.mu.Lock()
	s.dropOldestLocked()
	s.mu.Unlock()
	return s, nil
}

func (s *Spool) rescan() error {
	ents, err := os.ReadDir(s.dir)
	if err != nil {
		return err
	}
	s.files = s.files[:0]
	s.bytes = 0
	for _, e := range ents {
		if e.IsDir() || !strings.HasSuffix(e.Name(), fileExt) {
			continue
		}
		base := strings.TrimSuffix(e.Name(), fileExt)
		seq, err := strconv.ParseUint(base, 10, 64)
		if err != nil {
			continue
		}
		info, err := e.Info()
		if err != nil {
			return err
		}
		s.files = append(s.files, entry{seq: seq, size: info.Size(), name: e.Name()})
		s.bytes += info.Size()
		if seq >= s.nextSeq {
			s.nextSeq = seq + 1
		}
	}
	sort.Slice(s.files, func(i, j int) bool { return s.files[i].seq < s.files[j].seq })
	return nil
}

func (s *Spool) Len() int {
	s.mu.Lock()
	defer s.mu.Unlock()
	return len(s.files)
}

func (s *Spool) Bytes() int64 {
	s.mu.Lock()
	defer s.mu.Unlock()
	return s.bytes
}

func (s *Spool) Enqueue(ev scheduler.Event) error {
	data, err := json.Marshal(ev)
	if err != nil {
		return err
	}
	s.mu.Lock()
	defer s.mu.Unlock()

	seq := s.nextSeq
	s.nextSeq++
	name := fmt.Sprintf("%014d%s", seq, fileExt)
	tmp := filepath.Join(s.dir, name+".tmp")
	final := filepath.Join(s.dir, name)
	if err := os.WriteFile(tmp, data, 0o600); err != nil {
		return err
	}
	if err := os.Rename(tmp, final); err != nil {
		_ = os.Remove(tmp)
		return err
	}
	size := int64(len(data))
	s.files = append(s.files, entry{seq: seq, size: size, name: name})
	s.bytes += size
	s.dropOldestLocked()
	return nil
}

func (s *Spool) dropOldestLocked() {
	for len(s.files) > s.maxEvents || s.bytes > s.maxBytes {
		if len(s.files) == 0 {
			return
		}
		old := s.files[0]
		_ = os.Remove(filepath.Join(s.dir, old.name))
		s.bytes -= old.size
		s.files = s.files[1:]
	}
}

// Peek returns up to limit oldest items. Does not remove them.
func (s *Spool) Peek(limit int) ([]Item, error) {
	s.mu.Lock()
	defer s.mu.Unlock()
	if limit > len(s.files) {
		limit = len(s.files)
	}
	out := make([]Item, 0, limit)
	corrupt := make(map[uint64]struct{})
	for i := 0; i < limit; i++ {
		e := s.files[i]
		b, err := os.ReadFile(filepath.Join(s.dir, e.name))
		if err != nil {
			if os.IsNotExist(err) {
				corrupt[e.seq] = struct{}{}
				continue
			}
			s.dropCorruptLocked(corrupt)
			return out, err
		}
		var ev scheduler.Event
		if err := json.Unmarshal(b, &ev); err != nil {
			_ = os.Remove(filepath.Join(s.dir, e.name))
			corrupt[e.seq] = struct{}{}
			continue
		}
		out = append(out, Item{Seq: e.seq, Event: ev})
	}
	s.dropCorruptLocked(corrupt)
	return out, nil
}

func (s *Spool) dropCorruptLocked(seqs map[uint64]struct{}) {
	if len(seqs) == 0 {
		return
	}
	kept := s.files[:0]
	for _, e := range s.files {
		if _, bad := seqs[e.seq]; bad {
			s.bytes -= e.size
			continue
		}
		kept = append(kept, e)
	}
	s.files = kept
}

// Commit removes items with the given sequence numbers.
func (s *Spool) Commit(seqs []uint64) {
	if len(seqs) == 0 {
		return
	}
	want := make(map[uint64]struct{}, len(seqs))
	for _, x := range seqs {
		want[x] = struct{}{}
	}
	s.mu.Lock()
	defer s.mu.Unlock()
	kept := s.files[:0]
	for _, e := range s.files {
		if _, ok := want[e.seq]; ok {
			_ = os.Remove(filepath.Join(s.dir, e.name))
			s.bytes -= e.size
			continue
		}
		kept = append(kept, e)
	}
	s.files = kept
}