orchestra/internal/store/store.go

package store

import (
	"bufio"
	"encoding/json"
	"errors"
	"fmt"
	"orchestra/internal/domain"
	"os"
	"path/filepath"
	"sync"
	"time"
)

type Store struct {
	mu       sync.Mutex
	path     string
	cas      string
	events   []domain.Event
	tasks    map[string]domain.Task
	external map[string]string
	snapshot string
}

func Open(dir string) (*Store, error) {
	if err := os.MkdirAll(dir, 0755); err != nil {
		return nil, err
	}
	s := &Store{path: filepath.Join(dir, "events.jsonl"), cas: filepath.Join(dir, "cas"), snapshot: filepath.Join(dir, "snapshot.json"), tasks: map[string]domain.Task{}, external: map[string]string{}}
	if err := os.MkdirAll(s.cas, 0755); err != nil {
		return nil, err
	}
	f, err := os.Open(s.path)
	if os.IsNotExist(err) {
		return s, nil
	}
	if err != nil {
		return nil, err
	}
	defer f.Close()
	sc := bufio.NewScanner(f)
	for sc.Scan() {
		var e domain.Event
		if err := json.Unmarshal(sc.Bytes(), &e); err == nil {
			s.events = append(s.events, e)
			if err := s.apply(e); err != nil {
				return nil, err
			}
		} else {
			return nil, err
		}
	}
	return s, sc.Err()
}
func (s *Store) apply(e domain.Event) error {
	var p map[string]any
	if err := json.Unmarshal(e.Payload, &p); err != nil {
		return err
	}
	t := s.tasks[e.TaskID]
	switch e.Type {
	case "TaskCreated":
		if err := domain.ValidateCreated(p); err != nil {
			return err
		}
		t = domain.Task{ID: e.TaskID, Source: p["source"].(string), ExternalID: p["external_id"].(string), Project: p["project"].(string), State: domain.StateQueued}
		if v, ok := p["capability"].([]any); ok {
			for _, x := range v {
				if z, ok := x.(string); ok {
					t.Capability = append(t.Capability, z)
				}
			}
		}
		if v, ok := p["title"].(string); ok {
			t.Title = v
		}
		s.external[t.Source+"\x00"+t.ExternalID] = t.ID
	case "TaskLeased":
		t.State = domain.StateLeased
		t.Lease = &domain.Lease{HarnessID: p["harness_id"].(string), Until: time.Unix(0, int64(p["until_ns"].(float64)))}
	case "TaskReleased":
		t.State = domain.StateQueued
		t.Lease = nil
	case "TaskCompleted":
		t.State = domain.StateCompleted
		t.Lease = nil
	case "TaskFailed":
		t.State = domain.StateFailed
		t.Lease = nil
	case "TaskBlocked":
		t.State = domain.StateBlocked
		t.Lease = nil
	case "TaskAmended":
		for k, v := range p {
			if k == "title" {
				t.Title, v = v.(string)
			}
		}
	}
	t.Version = e.Version
	s.tasks[e.TaskID] = t
	return nil
}
func (s *Store) Append(e domain.Event) error {
	s.mu.Lock()
	defer s.mu.Unlock()
	if err := domain.ValidateEvent(e); err != nil {
		return err
	}
	if e.At.IsZero() {
		e.At = time.Now().UTC()
	}
	if e.Seq == 0 {
		e.Seq = uint64(len(s.events) + 1)
	}
	if e.Type == "TaskCreated" {
		var p map[string]any
		if err := json.Unmarshal(e.Payload, &p); err != nil {
			return err
		}
		if id := s.external[p["source"].(string)+"\x00"+p["external_id"].(string)]; id != "" {
			return nil
		}
	}
	if t, ok := s.tasks[e.TaskID]; ok && e.Version != t.Version+1 {
		return domain.ErrConflict
	}
	if _, ok := s.tasks[e.TaskID]; !ok && e.Type != "TaskCreated" {
		return domain.ErrNotFound
	}
	if e.Type != "TaskCreated" && (e.Type == "TaskCompleted" || e.Type == "TaskBlocked" || e.Type == "TaskReleased") {
		var p map[string]any
		_ = json.Unmarshal(e.Payload, &p)
		for _, k := range []string{"handoff_ref", "report_ref"} {
			if ref, ok := p[k].(string); ok {
				if _, err := os.Stat(filepath.Join(s.cas, ref)); err != nil {
					return fmt.Errorf("%w: missing artifact %s", domain.ErrInvalid, ref)
				}
			}
		}
	}
	if err := s.apply(e); err != nil {
		return err
	}
	f, err := os.OpenFile(s.path, os.O_CREATE|os.O_APPEND|os.O_WRONLY, 0644)
	if err != nil {
		return err
	}
	defer f.Close()
	b, _ := json.Marshal(e)
	if _, err = f.Write(append(b, '\n')); err != nil {
		return err
	}
	if err = f.Sync(); err != nil {
		return err
	}
	s.events = append(s.events, e)
	if err := s.writeSnapshot(); err != nil {
		return err
	}
	return nil
}
func (s *Store) writeSnapshot() error {
	tasks := make([]domain.Task, 0, len(s.tasks))
	for _, t := range s.tasks {
		tasks = append(tasks, t)
	}
	b, err := json.Marshal(struct {
		Seq   uint64        `json:"seq"`
		Tasks []domain.Task `json:"tasks"`
	}{uint64(len(s.events)), tasks})
	if err != nil {
		return err
	}
	tmp := s.snapshot + ".tmp"
	if err = os.WriteFile(tmp, b, 0644); err != nil {
		return err
	}
	return os.Rename(tmp, s.snapshot)
}
func (s *Store) Tasks() []domain.Task {
	s.mu.Lock()
	defer s.mu.Unlock()
	out := make([]domain.Task, 0, len(s.tasks))
	for _, t := range s.tasks {
		out = append(out, t)
	}
	return out
}
func (s *Store) Events(since uint64) []domain.Event {
	s.mu.Lock()
	defer s.mu.Unlock()
	var out []domain.Event
	for _, e := range s.events {
		if e.Seq > since {
			out = append(out, e)
		}
	}
	return out
}
func (s *Store) PutArtifact(b []byte) (string, error) {
	h := domain.Hash(b)
	p := filepath.Join(s.cas, h)
	if _, err := os.Stat(p); errors.Is(err, os.ErrNotExist) {
		if err = os.WriteFile(p, b, 0644); err != nil {
			return "", err
		}
	}
	return h, nil
}

