Support pganalyze tracestate to set start time of the span

logs/querysample/tracing.go

@@ -5,6 +5,8 @@ import (
 	"encoding/binary"
 	"encoding/hex"
 	"fmt"
+	"strconv"
+	"strings"
 	"time"
 
 	"github.com/pganalyze/collector/state"
@@ -32,6 +34,41 @@ func urlToSample(server *state.Server, grant state.GrantLogs, sample state.Postg
 	)
 }
 
+func startAndEndTime(traceState trace.TraceState, sample state.PostgresQuerySample) (startTime time.Time, endTime time.Time) {
+	if pganalyzeState := traceState.Get("pganalyze"); pganalyzeState != "" {
+		// A pganalyze traceState allows the client to pass the query start time (sent time)
+		// on the client side, in nano second precision, like pganalyze=t:1697666938.6297212
+		// If there are multiple values in a pganalzye traceState, they are separated by semicolon
+		// like pganalyze=t:1697666938.6297212;x=123
+		for _, part := range strings.Split(strings.TrimSpace(pganalyzeState), ";") {
+			if strings.Contains(part, ":") {
+				keyAndValue := strings.SplitN(part, ":", 2)
+				if strings.TrimSpace(keyAndValue[0]) == "t" {
+					if startInSec, err := strconv.ParseFloat(keyAndValue[1], 64); err == nil {
+						startSec := int64(startInSec)
+						startNanoSec := int64(startInSec*1e9) - (startSec * 1e9)
+						startTime = time.Unix(startSec, startNanoSec).UTC()
+						// With this, we're adding the query duration to the start time.
+						// This could result creating inaccurate spans, as the start time passed
+						// from the client side using tracestate is the time of the query is sent
+						// from the client to the server.
+						// This means, we will ignore the network time between the client and the
+						// server, as well as the machine clock different between them.
+						endTime = startTime.Add(time.Duration(sample.RuntimeMs) * time.Millisecond)
+						return
+					}
+				}
+			}
+		}
+	}
+	// Calculate start and end time based on sample data
+	duration := time.Duration(sample.RuntimeMs) * time.Millisecond
+	startTime = sample.OccurredAt.Add(-1 * duration)
+	endTime = sample.OccurredAt
+
+	return
+}
+
 func ExportQuerySamplesAsTraceSpans(ctx context.Context, server *state.Server, logger *util.Logger, grant state.GrantLogs, samples []state.PostgresQuerySample) {
 	exportCount := 0
 	for _, sample := range samples {
@@ -49,9 +86,7 @@ func ExportQuerySamplesAsTraceSpans(ctx context.Context, server *state.Server, l
 				trace.WithInstrumentationVersion(util.CollectorVersion),
 				trace.WithSchemaURL(semconv.SchemaURL),
 			)
-			duration := -1 * time.Duration(sample.RuntimeMs) * time.Millisecond
-			startTime := sample.OccurredAt.Add(duration)
-			endTime := sample.OccurredAt
+			startTime, endTime := startAndEndTime(trace.SpanContextFromContext(ctx).TraceState(), sample)
 			_, span := tracer.Start(ctx, otelSpanName, trace.WithTimestamp(startTime))
 			// See https://opentelemetry.io/docs/specs/otel/trace/semantic_conventions/database/
 			// however note that "db.postgresql.plan" is non-standard.

logs/querysample/tracing_test.go

@@ -0,0 +1,86 @@
+package querysample
+
+import (
+	"testing"
+	"time"
+
+	"github.com/pganalyze/collector/state"
+	"go.opentelemetry.io/otel/trace"
+)
+
+type startAndEndTimeTestPair struct {
+	testName   string
+	traceState trace.TraceState
+	sample     state.PostgresQuerySample
+	startTime  time.Time
+	endTime    time.Time
+}
+
+func TestStartAndEndTime(t *testing.T) {
+	currentTime, err := time.Parse("2006-01-02", "2023-01-01")
+	if err != nil {
+		t.Fatalf("Failed to initialize object: %v", err)
+	}
+	traceState := trace.TraceState{}
+	otelTraceState, err := traceState.Insert("ot", "p:8;r:62")
+	if err != nil {
+		t.Fatalf("Failed to initialize object: %v", err)
+	}
+	pganalyzeTraceStateWithoutT, err := traceState.Insert("pganalyze", "x:foo;y:bar")
+	if err != nil {
+		t.Fatalf("Failed to initialize object: %v", err)
+	}
+	// inserting the same key will update the value
+	pganalyzeTraceState, err := traceState.Insert("pganalyze", "t:1697666938.6297212")
+	if err != nil {
+		t.Fatalf("Failed to initialize object: %v", err)
+	}
+	// 1697666938.6297212 = 2023-10-18 22:08:58.6297212
+	pganalyzeTime, err := time.Parse("2006-01-02T15:04:05.999999999", "2023-10-18T22:08:58.6297212")
+	if err != nil {
+		t.Fatalf("Failed to initialize object: %v", err)
+	}
+	// due to the limitation of the floating point, the result won't exactly like above, so tweaking to pass the test
+	pganalyzeTime = pganalyzeTime.Add(-1 * 112)
+
+	var startAndEndTimeTests = []startAndEndTimeTestPair{
+		{
+			"No trace state",
+			trace.TraceState{},
+			state.PostgresQuerySample{RuntimeMs: 1000, OccurredAt: currentTime},
+			currentTime.Add(-1 * 1000 * time.Millisecond),
+			currentTime,
+		},
+		{
+			"No pganalyze trace state",
+			otelTraceState,
+			state.PostgresQuerySample{RuntimeMs: 1000, OccurredAt: currentTime},
+			currentTime.Add(-1 * 1000 * time.Millisecond),
+			currentTime,
+		},
+		{
+			"pganalyze trace state without t",
+			pganalyzeTraceStateWithoutT,
+			state.PostgresQuerySample{RuntimeMs: 1000, OccurredAt: currentTime},
+			currentTime.Add(-1 * 1000 * time.Millisecond),
+			currentTime,
+		},
+		{
+			"pganalyze trace state",
+			pganalyzeTraceState,
+			state.PostgresQuerySample{RuntimeMs: 1000, OccurredAt: currentTime},
+			pganalyzeTime,
+			pganalyzeTime.Add(1000 * time.Millisecond),
+		},
+	}
+
+	for _, pair := range startAndEndTimeTests {
+		startTime, endTime := startAndEndTime(pair.traceState, pair.sample)
+		if pair.startTime != startTime {
+			t.Errorf("For %s: expected startTime to be %v, but was %v\n", pair.testName, pair.startTime, startTime)
+		}
+		if pair.endTime != endTime {
+			t.Errorf("For %s: expected endTime to be %v, but was %v\n", pair.testName, pair.endTime, endTime)
+		}
+	}
+}