run/spool/log_spool.go

package spool

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

	"browser.local/run/protocol"
)

type LogSpool struct {
	dir        string
	statePath  string
	mu         *sync.Mutex
	watermarks map[string]logStreamWatermark
	inflight   map[string]bool
}

const maxAggregatedLogEntries = 128

type logStreamWatermark struct {
	Allocated    uint64 `json:"allocated"`
	Acknowledged uint64 `json:"acknowledged"`
	SourceOffset int64  `json:"sourceOffset,omitempty"`
}
type logSpoolState struct {
	Streams map[string]logStreamWatermark `json:"streams"`
}

type LogSourceCursor struct {
	StartOffset int64 `json:"startOffset"`
	EndOffset   int64 `json:"endOffset"`
}

type durableLogBatch struct {
	protocol.LogBatchIngestRequest
	SourceCursor *LogSourceCursor `json:"_sourceCursor,omitempty"`
}

type pendingLogSegment struct {
	path  string
	batch durableLogBatch
}

func NewLogSpool(dir string) (LogSpool, error) {
	if strings.TrimSpace(dir) == "" {
		return LogSpool{}, fmt.Errorf("spool directory is required")
	}
	logDir := filepath.Join(dir, "logs")
	if err := os.MkdirAll(logDir, 0o755); err != nil {
		return LogSpool{}, fmt.Errorf("create log spool: %w", err)
	}
	spool := LogSpool{dir: logDir, statePath: filepath.Join(dir, "log-watermarks.json"), mu: &sync.Mutex{}, watermarks: map[string]logStreamWatermark{}, inflight: map[string]bool{}}
	if err := spool.loadWatermarks(); err != nil {
		return LogSpool{}, err
	}
	if err := spool.restorePendingWatermarks(); err != nil {
		return LogSpool{}, err
	}
	return spool, nil
}

// restorePendingWatermarks supports upgrades from spools created before
// per-stream watermarks existed. Pending durable segments are still an
// allocation fact and must win over a remote progress lookup.
func (spool *LogSpool) restorePendingWatermarks() error {
	spool.mu.Lock()
	defer spool.mu.Unlock()
	pending, err := spool.pendingSegmentsLocked()
	if err != nil {
		return err
	}
	changed := false
	for _, segment := range pending {
		batch := segment.batch
		watermark := spool.watermarks[batch.LogStreamID]
		if batch.LastSeq > watermark.Allocated {
			watermark.Allocated = batch.LastSeq
			changed = true
		}
		if batch.SourceCursor != nil && batch.SourceCursor.EndOffset > watermark.SourceOffset {
			watermark.SourceOffset = batch.SourceCursor.EndOffset
			changed = true
		}
		spool.watermarks[batch.LogStreamID] = watermark
	}
	if changed {
		return spool.persistWatermarksLocked()
	}
	return nil
}

func (spool LogSpool) NextSequence(ctx context.Context, streamID string, recover func(context.Context, string) (uint64, error)) (uint64, error) {
	spool.mu.Lock()
	defer spool.mu.Unlock()
	watermark, known := spool.watermarks[streamID]
	if !known && recover != nil {
		latest, err := recover(ctx, streamID)
		if err != nil {
			return 0, err
		}
		watermark = logStreamWatermark{Allocated: latest, Acknowledged: latest}
	}
	watermark.Allocated++
	spool.watermarks[streamID] = watermark
	if err := spool.persistWatermarksLocked(); err != nil {
		return 0, err
	}
	return watermark.Allocated, nil
}

func (spool LogSpool) Enqueue(batch protocol.LogBatchIngestRequest) error {
	return spool.enqueue(batch, nil)
}

// EnqueueAggregated extends the newest compatible durable segment so callers
// do not create one upload request for every process-output line.
func (spool LogSpool) EnqueueAggregated(batch protocol.LogBatchIngestRequest, checksum func([]protocol.LogEntry) (string, error)) error {
	return spool.enqueue(batch, checksum)
}