// Artifact returns a CAS artifact after verifying its content address.
func (s *Store) Artifact(ref string) ([]byte, error) {
	if len(ref) != 64 {
		return nil, fmt.Errorf("%w: invalid artifact reference", domain.ErrInvalid)
	}
	b, err := os.ReadFile(filepath.Join(s.cas, ref))
	if err != nil {
		return nil, err
	}
	if domain.Hash(b) != ref {
		return nil, fmt.Errorf("%w: corrupt artifact", domain.ErrInvalid)
	}
	return b, nil
}

func (s *Store) Task(id string) (domain.Task, bool) {
	s.mu.Lock()
	defer s.mu.Unlock()
	t, ok := s.tasks[id]
	return t, ok
}

func (s *Store) Lease(id, harness string, ttl time.Duration) (domain.Event, error) {
	t, ok := s.Task(id)
	if !ok {
		return domain.Event{}, domain.ErrNotFound
	}
	if t.State != domain.StateQueued {
		return domain.Event{}, domain.ErrConflict
	}
	p, _ := json.Marshal(map[string]any{"harness_id": harness, "until_ns": time.Now().Add(ttl).UnixNano()})
	e := domain.Event{ID: id, Type: "TaskLeased", TaskID: id, Version: t.Version + 1, Payload: p}
	return e, s.Append(e)
}

func (s *Store) ExpireLeases(now time.Time) ([]domain.Event, error) {
	var out []domain.Event
	for _, t := range s.Tasks() {
		if t.State == domain.StateLeased && t.Lease != nil && !t.Lease.Until.After(now) {
			p, _ := json.Marshal(map[string]any{"reason": "lease_expired", "harness_id": t.Lease.HarnessID})
			e := domain.Event{ID: t.ID, Type: "TaskReleased", TaskID: t.ID, Version: t.Version + 1, Payload: p}
			if err := s.Append(e); err != nil {
				return out, err
			}
			out = append(out, e)
		}
	}
	return out, nil
}