feat: Utility for passing debug correlation id to query implementations

Author: johanandren

akka-persistence-query/src/main/scala/akka/persistence/query/QueryCorrelationId.scala

@@ -0,0 +1,39 @@
+/*
+ * Copyright (C) 2009-2025 Lightbend Inc. <https://www.lightbend.com>
+ */
+
+package akka.persistence.query
+
+/**
+ * (Optional) mechanism for query implementations to pick up a correlation id from the caller, to use in logging and
+ * error messages. Used by akka-projections to make correlating projection logs with debug and trace logging from the
+ * underlying akka persistence query implementations possible.
+ */
+object QueryCorrelationId {
+
+  private val threadLocal = new ThreadLocal[String]
+
+  /**
+   * Expected to be used "around" calls to plugin query method, will clear the correlation id from thread local
+   * to make sure there is no leak between logic executed on shared threads.
+   *
+   * @param correlationId
+   */
+  def withCorrelationId[T](correlationId: String)(block: () => T): T = {
+    threadLocal.set(correlationId)
+    try {
+      block()
+    } finally {
+      threadLocal.remove()
+    }
+  }
+
+  /**
+   * @return Expected to be called directly after receiving a query call, before starting any asynchronous tasks,
+   *         returns and clears out the correlation id to make sure there is no leak between tasks. Further passing
+   *         around of the uuid inside the query plugin implementation is up to the implementer.
+   */
+  def get(): Option[String] =
+    Option(threadLocal.get)
+
+}

akka-persistence-query/src/test/scala/akka/persistence/query/QueryCorrelationIdSpec.scala

@@ -0,0 +1,46 @@
+/*
+ * Copyright (C) 2009-2025 Lightbend Inc. <https://www.lightbend.com>
+ */
+
+package akka.persistence.query
+
+import akka.testkit.TestException
+import org.scalatest.matchers.should.Matchers
+import org.scalatest.wordspec.AnyWordSpecLike
+
+import java.util.UUID
+
+class QueryCorrelationIdSpec extends AnyWordSpecLike with Matchers {
+
+  def pretendQueryMethod(): Option[String] =
+    QueryCorrelationId.get()
+
+  "The query correlation id utility" should {
+
+    "pass and clear correlation id" in {
+      val uuid = UUID.randomUUID().toString
+      val observed =
+        QueryCorrelationId.withCorrelationId(uuid) { () =>
+          pretendQueryMethod()
+        }
+      observed shouldEqual Some(uuid)
+
+      // cleared after returning
+      QueryCorrelationId.get() shouldBe None
+    }
+
+    "clear correlation id when call fails" in {
+      val uuid = UUID.randomUUID().toString
+      intercept[TestException] {
+        QueryCorrelationId.withCorrelationId(uuid) { () =>
+          throw TestException("boom")
+        }
+      }
+
+      // cleared after throwing
+      QueryCorrelationId.get() shouldBe None
+    }
+
+  }
+
+}