// EnqueueNextAggregated allocates the next sequence and commits its batch
// under one spool lock. The durable segment is the source of truth; startup
// restores its watermark if the separate watermark snapshot was interrupted.
func (spool LogSpool) EnqueueNextAggregated(ctx context.Context, batch protocol.LogBatchIngestRequest, cursor *LogSourceCursor, recover func(context.Context, string) (uint64, error), checksum func([]protocol.LogEntry) (string, error)) (uint64, bool, error) {
	if len(batch.Entries) != 1 {
		return 0, false, fmt.Errorf("next aggregated log batch requires exactly one entry")
	}
	if cursor != nil && (cursor.StartOffset < 0 || cursor.EndOffset <= cursor.StartOffset) {
		return 0, false, fmt.Errorf("source cursor range is invalid")
	}
	spool.mu.Lock()
	defer spool.mu.Unlock()
	watermark, known := spool.watermarks[batch.LogStreamID]
	if !known && recover != nil {
		latest, err := recover(ctx, batch.LogStreamID)
		if err != nil {
			return 0, false, err
		}
		watermark = logStreamWatermark{Allocated: latest, Acknowledged: latest}
	}
	if cursor != nil && cursor.EndOffset <= watermark.SourceOffset {
		return watermark.Allocated, false, nil
	}
	sequence := watermark.Allocated + 1
	batch.FirstSeq = sequence
	batch.LastSeq = sequence
	batch.Entries[0].Seq = sequence
	var err error
	batch.Checksum, err = checksum(batch.Entries)
	if err != nil {
		return 0, false, err
	}
	if err := spool.writeAggregatedLocked(durableLogBatch{LogBatchIngestRequest: batch, SourceCursor: cursor}, checksum); err != nil {
		return 0, false, err
	}
	watermark.Allocated = sequence
	if cursor != nil {
		watermark.SourceOffset = cursor.EndOffset
	}
	spool.watermarks[batch.LogStreamID] = watermark
	// The batch is already synced and durable. A later startup reconstructs
	// this watermark from pending segments if this snapshot cannot be written.
	_ = spool.persistWatermarksLocked()
	return sequence, true, nil
}

func (spool LogSpool) enqueue(batch protocol.LogBatchIngestRequest, checksum func([]protocol.LogEntry) (string, error)) error {
	spool.mu.Lock()
	defer spool.mu.Unlock()
	watermark := spool.watermarks[batch.LogStreamID]
	if batch.LastSeq > watermark.Allocated {
		watermark.Allocated = batch.LastSeq
		spool.watermarks[batch.LogStreamID] = watermark
		if err := spool.persistWatermarksLocked(); err != nil {
			return err
		}
	}
	return spool.writeAggregatedLocked(durableLogBatch{LogBatchIngestRequest: batch}, checksum)
}

func (spool LogSpool) writeAggregatedLocked(batch durableLogBatch, checksum func([]protocol.LogEntry) (string, error)) error {
	if checksum != nil && len(batch.Entries) == 1 {
		if merged, previous, ok, err := spool.mergeLatest(batch, checksum); err != nil {
			return err
		} else if ok {
			return spool.replaceBatch(previous.path, merged)
		}
	}
	return spool.writeBatch(batch)
}

func (spool LogSpool) mergeLatest(batch durableLogBatch, checksum func([]protocol.LogEntry) (string, error)) (durableLogBatch, pendingLogSegment, bool, error) {
	pending, err := spool.pendingSegmentsLocked()
	if err != nil {
		return durableLogBatch{}, pendingLogSegment{}, false, err
	}
	for index := len(pending) - 1; index >= 0; index-- {
		previous := pending[index]
		if spool.inflight[previous.path] || !compatibleLogBatch(previous.batch, batch) || previous.batch.LastSeq+1 != batch.FirstSeq || len(previous.batch.Entries)+len(batch.Entries) > maxAggregatedLogEntries || !compatibleSourceCursor(previous.batch.SourceCursor, batch.SourceCursor) {
			continue
		}
		merged := previous.batch
		merged.LastSeq = batch.LastSeq
		merged.Entries = append(append([]protocol.LogEntry(nil), previous.batch.Entries...), batch.Entries...)
		if merged.SourceCursor != nil {
			cursor := *merged.SourceCursor
			cursor.EndOffset = batch.SourceCursor.EndOffset
			merged.SourceCursor = &cursor
		}
		merged.Checksum, err = checksum(merged.Entries)
		if err != nil {
			return durableLogBatch{}, pendingLogSegment{}, false, err
		}
		return merged, previous, true, nil
	}
	return durableLogBatch{}, pendingLogSegment{}, false, nil
}

