From e06355741daf9abf559eba2182d6210c0496bc44 Mon Sep 17 00:00:00 2001
From: Daniel Fangl <daniel.fangl@localstack.cloud>
Date: Tue, 6 Feb 2024 10:51:37 +0100
Subject: [PATCH] Add new localstack tracer to return tracing headers set by
 invoke (#31)

---
 cmd/localstack/main.go   |  4 ++-
 cmd/localstack/tracer.go | 63 ++++++++++++++++++++++++++++++++++++++++
 2 files changed, 66 insertions(+), 1 deletion(-)
 create mode 100644 cmd/localstack/tracer.go

diff --git a/cmd/localstack/main.go b/cmd/localstack/main.go
index f03e3b6..02bf247 100644
--- a/cmd/localstack/main.go
+++ b/cmd/localstack/main.go
@@ -178,6 +178,7 @@ func main() {
 
 	logCollector := NewLogCollector()
 	localStackLogsEgressApi := NewLocalStackLogsEgressAPI(logCollector)
+	tracer := NewLocalStackTracer()
 
 	// build sandbox
 	sandbox := rapidcore.
@@ -191,7 +192,8 @@ func main() {
 		}).
 		SetExtensionsFlag(true).
 		SetInitCachingFlag(true).
-		SetLogsEgressAPI(localStackLogsEgressApi)
+		SetLogsEgressAPI(localStackLogsEgressApi).
+		SetTracer(tracer)
 
 	// xray daemon
 	endpoint := "http://" + lsOpts.LocalstackIP + ":" + lsOpts.EdgePort
diff --git a/cmd/localstack/tracer.go b/cmd/localstack/tracer.go
new file mode 100644
index 0000000..8506a9a
--- /dev/null
+++ b/cmd/localstack/tracer.go
@@ -0,0 +1,63 @@
+package main
+
+import (
+	"context"
+	"encoding/json"
+	"go.amzn.com/lambda/appctx"
+	"go.amzn.com/lambda/interop"
+)
+
+type LocalStackTracer struct {
+	invoke *interop.Invoke
+}
+
+func (t *LocalStackTracer) Configure(invoke *interop.Invoke) {
+	t.invoke = invoke
+}
+
+func (t *LocalStackTracer) CaptureInvokeSegment(ctx context.Context, criticalFunction func(context.Context) error) error {
+	return criticalFunction(ctx)
+}
+
+func (t *LocalStackTracer) CaptureInitSubsegment(ctx context.Context, criticalFunction func(context.Context) error) error {
+	return criticalFunction(ctx)
+}
+
+func (t *LocalStackTracer) CaptureInvokeSubsegment(ctx context.Context, criticalFunction func(context.Context) error) error {
+	return criticalFunction(ctx)
+}
+
+func (t *LocalStackTracer) CaptureOverheadSubsegment(ctx context.Context, criticalFunction func(context.Context) error) error {
+	return criticalFunction(ctx)
+}
+
+func (t *LocalStackTracer) RecordInitStartTime()                                             {}
+func (t *LocalStackTracer) RecordInitEndTime()                                               {}
+func (t *LocalStackTracer) SendInitSubsegmentWithRecordedTimesOnce(ctx context.Context)      {}
+func (t *LocalStackTracer) SendRestoreSubsegmentWithRecordedTimesOnce(ctx context.Context)   {}
+func (t *LocalStackTracer) MarkError(ctx context.Context)                                    {}
+func (t *LocalStackTracer) AttachErrorCause(ctx context.Context, errorCause json.RawMessage) {}
+
+func (t *LocalStackTracer) WithErrorCause(ctx context.Context, appCtx appctx.ApplicationContext, criticalFunction func(ctx context.Context) error) func(ctx context.Context) error {
+	return criticalFunction
+}
+func (t *LocalStackTracer) WithError(ctx context.Context, appCtx appctx.ApplicationContext, criticalFunction func(ctx context.Context) error) func(ctx context.Context) error {
+	return criticalFunction
+}
+func (t *LocalStackTracer) BuildTracingHeader() func(context.Context) string {
+	// extract root trace ID and parent from context and build the tracing header
+	return func(ctx context.Context) string {
+		return t.invoke.TraceID
+	}
+}
+
+func (t *LocalStackTracer) BuildTracingCtxForStart() *interop.TracingCtx {
+	return nil
+}
+func (t *LocalStackTracer) BuildTracingCtxAfterInvokeComplete() *interop.TracingCtx {
+	return nil
+}
+
+func NewLocalStackTracer() *LocalStackTracer {
+	return &LocalStackTracer{}
+}