From 584c4f5a49651aafb42778d300aa2fa2ddc45771 Mon Sep 17 00:00:00 2001
From: Kyle Mathews <mathews.kyle@gmail.com>
Date: Thu, 21 Nov 2024 01:20:07 -0700
Subject: [PATCH] feat: use traceparent header from incoming shape requests to
 set parent span (#2000)

Fix https://github.com/electric-sql/electric/issues/1914

- https://github.com/open-telemetry/opentelemetry-erlang/discussions/518

---------

Co-authored-by: Ilia Borovitinov <ilia@electric-sql.com>
---
 .changeset/odd-walls-walk.md                  |  5 ++++
 .../sync-service/lib/electric/plug/router.ex  |  1 +
 .../lib/electric/plug/trace_context_plug.ex   | 28 +++++++++++++++++++
 3 files changed, 34 insertions(+)
 create mode 100644 .changeset/odd-walls-walk.md
 create mode 100644 packages/sync-service/lib/electric/plug/trace_context_plug.ex

diff --git a/.changeset/odd-walls-walk.md b/.changeset/odd-walls-walk.md
new file mode 100644
index 0000000000..9552865a83
--- /dev/null
+++ b/.changeset/odd-walls-walk.md
@@ -0,0 +1,5 @@
+---
+"@core/sync-service": patch
+---
+
+use traceparent header from incoming shape requests to set parent span
diff --git a/packages/sync-service/lib/electric/plug/router.ex b/packages/sync-service/lib/electric/plug/router.ex
index 126c9ded4d..4cd912a69f 100644
--- a/packages/sync-service/lib/electric/plug/router.ex
+++ b/packages/sync-service/lib/electric/plug/router.ex
@@ -10,6 +10,7 @@ defmodule Electric.Plug.Router do
   plug Plug.Head
   plug :match
   plug Electric.Plug.LabelProcessPlug
+  plug Electric.Plug.TraceContextPlug
   plug Plug.Telemetry, event_prefix: [:electric, :routing]
   plug Plug.Logger
   plug :put_cors_headers
diff --git a/packages/sync-service/lib/electric/plug/trace_context_plug.ex b/packages/sync-service/lib/electric/plug/trace_context_plug.ex
new file mode 100644
index 0000000000..180f8362b8
--- /dev/null
+++ b/packages/sync-service/lib/electric/plug/trace_context_plug.ex
@@ -0,0 +1,28 @@
+defmodule Electric.Plug.TraceContextPlug do
+  @moduledoc """
+  A plug that extracts trace context from incoming HTTP headers and sets it as the parent span.
+  """
+  @behaviour Plug
+
+  require Logger
+
+  def init(opts), do: opts
+
+  def call(%Plug.Conn{req_headers: headers} = conn, _opts) do
+    # Extract function expects a list of headers as tuples and knows
+    # the expected header keys, so we don't have to prefilter.
+    ctx = :otel_propagator_text_map.extract_to(:otel_ctx.new(), headers)
+
+    # Get the span context from the extracted context
+    case :otel_tracer.current_span_ctx(ctx) do
+      :undefined ->
+        # No parent, continue as-is
+        conn
+
+      span_ctx ->
+        # Parent found, set as current span
+        :otel_tracer.set_current_span(span_ctx)
+        conn
+    end
+  end
+end