func compatibleLogBatch(previous durableLogBatch, next durableLogBatch) bool {
	return previous.LogStreamID == next.LogStreamID && previous.ServerInstanceID == next.ServerInstanceID && previous.StreamKey == next.StreamKey && previous.Source == next.Source && previous.Compression == next.Compression && previous.LogSessionID == next.LogSessionID && previous.SessionStartedAt.Equal(next.SessionStartedAt)
}

func compatibleSourceCursor(previous *LogSourceCursor, next *LogSourceCursor) bool {
	if previous == nil || next == nil {
		return previous == nil && next == nil
	}
	return next.StartOffset >= previous.EndOffset
}

func (spool LogSpool) writeBatch(batch durableLogBatch) error {
	path := spool.batchPath(batch)
	if existing, err := readDurableLogBatch(path); err == nil {
		if existing.LogStreamID == batch.LogStreamID && existing.FirstSeq == batch.FirstSeq && existing.LastSeq == batch.LastSeq && existing.Checksum == batch.Checksum {
			return nil
		}
		return fmt.Errorf("log spool segment conflicts with committed batch")
	} else if !os.IsNotExist(err) {
		return err
	}
	return spool.writeBatchAt(path, batch)
}

func (spool LogSpool) replaceBatch(path string, batch durableLogBatch) error {
	if strings.TrimSpace(path) == "" {
		return fmt.Errorf("log spool segment path is required")
	}
	return spool.writeBatchAt(path, batch)
}

func (spool LogSpool) writeBatchAt(path string, batch durableLogBatch) error {
	tmp := path + ".tmp"
	file, err := os.OpenFile(tmp, os.O_CREATE|os.O_TRUNC|os.O_WRONLY, 0o600)
	if err != nil {
		return fmt.Errorf("open log spool segment: %w", err)
	}
	encodeErr := json.NewEncoder(file).Encode(batch)
	closeErr := file.Close()
	if encodeErr != nil {
		_ = os.Remove(tmp)
		return fmt.Errorf("encode log spool segment: %w", encodeErr)
	}
	if closeErr != nil {
		_ = os.Remove(tmp)
		return fmt.Errorf("close log spool segment: %w", closeErr)
	}
	if err := syncFile(tmp); err != nil {
		_ = os.Remove(tmp)
		return err
	}
	if err := os.Rename(tmp, path); err != nil {
		_ = os.Remove(tmp)
		return fmt.Errorf("commit log spool segment: %w", err)
	}
	return nil
}

func (spool LogSpool) Pending() ([]protocol.LogBatchIngestRequest, error) {
	spool.mu.Lock()
	defer spool.mu.Unlock()
	segments, err := spool.pendingSegmentsLocked()
	if err != nil {
		return nil, err
	}
	batches := make([]protocol.LogBatchIngestRequest, 0, len(segments))
	for _, segment := range segments {
		batches = append(batches, segment.batch.LogBatchIngestRequest)
	}
	return batches, nil
}

func (spool LogSpool) pendingSegmentsLocked() ([]pendingLogSegment, error) {
	entries, err := os.ReadDir(spool.dir)
	if err != nil {
		return nil, fmt.Errorf("read log spool: %w", err)
	}
	paths := make([]string, 0, len(entries))
	for _, entry := range entries {
		if entry.IsDir() || !strings.HasSuffix(entry.Name(), ".json") {
			continue
		}
		paths = append(paths, filepath.Join(spool.dir, entry.Name()))
	}
	sort.Strings(paths)
	segments := make([]pendingLogSegment, 0, len(paths))
	for _, path := range paths {
		batch, err := readDurableLogBatch(path)
		if err != nil {
			if os.IsNotExist(err) {
				continue
			}
			return nil, err
		}
		segments = append(segments, pendingLogSegment{path: path, batch: batch})
	}
	return segments, nil
}

