package caller

caller.go

package caller

import (
	"encoding/json"
	"fmt"
	"io"
	"runtime"
	"strconv"
	"strings"

	"github.com/go-kit/log"
)

type Context struct {
	help runtime.Frame
}

func New() *Context {
	return &Context{}
}

func (c *Context) SetHelper() {
	c.help = getFrames(3)[0]
}

func (c *Context) ClearHelper() {
	c.help = zeroFrame
}

func (c *Context) Value() log.Valuer {
	return c.ValueFunc(func(f runtime.Frame) string {
		file := f.File
		if index := strings.LastIndex(file, "/"); index > 0 {
			file = file[index+1:]
		}
		return file + ":" + strconv.Itoa(f.Line)
	})
}

func (c *Context) ValueFunc(f func(runtime.Frame) string) log.Valuer {
	return func() interface{} {
		return lazyValuer(func() interface{} {
			switch {
			case c.help != zeroFrame:
				return f(c.help)
			default:
				return f(filterFrames(getFrames(3), firstPackageLogFrame)[0])
			}
		})
	}
}

func (c *Context) valueWithCallback(cb func([]runtime.Frame)) log.Valuer {
	return func() interface{} {
		return lazyValuer(func() interface{} {
			cb(getFrames(3))
			return nil
		})
	}
}

type lazyValuer func() interface{}

func (z lazyValuer) String() string {
	return fmt.Sprintf("%v", z()) // TODO: safety
}

func (z lazyValuer) MarshalJSON() ([]byte, error) {
	return json.Marshal(z()) // TODO: safety
}

func getFrames(skip int) []runtime.Frame {
	pc := make([]uintptr, 1000) // TODO: sync.Pool

	n := runtime.Callers(skip, pc)
	if n <= 0 {
		return nil
	}

	fs := []runtime.Frame{}
	frames := runtime.CallersFrames(pc[:n])

	for {
		frame, more := frames.Next()
		if frame == zeroFrame {
			continue
		}

		fs = append(fs, frame)

		if !more {
			break
		}
	}

	return fs
}

var zeroFrame runtime.Frame

var firstPackageLogFrame = func() runtime.Frame {
	var (
		context  = New()
		frames   = []runtime.Frame{}
		callback = func(fs []runtime.Frame) { frames = fs }
		value    = context.valueWithCallback(callback)
		logger   = log.NewLogfmtLogger(io.Discard)
	)

	log.With(logger, "caller", value).Log()

	for i := len(frames) - 1; i >= 0; i-- {
		if f := frames[i]; strings.HasPrefix(f.Function, `github.com/go-kit/log.`) {
			return f
		}
	}

	return zeroFrame
}()

func filterFrames(fs []runtime.Frame, target runtime.Frame) []runtime.Frame {
	for i, f := range fs {
		if f == target {
			return fs[i+1:]
		}
	}
	return fs
}

caller_test.go

package caller

import (
	"os"
	"testing"

	"github.com/go-kit/log"
)

func foo(lg log.Logger) {
	lg.Log("func", "foo", "src", "caller_test.go:11")
}

func TestContext(t *testing.T) {
	context := New()
	logger := log.With(log.NewLogfmtLogger(os.Stderr), "caller", context.Value())
	logger.Log("msg", "calling foo", "src", "caller_test.go:17", "want_next_caller", "caller_test.go:11")
	foo(logger)
	logger.Log("msg", "setting helper", "src", "caller_test.go:19")
	context.SetHelper()
	logger.Log("msg", "calling foo", "src", "caller_test.go:21", "want_next_caller", "caller_test.go:20")
	foo(logger)
	logger.Log("msg", "removing helper", "src", "caller_test.go:23")
	context.ClearHelper()
	logger.Log("msg", "calling foo with prefix", "src", "caller_test.go:25", "want_next_caller", "caller_test.go:11")
	foo(log.With(logger, "prefix", "value"))
	logger.Log("msg", "calling foo with decorator", "src", "caller_test.go:27", "want_next_caller", "caller_test.go:11")
	foo(logDecorator{logger}) // TODO: problem -- fix
}

type logDecorator struct{ lg log.Logger }

func (d logDecorator) Log(keyvals ...any) error {
	return d.lg.Log(append(keyvals, "decorator", true)...)
}