agents_and_robots/shell/logger/writer.go

package logger

import (
	"compress/gzip"
	"fmt"
	"io"
	"os"
	"path/filepath"
	"sync"
	"time"
)

// DailyRotatingWriter is an io.Writer that rotates log files daily and by
// size. Files are named <baseDir>/<agentID>/YYYY-MM-DD.jsonl with optional
// numeric suffixes for size-based splits within the same day.
type DailyRotatingWriter struct {
	baseDir   string
	agentID   string
	maxSize   int64 // bytes
	compress  bool
	nowFunc   func() time.Time // for testing; defaults to time.Now().UTC
	dir       string           // resolved agent log directory

	mu         sync.Mutex
	current    *os.File
	written    int64
	currentDay string
	suffix     int
}

// NewDailyRotatingWriter creates a writer that stores logs under
// baseDir/agentID/. It creates the directory if needed and opens the first
// log file for today.
func NewDailyRotatingWriter(baseDir, agentID string, maxSizeMB int64, compress bool) (*DailyRotatingWriter, error) {
	dir := filepath.Join(baseDir, agentID)
	if err := os.MkdirAll(dir, 0o755); err != nil {
		return nil, fmt.Errorf("create log dir %s: %w", dir, err)
	}

	w := &DailyRotatingWriter{
		baseDir:  baseDir,
		agentID:  agentID,
		maxSize:  maxSizeMB * 1024 * 1024,
		compress: compress,
		nowFunc:  func() time.Time { return time.Now().UTC() },
		dir:      dir,
	}

	if err := w.openFile(); err != nil {
		return nil, err
	}
	return w, nil
}

// Write implements io.Writer with daily and size-based rotation.
func (w *DailyRotatingWriter) Write(p []byte) (int, error) {
	w.mu.Lock()
	defer w.mu.Unlock()

	today := w.nowFunc().Format("2006-01-02")

	// Day changed → rotate to new day file.
	if today != w.currentDay {
		if err := w.rotate(today, 0); err != nil {
			return 0, err
		}
	}

	// Size exceeded → split within same day.
	if w.written+int64(len(p)) > w.maxSize && w.written > 0 {
		w.suffix++
		if err := w.rotate(today, w.suffix); err != nil {
			return 0, err
		}
	}

	n, err := w.current.Write(p)
	w.written += int64(n)
	return n, err
}

// Close closes the current log file.
func (w *DailyRotatingWriter) Close() error {
	w.mu.Lock()
	defer w.mu.Unlock()
	if w.current != nil {
		return w.current.Close()
	}
	return nil
}

// rotate closes the current file (optionally compressing it) and opens a new one.
func (w *DailyRotatingWriter) rotate(day string, suffix int) error {
	prev := w.current
	prevPath := ""
	if prev != nil {
		prevPath = prev.Name()
		prev.Close()
	}

	// Compress the previous file in the background if enabled and it's from a
	// different day (we don't compress intra-day splits until day rotates).
	if w.compress && prevPath != "" && day != w.currentDay {
		go compressFile(prevPath)
	}

	w.currentDay = day
	w.suffix = suffix
	w.written = 0

	return w.openFile()
}

// openFile opens (or creates) the log file for the current day/suffix.
func (w *DailyRotatingWriter) openFile() error {
	w.currentDay = w.nowFunc().Format("2006-01-02")
	name := w.filename(w.currentDay, w.suffix)

	f, err := os.OpenFile(name, os.O_CREATE|os.O_WRONLY|os.O_APPEND, 0o644)
	if err != nil {
		return fmt.Errorf("open log file %s: %w", name, err)
	}

	// Track how much has already been written (append mode).
	info, err := f.Stat()
	if err != nil {
		f.Close()
		return err
	}

	w.current = f
	w.written = info.Size()
	return nil
}

// filename returns the full path for a given day and suffix.
func (w *DailyRotatingWriter) filename(day string, suffix int) string {
	if suffix == 0 {
		return filepath.Join(w.dir, day+".jsonl")
	}
	return filepath.Join(w.dir, fmt.Sprintf("%s.%d.jsonl", day, suffix))
}

// compressFile gzips src to src.gz and removes the original.
func compressFile(src string) {
	in, err := os.Open(src)
	if err != nil {
		return
	}
	defer in.Close()

	out, err := os.Create(src + ".gz")
	if err != nil {
		return
	}

	gz := gzip.NewWriter(out)
	if _, err := io.Copy(gz, in); err != nil {
		gz.Close()
		out.Close()
		os.Remove(src + ".gz")
		return
	}
	gz.Close()
	out.Close()
	in.Close()
	os.Remove(src)
}