func (spool LogSpool) Ack(response protocol.LogBatchIngestResponse) error {
	if response.AcceptedFrom == 0 || response.AcceptedTo < response.AcceptedFrom {
		return nil
	}
	spool.mu.Lock()
	defer spool.mu.Unlock()
	watermark := spool.watermarks[response.LogStreamID]
	if response.AcceptedTo > watermark.Acknowledged {
		watermark.Acknowledged = response.AcceptedTo
	}
	if watermark.Allocated < watermark.Acknowledged {
		watermark.Allocated = watermark.Acknowledged
	}
	spool.watermarks[response.LogStreamID] = watermark
	if err := spool.persistWatermarksLocked(); err != nil {
		return err
	}
	segments, err := spool.pendingSegmentsLocked()
	if err != nil {
		return err
	}
	for _, segment := range segments {
		batch := segment.batch
		if batch.LogStreamID == response.LogStreamID && batch.FirstSeq >= response.AcceptedFrom && batch.LastSeq <= response.AcceptedTo {
			if err := os.Remove(segment.path); err != nil && !os.IsNotExist(err) {
				return fmt.Errorf("remove acknowledged log spool segment: %w", err)
			}
		}
	}
	return nil
}

func (spool *LogSpool) loadWatermarks() error {
	body, err := os.ReadFile(spool.statePath)
	if os.IsNotExist(err) {
		return nil
	}
	if err != nil {
		return fmt.Errorf("read log watermarks: %w", err)
	}
	var state logSpoolState
	if err := json.Unmarshal(body, &state); err != nil {
		return fmt.Errorf("decode log watermarks: %w", err)
	}
	if state.Streams != nil {
		spool.watermarks = state.Streams
	}
	return nil
}

func (spool LogSpool) persistWatermarksLocked() error {
	body, err := json.Marshal(logSpoolState{Streams: spool.watermarks})
	if err != nil {
		return err
	}
	temporary := spool.statePath + ".tmp"
	if err := os.WriteFile(temporary, body, 0o600); err != nil {
		return fmt.Errorf("write log watermarks: %w", err)
	}
	if err := syncFile(temporary); err != nil {
		_ = os.Remove(temporary)
		return err
	}
	if err := os.Rename(temporary, spool.statePath); err != nil {
		_ = os.Remove(temporary)
		return fmt.Errorf("commit log watermarks: %w", err)
	}
	return nil
}

type LogBatchClient interface {
	IngestLogBatch(context.Context, protocol.LogBatchIngestRequest) (protocol.LogBatchIngestResponse, error)
}

type PermanentLogBatchError struct {
	Reason string
	Err    error
}

func (err PermanentLogBatchError) Error() string {
	if err.Err == nil {
		return "permanent log batch rejection: " + err.Reason
	}
	return "permanent log batch rejection: " + err.Reason + ": " + err.Err.Error()
}

func (err PermanentLogBatchError) Unwrap() error {
	return err.Err
}

func PermanentLogBatchRejection(reason string, err error) error {
	return PermanentLogBatchError{Reason: sanitizeSegmentName(reason), Err: err}
}

func (spool LogSpool) Flush(ctx context.Context, client LogBatchClient) (int, error) {
	if client == nil {
		return 0, fmt.Errorf("log batch client is required")
	}
	spool.mu.Lock()
	pending, err := spool.pendingSegmentsLocked()
	if err != nil {
		spool.mu.Unlock()
		return 0, err
	}
	for _, segment := range pending {
		spool.inflight[segment.path] = true
	}
	spool.mu.Unlock()
	defer func() {
		spool.mu.Lock()
		defer spool.mu.Unlock()
		for _, segment := range pending {
			delete(spool.inflight, segment.path)
		}
	}()
	acknowledged := 0
	for _, segment := range pending {
		batch := segment.batch.LogBatchIngestRequest
		if err := ctx.Err(); err != nil {
			return acknowledged, err
		}
		response, err := client.IngestLogBatch(ctx, batch)
		if err != nil {
			var permanent PermanentLogBatchError
			if errors.As(err, &permanent) {
				if rejectErr := spool.Reject(batch, permanent.Reason); rejectErr != nil {
					return acknowledged, rejectErr
				}
				acknowledged++
				continue
			}
			return acknowledged, err
		}
		if !response.Accepted || response.LogStreamID != batch.LogStreamID || response.AcceptedFrom > batch.FirstSeq || response.AcceptedTo < batch.LastSeq {
			return acknowledged, fmt.Errorf("platform log acknowledgement does not cover pending batch")
		}
		if err := spool.Ack(response); err != nil {
			return acknowledged, err
		}
		acknowledged++
	}
	return acknowledged, nil
}

