Add system to track only Autometricized parents

This system will likely fail whenever a Trace forks into multiple goroutines, this is something to take care of. And likely the reason we will need some kind of goroutine local storage on top of the global state

Fixes: AM-42

otel/autometrics/instrument.go

@@ -47,10 +47,10 @@ func Instrument(ctx context.Context, err *error) {
 
 	functionCallsCount.Add(ctx, 1,
 		metric.WithAttributes([]attribute.KeyValue{
-			attribute.Key(FunctionLabel).String(callInfo.FuncName),
-			attribute.Key(ModuleLabel).String(callInfo.ModuleName),
-			attribute.Key(CallerFunctionLabel).String(callInfo.ParentFuncName),
-			attribute.Key(CallerModuleLabel).String(callInfo.ParentModuleName),
+			attribute.Key(FunctionLabel).String(callInfo.Current.Function),
+			attribute.Key(ModuleLabel).String(callInfo.Current.Module),
+			attribute.Key(CallerFunctionLabel).String(callInfo.Parent.Function),
+			attribute.Key(CallerModuleLabel).String(callInfo.Parent.Module),
 			attribute.Key(ResultLabel).String(result),
 			attribute.Key(TargetSuccessRateLabel).String(successObjective),
 			attribute.Key(SloNameLabel).String(sloName),
@@ -62,10 +62,10 @@ func Instrument(ctx context.Context, err *error) {
 		}...))
 	functionCallsDuration.Record(ctx, time.Since(am.GetStartTime(ctx)).Seconds(),
 		metric.WithAttributes([]attribute.KeyValue{
-			attribute.Key(FunctionLabel).String(callInfo.FuncName),
-			attribute.Key(ModuleLabel).String(callInfo.ModuleName),
-			attribute.Key(CallerFunctionLabel).String(callInfo.ParentFuncName),
-			attribute.Key(CallerModuleLabel).String(callInfo.ParentModuleName),
+			attribute.Key(FunctionLabel).String(callInfo.Current.Function),
+			attribute.Key(ModuleLabel).String(callInfo.Current.Module),
+			attribute.Key(CallerFunctionLabel).String(callInfo.Parent.Function),
+			attribute.Key(CallerModuleLabel).String(callInfo.Parent.Module),
 			attribute.Key(TargetLatencyLabel).String(latencyTarget),
 			attribute.Key(TargetSuccessRateLabel).String(latencyObjective),
 			attribute.Key(SloNameLabel).String(sloName),
@@ -79,17 +79,19 @@ func Instrument(ctx context.Context, err *error) {
 	if am.GetTrackConcurrentCalls(ctx) {
 		functionCallsConcurrent.Add(ctx, -1,
 			metric.WithAttributes([]attribute.KeyValue{
-				attribute.Key(FunctionLabel).String(callInfo.FuncName),
-				attribute.Key(ModuleLabel).String(callInfo.ModuleName),
-				attribute.Key(CallerFunctionLabel).String(callInfo.ParentFuncName),
-				attribute.Key(CallerModuleLabel).String(callInfo.ParentModuleName),
+				attribute.Key(FunctionLabel).String(callInfo.Current.Function),
+				attribute.Key(ModuleLabel).String(callInfo.Current.Module),
+				attribute.Key(CallerFunctionLabel).String(callInfo.Parent.Function),
+				attribute.Key(CallerModuleLabel).String(callInfo.Parent.Module),
 				attribute.Key(CommitLabel).String(buildInfo.Commit),
 				attribute.Key(VersionLabel).String(buildInfo.Version),
 				attribute.Key(BranchLabel).String(buildInfo.Branch),
 				attribute.Key(ServiceNameLabel).String(buildInfo.Service),
 				attribute.Key(JobNameLabel).String(am.GetPushJobName()),
 			}...))
 	}
+
+	am.PopFunctionName(ctx)
 }
 
 // PreInstrument runs the "before wrappee" part of instrumentation.
@@ -101,19 +103,20 @@ func PreInstrument(ctx context.Context) context.Context {
 		return nil
 	}
 
-	callInfo := am.CallerInfo()
-	ctx = am.SetCallInfo(ctx, callInfo)
+	ctx = am.FillTracingAndCallerInfo(ctx)
 	ctx = am.FillBuildInfo(ctx)
-	ctx = am.FillTracingInfo(ctx)
+
+	callInfo := am.GetCallInfo(ctx)
+	am.PushFunctionName(ctx, callInfo.Current)
 
 	if am.GetTrackConcurrentCalls(ctx) {
 		buildInfo := am.GetBuildInfo(ctx)
 		functionCallsConcurrent.Add(ctx, 1,
 			metric.WithAttributes([]attribute.KeyValue{
-				attribute.Key(FunctionLabel).String(callInfo.FuncName),
-				attribute.Key(ModuleLabel).String(callInfo.ModuleName),
-				attribute.Key(CallerFunctionLabel).String(callInfo.ParentFuncName),
-				attribute.Key(CallerModuleLabel).String(callInfo.ParentModuleName),
+				attribute.Key(FunctionLabel).String(callInfo.Current.Function),
+				attribute.Key(ModuleLabel).String(callInfo.Current.Module),
+				attribute.Key(CallerFunctionLabel).String(callInfo.Parent.Function),
+				attribute.Key(CallerModuleLabel).String(callInfo.Parent.Module),
 				attribute.Key(CommitLabel).String(buildInfo.Commit),
 				attribute.Key(VersionLabel).String(buildInfo.Version),
 				attribute.Key(BranchLabel).String(buildInfo.Branch),

pkg/autometrics/ctx.go

@@ -2,6 +2,7 @@ package autometrics // import "github.com/autometrics-dev/autometrics-go/pkg/aut
 
 import (
 	"context"
+	"errors"
 	"log"
 	"math/rand"
 	"time"
@@ -229,12 +230,12 @@ func GetParentSpanID(c context.Context) (SpanID, bool) {
 	return sid, ok
 }
 
-// FillTracingInfo ensures the context has a traceID and a spanID.
+// FillTracingAndCallerInfo ensures the context has a traceID and a spanID, and looks for relevant caller information to add in the context as well.
 // If they do not have this information, this method adds randomly
 // generated IDs in the context to be used later for exemplars
 //
 // The random generator is a PRNG, seeded with the timestamp of the first time new IDs are needed.
-func FillTracingInfo(ctx context.Context) context.Context {
+func FillTracingAndCallerInfo(ctx context.Context) context.Context {
 	// We are using a PRNG because FillTracingInfo is expected to be called in PreInstrument.
 	// Therefore it can have a noticeable impact on the performance of instrumented code.
 	// Pseudo randomness should be enough for our use cases, true randomness might introduce too much latency.
@@ -258,6 +259,9 @@ func FillTracingInfo(ctx context.Context) context.Context {
 		ctx = SetTraceID(ctx, tid)
 	}
 
+	callInfo := callerInfo(ctx)
+	ctx = SetCallInfo(ctx, callInfo)
+
 	return ctx
 }
 
@@ -336,3 +340,32 @@ func GetValidHttpCodeRanges(c context.Context) []InclusiveIntRange {
 
 	return ranges
 }
+
+func PeekFunctionName(ctx context.Context) (FunctionID, error) {
+	tid, ok := GetTraceID(ctx)
+	if !ok {
+		return FunctionID{}, errors.New("context does not have any trace ID to follow.")
+	}
+
+	return peekFunctionName(tid)
+}
+
+func PushFunctionName(ctx context.Context, functionID FunctionID) error {
+	tid, ok := GetTraceID(ctx)
+	if !ok {
+		return errors.New("context does not have any trace ID to follow.")
+	}
+
+	pushFunctionName(tid, functionID)
+
+	return nil
+}
+
+func PopFunctionName(ctx context.Context) (FunctionID, error) {
+	tid, ok := GetTraceID(ctx)
+	if !ok {
+		return FunctionID{}, errors.New("context does not have any trace ID to follow.")
+	}
+
+	return popFunctionName(tid)
+}

pkg/autometrics/global_state.go

@@ -1,5 +1,9 @@
 package autometrics // import "github.com/autometrics-dev/autometrics-go/pkg/autometrics"
 
+import (
+	"fmt"
+)
+
 // These variables are describing the state of the application being autometricized,
 // _not_ the build information of the binary
 
@@ -15,12 +19,13 @@ const (
 )
 
 var (
-	version     string
-	commit      string
-	branch      string
-	service     string
-	pushJobName string
-	pushJobURL  string
+	version                 string
+	commit                  string
+	branch                  string
+	service                 string
+	pushJobName             string
+	pushJobURL              string
+	instrumentedNamesStacks map[TraceID]*InstrumentedFunctionsCallStack = make(map[TraceID]*InstrumentedFunctionsCallStack)
 )
 
 // GetVersion returns the version of the codebase being instrumented.
@@ -82,3 +87,40 @@ func GetPushJobURL() string {
 func SetPushJobURL(newPushJobURL string) {
 	pushJobURL = newPushJobURL
 }
+
+func peekFunctionName(traceID TraceID) (FunctionID, error) {
+	stack := instrumentedNamesStacks[traceID]
+	if stack == nil {
+		return FunctionID{}, fmt.Errorf("%v is not a known traceID now", traceID)
+	}
+
+	return stack.Peek(0)
+}
+
+func pushFunctionName(traceID TraceID, functionID FunctionID) {
+	stack := instrumentedNamesStacks[traceID]
+	if stack == nil {
+		instrumentedNamesStacks[traceID] = NewInstrumentedFunctionsCallStack()
+		stack = instrumentedNamesStacks[traceID]
+	}
+
+	stack.Push(functionID)
+}
+
+func popFunctionName(traceID TraceID) (FunctionID, error) {
+	stack := instrumentedNamesStacks[traceID]
+	if stack == nil {
+		return FunctionID{}, fmt.Errorf("%v is not a known traceID now", traceID)
+	}
+
+	val, err := stack.Pop()
+	if err != nil {
+		return val, fmt.Errorf("issue popping the element: %w", err)
+	}
+
+	if stack.IsEmpty() {
+		delete(instrumentedNamesStacks, traceID)
+	}
+
+	return val, nil
+}

pkg/autometrics/id_stack.go

@@ -0,0 +1,50 @@
+package autometrics // import "github.com/autometrics-dev/autometrics-go/pkg/autometrics"
+
+import (
+	"errors"
+	"fmt"
+)
+
+// InstrumentedFunctionsCallStack is a stack containing only function IDs (pairs of function name and module name) for
+// autometricized functions in a given callstack.
+type InstrumentedFunctionsCallStack struct {
+	storage []FunctionID
+}
+
+func NewInstrumentedFunctionsCallStack() *InstrumentedFunctionsCallStack {
+	return &InstrumentedFunctionsCallStack{
+		storage: make([]FunctionID, 0, 16),
+	}
+}
+
+func (stack InstrumentedFunctionsCallStack) IsEmpty() bool {
+	return len(stack.storage) == 0
+}
+
+func (stack InstrumentedFunctionsCallStack) Peek(fromTop int) (FunctionID, error) {
+	l := len(stack.storage)
+	if fromTop >= l {
+		return FunctionID{}, fmt.Errorf("out of bounds access: %v while length is %v", fromTop, l)
+	}
+
+	return stack.storage[l-1-fromTop], nil
+}
+
+func (stack *InstrumentedFunctionsCallStack) Push(functionID FunctionID) int {
+	stack.storage = append(stack.storage, functionID)
+
+	return len(stack.storage)
+}
+
+func (stack *InstrumentedFunctionsCallStack) Pop() (FunctionID, error) {
+	l := len(stack.storage)
+
+	if l == 0 {
+		return FunctionID{}, errors.New("out of bounds access: stack is empty")
+	}
+
+	topElem := stack.storage[l-1]
+	stack.storage = stack.storage[:l]
+
+	return topElem, nil
+}

pkg/autometrics/instrument.go

@@ -1,14 +1,15 @@
 package autometrics
 
 import (
+	"context"
 	"reflect"
 	"runtime"
 	"strings"
 )
 
 // CallerInfo returns the (method name, module name) of the function that called the function that called this function.
 //
-// It also returns the information about its grandparent.
+// It also returns the information about its autometricized grandparent.
 //
 // The module name and the parent module names are cropped to their last part, because the generator we use
 // only has access to the last "package" name in `GOPACKAGE` environment variable.
@@ -17,7 +18,7 @@ import (
 // then we can lift this artificial limitation here and use the full "module name" from the caller information.
 // Currently this compromise is the only way to have the documentation links generator creating correct
 // queries.
-func CallerInfo() (callInfo CallInfo) {
+func callerInfo(ctx context.Context) (callInfo CallInfo) {
 	programCounters := make([]uintptr, 15)
 
 	// skip 3 frames to start with:
@@ -27,45 +28,30 @@ func CallerInfo() (callInfo CallInfo) {
 	entries := runtime.Callers(3, programCounters)
 
 	frames := runtime.CallersFrames(programCounters[:entries])
-	frame, hasParent := frames.Next()
+	frame, _ := frames.Next()
 
 	functionName := frame.Function
 	index := strings.LastIndex(functionName, ".")
 
 	if index == -1 {
-		callInfo.FuncName = functionName
+		callInfo.Current.Function = functionName
 	} else {
-		callInfo.ModuleName = strings.ReplaceAll(strings.ReplaceAll(strings.ReplaceAll(
+		callInfo.Current.Module = strings.ReplaceAll(strings.ReplaceAll(strings.ReplaceAll(
 			functionName[:index],
 			"(", ""),
 			")", ""),
 			"*", "")
 
-		callInfo.FuncName = functionName[index+1:]
+		callInfo.Current.Function = functionName[index+1:]
 	}
 
-	if !hasParent {
+	parent, err := PeekFunctionName(ctx)
+
+	if err != nil {
 		return
 	}
 
-	// Do the same with the parent
-	parentFrame, _ := frames.Next()
-
-	parentFunctionName := parentFrame.Function
-	index = strings.LastIndex(parentFunctionName, ".")
-
-	if index == -1 {
-		callInfo.ParentFuncName = parentFunctionName
-	} else {
-		moduleIndex := strings.LastIndex(parentFunctionName[:index], ".")
-		if moduleIndex == -1 {
-			callInfo.ParentModuleName = parentFunctionName[:index]
-		} else {
-			callInfo.ParentModuleName = parentFunctionName[moduleIndex+1 : index]
-		}
-
-		callInfo.ParentFuncName = parentFunctionName[index+1:]
-	}
+	callInfo.Parent = parent
 
 	return
 }
@@ -80,15 +66,15 @@ func ReflectFunctionModuleName(f interface{}) (callInfo CallInfo) {
 
 	index := strings.LastIndex(functionName, ".")
 	if index == -1 {
-		callInfo.FuncName = functionName
+		callInfo.Current.Function = functionName
 	} else {
-		moduleIndex := strings.LastIndex(functionName[:index], ".")
-		if moduleIndex == -1 {
-			callInfo.ModuleName = functionName[:index]
-		} else {
-			callInfo.ModuleName = functionName[moduleIndex+1 : index]
-		}
-		callInfo.FuncName = functionName[index+1:]
+		callInfo.Current.Module = strings.ReplaceAll(strings.ReplaceAll(strings.ReplaceAll(
+			functionName[:index],
+			"(", ""),
+			")", ""),
+			"*", "")
+
+		callInfo.Current.Function = functionName[index+1:]
 	}
 
 	return callInfo