bugabinga.net

package store

import (
	"crypto/rand"
	"crypto/sha256"
	"encoding/hex"
	"encoding/json"
	"fmt"
	"os"
	"path/filepath"
	"syscall"
	"time"
)

type Phase string

const (
	PhasePreparing   Phase = "preparing"
	PhaseActivating  Phase = "activating"
	PhaseProbation   Phase = "probation"
	PhasePromoting   Phase = "promoting"
	PhasePromoted    Phase = "promoted"
	PhaseRollingBack Phase = "rolling-back"
	PhaseRolledBack  Phase = "rolled-back"
	PhaseFailed      Phase = "failed"
)

func (p Phase) Terminal() bool {
	return p == PhasePromoted || p == PhaseRolledBack || p == PhaseFailed
}

type Revision struct {
	Digest    string `json:"digest"`
	Reference string `json:"reference"`
	Unit      string `json:"unit"`
}

type State struct {
	Generation  uint64    `json:"generation"`
	OperationID string    `json:"operationID,omitempty"`
	Current     *Revision `json:"current,omitempty"`
	Rollback    *Revision `json:"rollback,omitempty"`
	UpdatedAt   time.Time `json:"updatedAt"`
}

type Operation struct {
	ID                 string     `json:"id"`
	Service            string     `json:"service"`
	Phase              Phase      `json:"phase"`
	ExpectedGeneration uint64     `json:"expectedGeneration"`
	Generation         uint64     `json:"generation,omitempty"`
	IdempotencyKey     string     `json:"idempotencyKey"`
	Candidate          Revision   `json:"candidate"`
	Rollback           *Revision  `json:"rollback,omitempty"`
	Restarts           int        `json:"restarts"`
	Reason             string     `json:"reason,omitempty"`
	StartedAt          time.Time  `json:"startedAt"`
	UpdatedAt          time.Time  `json:"updatedAt"`
	Deadline           time.Time  `json:"deadline"`
	EndedAt            *time.Time `json:"endedAt,omitempty"`
}

type Store struct {
	Root string
}

func (s Store) serviceDir(service string) string { return filepath.Join(s.Root, "services", service) }

func (s Store) Ensure(service string) error {
	for _, dir := range []string{"operations", "idempotency"} {
		if err := os.MkdirAll(filepath.Join(s.serviceDir(service), dir), 0o750); err != nil {
			return err
		}
	}
	return nil
}

// Lock takes the exclusive single-writer lock for one service.
func (s Store) Lock(service string) (func(), error) {
	if err := s.Ensure(service); err != nil {
		return nil, err
	}
	file, err := os.OpenFile(filepath.Join(s.serviceDir(service), "lock"), os.O_CREATE|os.O_RDWR, 0o640)
	if err != nil {
		return nil, err
	}
	if err := syscall.Flock(int(file.Fd()), syscall.LOCK_EX|syscall.LOCK_NB); err != nil {
		file.Close()
		return nil, fmt.Errorf("service %s is locked by another writer: %w", service, err)
	}
	return func() {
		_ = syscall.Flock(int(file.Fd()), syscall.LOCK_UN)
		_ = file.Close()
	}, nil
}

func (s Store) ReadState(service string) (State, error) {
	var state State
	data, err := os.ReadFile(filepath.Join(s.serviceDir(service), "state.json"))
	if os.IsNotExist(err) {
		return State{}, nil
	}
	if err != nil {
		return State{}, err
	}
	if err := json.Unmarshal(data, &state); err != nil {
		return State{}, fmt.Errorf("service %s state: %w", service, err)
	}
	return state, nil
}

func (s Store) WriteState(service string, state State) error {
	if err := s.Ensure(service); err != nil {
		return err
	}
	state.UpdatedAt = time.Now().UTC()
	return writeAtomic(filepath.Join(s.serviceDir(service), "state.json"), encode(state))
}

func (s Store) ReadOperation(service, id string) (Operation, error) {
	var operation Operation
	if !safeID(id) {
		return Operation{}, fmt.Errorf("invalid operation id %q", id)
	}
	data, err := os.ReadFile(filepath.Join(s.serviceDir(service), "operations", id+".json"))
	if err != nil {
		return Operation{}, err
	}
	if err := json.Unmarshal(data, &operation); err != nil {
		return Operation{}, err
	}
	return operation, nil
}

