lambdalog.go

// Package lambdalog provides a small structured logger tuned for AWS Lambda.
//
// The logger writes one JSON object per line to an io.Writer (typically
// os.Stdout). Each record carries a timestamp, level, message, the Lambda
// request ID taken from context.Context, and any user-attached attributes.
//
// See mercemay.top/src/lambdalog/ for the rendered source tree.
package lambdalog

import (
	"bytes"
	"context"
	"encoding/json"
	"fmt"
	"io"
	"os"
	"sync"
	"time"
)

// Level describes the severity of a record.
type Level string

const (
	LevelDebug Level = "debug"
	LevelInfo  Level = "info"
	LevelWarn  Level = "warn"
	LevelError Level = "error"
)

// attr is a single key/value pair. We keep attributes as an ordered slice
// rather than a map so JSON output order is deterministic.
type attr struct {
	Key   string
	Value any
}

// Logger emits JSON records to an io.Writer. Loggers are safe for concurrent
// use; writes are serialized through an internal mutex.
type Logger struct {
	mu       *sync.Mutex
	w        io.Writer
	attrs    []attr
	sampler  *Sampler
	minLevel Level
	now      func() time.Time
}

// New constructs a Logger writing to w. The level threshold defaults to Info;
// use WithLevel to change it.
func New(w io.Writer) *Logger {
	if w == nil {
		w = os.Stdout
	}
	return &Logger{
		mu:       &sync.Mutex{},
		w:        w,
		minLevel: LevelInfo,
		now:      time.Now,
	}
}

// With returns a child logger carrying an additional attribute. The parent's
// attribute slice is deep-copied so mutations on either logger are isolated.
func (l *Logger) With(key string, value any) *Logger {
	c := l.clone()
	c.attrs = append(c.attrs, attr{Key: key, Value: value})
	return c
}

// WithFields returns a child logger carrying several attributes at once.
func (l *Logger) WithFields(kv ...any) *Logger {
	c := l.clone()
	for i := 0; i+1 < len(kv); i += 2 {
		k, ok := kv[i].(string)
		if !ok {
			k = fmt.Sprintf("%v", kv[i])
		}
		c.attrs = append(c.attrs, attr{Key: k, Value: kv[i+1]})
	}
	if len(kv)%2 == 1 {
		c.attrs = append(c.attrs, attr{Key: "MISSING", Value: kv[len(kv)-1]})
	}
	return c
}

// WithLevel returns a child logger with a new minimum level threshold.
func (l *Logger) WithLevel(lv Level) *Logger {
	c := l.clone()
	c.minLevel = lv
	return c
}

// Sampled returns a child logger that throttles records by an adaptive
// 1-in-N sampler keyed by name. See Sampler for the decision policy.
func (l *Logger) Sampled(name string, n int) *Logger {
	c := l.clone()
	c.sampler = NewSampler(name, n)
	return c
}

// FromContext returns a child logger enriched with context-bound fields,
// notably the Lambda request ID under key "rid".
func (l *Logger) FromContext(ctx context.Context) *Logger {
	if ctx == nil {
		return l
	}
	c := l.clone()
	if rid := requestIDFromContext(ctx); rid != "" {
		c.attrs = append(c.attrs, attr{Key: "rid", Value: rid})
	}
	for _, a := range fieldsFromContext(ctx) {
		c.attrs = append(c.attrs, a)
	}
	return c
}

// Debug emits a record at LevelDebug.
func (l *Logger) Debug(msg string, kv ...any) { l.log(LevelDebug, msg, kv) }

// Info emits a record at LevelInfo.
func (l *Logger) Info(msg string, kv ...any) { l.log(LevelInfo, msg, kv) }

// Warn emits a record at LevelWarn.
func (l *Logger) Warn(msg string, kv ...any) { l.log(LevelWarn, msg, kv) }

// Error emits a record at LevelError.
func (l *Logger) Error(msg string, kv ...any) { l.log(LevelError, msg, kv) }

func (l *Logger) log(lv Level, msg string, kv []any) {
	if !levelEnabled(lv, l.minLevel) {
		return
	}
	if l.sampler != nil && !l.sampler.Allow() {
		return
	}

	var buf bytes.Buffer
	buf.Grow(256)
	buf.WriteByte('{')
	writeField(&buf, "ts", l.now().UTC().Format(time.RFC3339Nano))
	buf.WriteByte(',')
	writeField(&buf, "level", string(lv))
	buf.WriteByte(',')
	writeField(&buf, "msg", msg)
	for _, a := range l.attrs {
		buf.WriteByte(',')
		writeField(&buf, a.Key, a.Value)
	}
	for i := 0; i+1 < len(kv); i += 2 {
		buf.WriteByte(',')
		k, ok := kv[i].(string)
		if !ok {
			k = fmt.Sprintf("%v", kv[i])
		}
		writeField(&buf, k, kv[i+1])
	}
	if len(kv)%2 == 1 {
		buf.WriteByte(',')
		writeField(&buf, "MISSING", kv[len(kv)-1])
	}
	buf.WriteByte('}')
	buf.WriteByte('\n')

	l.mu.Lock()
	defer l.mu.Unlock()
	_, _ = l.w.Write(buf.Bytes())
}

func (l *Logger) clone() *Logger {
	attrs := make([]attr, len(l.attrs))
	copy(attrs, l.attrs)
	return &Logger{
		mu:       l.mu,
		w:        l.w,
		attrs:    attrs,
		sampler:  l.sampler,
		minLevel: l.minLevel,
		now:      l.now,
	}
}

func writeField(buf *bytes.Buffer, k string, v any) {
	kb, _ := json.Marshal(k)
	buf.Write(kb)
	buf.WriteByte(':')
	// Marshal the value via encoding/json for robust escaping of strings,
	// nested maps, errors rendered via their Error() method, etc.
	vb, err := json.Marshal(v)
	if err != nil {
		// Fall back to a stringified form so a single bad value cannot
		// corrupt the stream.
		fallback, _ := json.Marshal(fmt.Sprintf("<unmarshallable: %v>", err))
		buf.Write(fallback)
		return
	}
	buf.Write(vb)
}

func levelEnabled(have, min Level) bool {
	return levelRank(have) >= levelRank(min)
}

func levelRank(l Level) int {
	switch l {
	case LevelDebug:
		return 0
	case LevelInfo:
		return 1
	case LevelWarn:
		return 2
	case LevelError:
		return 3
	default:
		return 1
	}
}