context/context.go

// Package context houses context.Context helpers that are specific to
// lambdalog but that should not live in the root package (keeping the
// public surface there small). Sub-packages add request-ID and tracing
// support.
//
// See mercemay.top/src/lambdalog/context/.
package context

import (
	"context"
	"time"
)

// key is the unexported type used for every value stored in ctx by this
// package. Using an unexported struct type guards against collisions with
// user code.
type key struct{ name string }

var (
	attributesKey = key{"attributes"}
	deadlineKey   = key{"soft-deadline"}
)

// Attributes is the map-like structure attached to ctx by WithAttributes.
// Values are kept as any so callers can attach arbitrary types; the
// encoder will decide how to serialise them.
type Attributes map[string]any

// Clone returns a shallow copy so a receiver can safely mutate without
// racing with the original.
func (a Attributes) Clone() Attributes {
	if len(a) == 0 {
		return nil
	}
	out := make(Attributes, len(a))
	for k, v := range a {
		out[k] = v
	}
	return out
}

// Merge returns a new Attributes containing keys from a with any matching
// keys overwritten by other.
func (a Attributes) Merge(other Attributes) Attributes {
	if len(a) == 0 {
		return other.Clone()
	}
	if len(other) == 0 {
		return a.Clone()
	}
	out := make(Attributes, len(a)+len(other))
	for k, v := range a {
		out[k] = v
	}
	for k, v := range other {
		out[k] = v
	}
	return out
}

// WithAttributes returns a copy of ctx with the given attributes merged in.
// If ctx already holds attributes, the returned ctx holds the combined map.
func WithAttributes(ctx context.Context, attrs Attributes) context.Context {
	if len(attrs) == 0 {
		return ctx
	}
	current, _ := ctx.Value(attributesKey).(Attributes)
	return context.WithValue(ctx, attributesKey, current.Merge(attrs))
}

// AttributesFromContext returns the attributes stored by WithAttributes, or
// nil if none were attached.
func AttributesFromContext(ctx context.Context) Attributes {
	if ctx == nil {
		return nil
	}
	a, _ := ctx.Value(attributesKey).(Attributes)
	return a
}

// WithSoftDeadline attaches a user-facing deadline that is tighter than the
// Lambda invocation deadline. Unlike context.WithDeadline, the returned ctx
// does not fire a Done channel early; the value is informational only so
// log records can report "time remaining until soft deadline".
func WithSoftDeadline(ctx context.Context, at time.Time) context.Context {
	return context.WithValue(ctx, deadlineKey, at)
}

// SoftDeadlineFromContext returns the time set by WithSoftDeadline, along
// with ok==true if one was set.
func SoftDeadlineFromContext(ctx context.Context) (time.Time, bool) {
	if ctx == nil {
		return time.Time{}, false
	}
	v, ok := ctx.Value(deadlineKey).(time.Time)
	return v, ok
}