func (spool LogSpool) Reject(batch protocol.LogBatchIngestRequest, reason string) error {
	spool.mu.Lock()
	defer spool.mu.Unlock()
	segments, err := spool.pendingSegmentsLocked()
	if err != nil {
		return err
	}
	path := ""
	for _, segment := range segments {
		current := segment.batch
		if current.LogStreamID == batch.LogStreamID && current.FirstSeq == batch.FirstSeq && current.LastSeq == batch.LastSeq && current.Checksum == batch.Checksum {
			path = segment.path
			break
		}
	}
	if path == "" {
		return nil
	}
	rejectedDir := filepath.Join(filepath.Dir(spool.dir), "logs-rejected")
	if err := os.MkdirAll(rejectedDir, 0o755); err != nil {
		return fmt.Errorf("create rejected log spool directory: %w", err)
	}
	rejectedPath := filepath.Join(rejectedDir, fmt.Sprintf("%s.%s", filepath.Base(path), sanitizeSegmentName(reason)))
	if err := os.Rename(path, rejectedPath); err != nil {
		return fmt.Errorf("move rejected log spool segment: %w", err)
	}
	return nil
}

func readDurableLogBatch(path string) (durableLogBatch, error) {
	file, err := os.Open(path)
	if err != nil {
		return durableLogBatch{}, err
	}
	defer file.Close()
	var batch durableLogBatch
	if err := json.NewDecoder(file).Decode(&batch); err != nil {
		return durableLogBatch{}, fmt.Errorf("decode log spool segment: %w", err)
	}
	return batch, nil
}

func syncFile(path string) error {
	file, err := os.OpenFile(path, os.O_RDWR, 0)
	if err != nil {
		return fmt.Errorf("open spool file for sync: %w", err)
	}
	defer file.Close()
	if err := file.Sync(); err != nil {
		return fmt.Errorf("sync spool file: %w", err)
	}
	return nil
}

func (spool LogSpool) batchPath(batch durableLogBatch) string {
	streamID := sanitizeSegmentName(batch.LogStreamID)
	return filepath.Join(spool.dir, fmt.Sprintf("%s-%020d-%020d.json", streamID, batch.FirstSeq, batch.LastSeq))
}

func (spool LogSpool) watermarkPath(streamID string) string {
	return filepath.Join(filepath.Dir(spool.dir), "log-watermarks", sanitizeSegmentName(streamID)+".json")
}

func (spool LogSpool) readWatermark(streamID string) (uint64, error) {
	file, err := os.Open(spool.watermarkPath(streamID))
	if err != nil {
		return 0, err
	}
	defer file.Close()
	var watermark struct {
		LatestSeq uint64 `json:"latestSeq"`
	}
	if err := json.NewDecoder(file).Decode(&watermark); err != nil {
		return 0, fmt.Errorf("decode log watermark: %w", err)
	}
	return watermark.LatestSeq, nil
}

func sanitizeSegmentName(value string) string {
	var builder strings.Builder
	for _, r := range value {
		if r >= 'a' && r <= 'z' || r >= 'A' && r <= 'Z' || r >= '0' && r <= '9' || r == '-' || r == '_' || r == '.' {
			builder.WriteRune(r)
			continue
		}
		builder.WriteByte('_')
	}
	if builder.Len() == 0 {
		return "stream"
	}
	return builder.String()
}