From 054e5e38afc7139a543d98b98d649b0c7b735d04 Mon Sep 17 00:00:00 2001
From: r1viollet <74836499+r1viollet@users.noreply.github.com>
Date: Tue, 10 Sep 2024 17:55:43 +0200
Subject: [PATCH] Timeline (#55)

* Timeline
- Allow timestamps to be added to the samples as labels.
- Ensure process name is added as a label to help with thread grouping logics.
---
 cli_flags.go                 |  2 ++
 main.go                      |  1 +
 reporter/config.go           |  3 +++
 reporter/datadog_reporter.go | 32 ++++++++++++++++++++++++++------
 4 files changed, 32 insertions(+), 6 deletions(-)

diff --git a/cli_flags.go b/cli_flags.go
index 03574bc..80afb50 100644
--- a/cli_flags.go
+++ b/cli_flags.go
@@ -99,6 +99,7 @@ var (
 	argProbabilisticInterval  time.Duration
 	argPprofAddr              string
 	argSaveCPUProfile         bool
+	argTimeline               bool
 
 	// "internal" flag variables.
 	// Flag variables that are configured in "internal" builds will have to be assigned
@@ -161,6 +162,7 @@ func parseArgs() error {
 		saveCPUProfileHelp)
 	fs.IntVar(&argSamplesPerSecond, "samples-per-second", defaultArgSamplesPerSecond,
 		samplesPerSecondHelp)
+	fs.BoolVar(&argTimeline, "timeline", false, "Enable timeline feature.")
 
 	fs.Usage = func() {
 		fs.PrintDefaults()
diff --git a/main.go b/main.go
index 3ab9bd4..f9aa63e 100644
--- a/main.go
+++ b/main.go
@@ -315,6 +315,7 @@ func mainWithExitCode() exitCode {
 		FallbackSymbolsMaxQueue: 1024,
 		DisableTLS:              argDisableTLS,
 		MaxGRPCRetries:          5,
+		Timeline:                argTimeline,
 		SamplesPerSecond:        conf.SamplesPerSecond,
 		SaveCPUProfile:          argSaveCPUProfile,
 		Times:                   times,
diff --git a/reporter/config.go b/reporter/config.go
index 5b42e47..337bd69 100644
--- a/reporter/config.go
+++ b/reporter/config.go
@@ -51,6 +51,9 @@ type Config struct {
 
 	SaveCPUProfile bool
 
+	// Whether to include timestamps on samples for the timeline feature
+	Timeline bool
+
 	Times Times
 
 	// gRPCInterceptor is the client gRPC interceptor, e.g., for sending gRPC metadata.
diff --git a/reporter/datadog_reporter.go b/reporter/datadog_reporter.go
index 5e2e619..1259a08 100644
--- a/reporter/datadog_reporter.go
+++ b/reporter/datadog_reporter.go
@@ -15,6 +15,7 @@ import (
 	"os"
 	"path"
 	"runtime"
+	"strconv"
 	"strings"
 	"time"
 
@@ -49,6 +50,8 @@ type DatadogReporter struct {
 
 	agentAddr string
 
+	timeline bool
+
 	samplingPeriod uint64
 
 	saveCPUProfile bool
@@ -255,6 +258,7 @@ func StartDatadog(mainCtx context.Context, cfg *Config) (Reporter, error) {
 		agentAddr:       cfg.CollAgentAddr,
 		samplingPeriod:  1000000000 / uint64(cfg.SamplesPerSecond),
 		saveCPUProfile:  cfg.SaveCPUProfile,
+		timeline:        cfg.Timeline,
 		fallbackSymbols: fallbackSymbols,
 		executables:     executables,
 		frames:          frames,
@@ -512,17 +516,21 @@ func (r *DatadogReporter) getPprofProfile() (profile *pprofile.Profile,
 			// location with the kernel as the function name.
 			execPath = "kernel"
 		}
+		baseExec := path.Base(execPath)
 
 		if execPath != "" {
-			base := path.Base(execPath)
 			loc := createPProfLocation(profile, 0)
-			m := createPprofFunctionEntry(funcMap, profile, base, execPath)
+			m := createPprofFunctionEntry(funcMap, profile, baseExec, execPath)
 			loc.Line = append(loc.Line, pprofile.Line{Function: m})
 			sample.Location = append(sample.Location, loc)
 		}
 
 		sample.Label = make(map[string][]string)
-		addTraceLabels(sample.Label, traceKey)
+		var timestamps []uint64
+		if r.timeline {
+			timestamps = traceInfo.timestamps
+		}
+		addTraceLabels(sample.Label, &traceKey, baseExec, timestamps) //nolint:gomemory
 
 		count := int64(len(traceInfo.timestamps))
 		sample.Value = append(sample.Value, count, count*int64(r.samplingPeriod))
@@ -564,12 +572,11 @@ func createPprofFunctionEntry(funcMap map[funcInfo]*pprofile.Function,
 	return function
 }
 
-//nolint:gocritic
-func addTraceLabels(labels map[string][]string, k traceAndMetaKey) {
+func addTraceLabels(labels map[string][]string, k *traceAndMetaKey,
+	baseExec string, timestamps []uint64) {
 	if k.comm != "" {
 		labels["thread_name"] = append(labels["thread_name"], k.comm)
 	}
-
 	if k.podName != "" {
 		labels["pod_name"] = append(labels["pod_name"], k.podName)
 	}
@@ -596,6 +603,19 @@ func addTraceLabels(labels map[string][]string, k traceAndMetaKey) {
 		// This is also consistent with ddprof.
 		labels["thread id"] = append(labels["thread id"], fmt.Sprintf("%d", k.tid))
 	}
+
+	if baseExec != "" {
+		labels["process_name"] = append(labels["process_name"], baseExec)
+	}
+
+	if len(timestamps) > 0 {
+		timestampStrs := make([]string, 0, len(timestamps))
+		for _, ts := range timestamps {
+			timestampStrs = append(timestampStrs, strconv.FormatUint(ts, 10))
+		}
+		// Assign all timestamps as a single label entry
+		labels["end_timestamp_ns"] = timestampStrs
+	}
 }
 
 // getDummyMappingIndex inserts or looks up a dummy entry for interpreted FileIDs.