func (s Store) WriteOperation(operation Operation) error {
	if err := s.Ensure(operation.Service); err != nil {
		return err
	}
	if !safeID(operation.ID) {
		return fmt.Errorf("invalid operation id %q", operation.ID)
	}
	operation.UpdatedAt = time.Now().UTC()
	path := filepath.Join(s.serviceDir(operation.Service), "operations", operation.ID+".json")
	return writeAtomic(path, encode(operation))
}

func (s Store) OperationIDs(service string) ([]string, error) {
	entries, err := os.ReadDir(filepath.Join(s.serviceDir(service), "operations"))
	if os.IsNotExist(err) {
		return nil, nil
	}
	if err != nil {
		return nil, err
	}
	var ids []string
	for _, entry := range entries {
		name := entry.Name()
		if filepath.Ext(name) == ".json" {
			ids = append(ids, name[:len(name)-len(".json")])
		}
	}
	return ids, nil
}

func (s Store) LookupIdempotency(service, key string) (string, bool, error) {
	data, err := os.ReadFile(s.idempotencyPath(service, key))
	if os.IsNotExist(err) {
		return "", false, nil
	}
	if err != nil {
		return "", false, err
	}
	return string(data), true, nil
}

func (s Store) RecordIdempotency(service, key, id string) error {
	if err := s.Ensure(service); err != nil {
		return err
	}
	return writeAtomic(s.idempotencyPath(service, key), []byte(id))
}

func (s Store) idempotencyPath(service, key string) string {
	sum := sha256.Sum256([]byte(key))
	return filepath.Join(s.serviceDir(service), "idempotency", hex.EncodeToString(sum[:]))
}

func (s Store) AppendLog(service, id, line string) error {
	if !safeID(id) {
		return fmt.Errorf("invalid operation id %q", id)
	}
	if err := s.Ensure(service); err != nil {
		return err
	}
	path := filepath.Join(s.serviceDir(service), "operations", id+".log")
	file, err := os.OpenFile(path, os.O_CREATE|os.O_WRONLY|os.O_APPEND, 0o640)
	if err != nil {
		return err
	}
	stamped := time.Now().UTC().Format(time.RFC3339) + " " + line + "\n"
	if _, err := file.WriteString(stamped); err != nil {
		file.Close()
		return err
	}
	return file.Close()
}

func (s Store) ReadLog(service, id string) ([]byte, error) {
	if !safeID(id) {
		return nil, fmt.Errorf("invalid operation id %q", id)
	}
	data, err := os.ReadFile(filepath.Join(s.serviceDir(service), "operations", id+".log"))
	if os.IsNotExist(err) {
		return nil, nil
	}
	return data, err
}

func NewOperationID(now time.Time) string {
	var suffix [4]byte
	if _, err := rand.Read(suffix[:]); err != nil {
		suffix = [4]byte{}
	}
	return fmt.Sprintf("%020d-%s", now.UTC().UnixNano(), hex.EncodeToString(suffix[:]))
}

func safeID(id string) bool {
	if id == "" || len(id) > 64 {
		return false
	}
	for _, r := range id {
		if !(r >= '0' && r <= '9' || r >= 'a' && r <= 'f' || r == '-') {
			return false
		}
	}
	return true
}

// WriteFileAtomic publishes a complete file or leaves the previous one untouched.
func WriteFileAtomic(path string, data []byte) error { return writeAtomic(path, data) }

func encode(value any) []byte {
	data, err := json.MarshalIndent(value, "", "  ")
	if err != nil {
		panic(err)
	}
	return append(data, '\n')
}

func writeAtomic(path string, data []byte) error {
	directory := filepath.Dir(path)
	file, err := os.CreateTemp(directory, ".tmp-")
	if err != nil {
		return err
	}
	temporary := file.Name()
	defer os.Remove(temporary)
	if _, err := file.Write(data); err != nil {
		file.Close()
		return err
	}
	if err := file.Sync(); err != nil {
		file.Close()
		return err
	}
	if err := file.Close(); err != nil {
		return err
	}
	if err := os.Chmod(temporary, 0o640); err != nil {
		return err
	}
	if err := os.Rename(temporary, path); err != nil {
		return err
	}
	handle, err := os.Open(directory)
	if err != nil {
		return err
	}
	syncErr := handle.Sync()
	closeErr := handle.Close()
	if syncErr != nil {
		return syncErr
	}
	return closeErr
}