init monlet repo with stage 0-2 (docs, contract, agent MVP)
This commit is contained in:
Stanislav Rossovskii
213
agent/internal/spool/spool.go
Normal file
213
agent/internal/spool/spool.go
Normal file
@@ -0,0 +1,213 @@
|
||||
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
|
||||
}
|
||||
129
agent/internal/spool/spool_test.go
Normal file
129
agent/internal/spool/spool_test.go
Normal file
@@ -0,0 +1,129 @@
|
||||
package spool
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"strings"
|
||||
"testing"
|
||||
|
||||
"github.com/monlet/agent/internal/scheduler"
|
||||
)
|
||||
|
||||
func mkEvent(id string) scheduler.Event {
|
||||
return scheduler.Event{EventID: id, CheckID: "c1", Status: "ok"}
|
||||
}
|
||||
|
||||
func TestEnqueuePeekCommit(t *testing.T) {
|
||||
s, err := Open(t.TempDir(), 10, 1<<20)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
for i := 0; i < 3; i++ {
|
||||
if err := s.Enqueue(mkEvent(fmt.Sprintf("e%d", i))); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
}
|
||||
items, err := s.Peek(10)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if len(items) != 3 {
|
||||
t.Fatalf("got %d items", len(items))
|
||||
}
|
||||
if items[0].Event.EventID != "e0" {
|
||||
t.Fatalf("FIFO violated: %q", items[0].Event.EventID)
|
||||
}
|
||||
s.Commit([]uint64{items[0].Seq, items[1].Seq})
|
||||
if s.Len() != 1 {
|
||||
t.Fatalf("len=%d", s.Len())
|
||||
}
|
||||
}
|
||||
|
||||
func TestDropOldestByCount(t *testing.T) {
|
||||
s, err := Open(t.TempDir(), 3, 1<<20)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
for i := 0; i < 5; i++ {
|
||||
_ = s.Enqueue(mkEvent(fmt.Sprintf("e%d", i)))
|
||||
}
|
||||
if s.Len() != 3 {
|
||||
t.Fatalf("len=%d", s.Len())
|
||||
}
|
||||
items, _ := s.Peek(10)
|
||||
if items[0].Event.EventID != "e2" {
|
||||
t.Fatalf("expected oldest dropped, got %q", items[0].Event.EventID)
|
||||
}
|
||||
}
|
||||
|
||||
func TestDropOldestByBytes(t *testing.T) {
|
||||
big := strings.Repeat("x", 1000)
|
||||
s, err := Open(t.TempDir(), 1_000_000, 3000)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
for i := 0; i < 10; i++ {
|
||||
ev := mkEvent(fmt.Sprintf("e%d", i))
|
||||
ev.Output = big
|
||||
_ = s.Enqueue(ev)
|
||||
}
|
||||
if s.Bytes() > 3000 {
|
||||
t.Fatalf("bytes=%d > 3000", s.Bytes())
|
||||
}
|
||||
}
|
||||
|
||||
func TestReopenAppliesLimits(t *testing.T) {
|
||||
dir := t.TempDir()
|
||||
s, _ := Open(dir, 10, 1<<20)
|
||||
for i := 0; i < 5; i++ {
|
||||
_ = s.Enqueue(mkEvent(fmt.Sprintf("e%d", i)))
|
||||
}
|
||||
// Reopen with tighter limit: should drop oldest down to maxEvents.
|
||||
s2, err := Open(dir, 2, 1<<20)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if s2.Len() != 2 {
|
||||
t.Fatalf("expected limits applied on Open, len=%d", s2.Len())
|
||||
}
|
||||
items, _ := s2.Peek(10)
|
||||
if items[0].Event.EventID != "e3" {
|
||||
t.Fatalf("expected oldest dropped, got %q", items[0].Event.EventID)
|
||||
}
|
||||
}
|
||||
|
||||
func TestPeekSkipsCorruptAndCleansList(t *testing.T) {
|
||||
dir := t.TempDir()
|
||||
s, _ := Open(dir, 10, 1<<20)
|
||||
_ = s.Enqueue(mkEvent("e0"))
|
||||
_ = s.Enqueue(mkEvent("e1"))
|
||||
// Corrupt the first file on disk.
|
||||
items, _ := s.Peek(10)
|
||||
if err := os.WriteFile(filepath.Join(dir, fmt.Sprintf("%014d.json", items[0].Seq)), []byte("not-json"), 0o600); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
out, _ := s.Peek(10)
|
||||
if len(out) != 1 || out[0].Event.EventID != "e1" {
|
||||
t.Fatalf("got %+v", out)
|
||||
}
|
||||
if s.Len() != 1 {
|
||||
t.Fatalf("corrupt entry not pruned from list, len=%d", s.Len())
|
||||
}
|
||||
}
|
||||
|
||||
func TestReopenPreservesOrder(t *testing.T) {
|
||||
dir := t.TempDir()
|
||||
s, _ := Open(dir, 100, 1<<20)
|
||||
for i := 0; i < 3; i++ {
|
||||
_ = s.Enqueue(mkEvent(fmt.Sprintf("e%d", i)))
|
||||
}
|
||||
s2, err := Open(dir, 100, 1<<20)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
items, _ := s2.Peek(10)
|
||||
if len(items) != 3 || items[0].Event.EventID != "e0" || items[2].Event.EventID != "e2" {
|
||||
t.Fatalf("bad reopen: %+v", items)
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user