From a7087025e51a1f6a314481c7e78786546fca7d93 Mon Sep 17 00:00:00 2001
From: Tim <hello@timsmart.co>
Date: Fri, 24 May 2024 12:17:42 +1200
Subject: [PATCH] add Layer.annotateLogs & Layer.annotateSpans (#2836)

---
 .changeset/purple-mangos-grin.md      | 17 ++++++++
 packages/effect/src/Layer.ts          | 22 ++++++++++
 packages/effect/src/internal/layer.ts | 63 +++++++++++++++++++++++++++
 3 files changed, 102 insertions(+)
 create mode 100644 .changeset/purple-mangos-grin.md

diff --git a/.changeset/purple-mangos-grin.md b/.changeset/purple-mangos-grin.md
new file mode 100644
index 0000000000..1e4f640331
--- /dev/null
+++ b/.changeset/purple-mangos-grin.md
@@ -0,0 +1,17 @@
+---
+"effect": minor
+---
+
+add Layer.annotateLogs & Layer.annotateSpans
+
+This allows you to add log & span annotation to a Layer.
+
+```ts
+import { Effect, Layer } from "effect";
+
+Layer.effectDiscard(Effect.log("hello")).pipe(
+  Layer.annotateLogs({
+    service: "my-service",
+  }),
+);
+```
diff --git a/packages/effect/src/Layer.ts b/packages/effect/src/Layer.ts
index e18181a1a8..0e6a1ebfe4 100644
--- a/packages/effect/src/Layer.ts
+++ b/packages/effect/src/Layer.ts
@@ -139,6 +139,28 @@ export const isLayer: (u: unknown) => u is Layer<unknown, unknown, unknown> = in
  */
 export const isFresh: <RIn, E, ROut>(self: Layer<ROut, E, RIn>) => boolean = internal.isFresh
 
+/**
+ * @since 3.3.0
+ * @category tracing
+ */
+export const annotateLogs: {
+  (key: string, value: unknown): <A, E, R>(self: Layer<A, E, R>) => Layer<A, E, R>
+  (values: Record<string, unknown>): <A, E, R>(self: Layer<A, E, R>) => Layer<A, E, R>
+  <A, E, R>(self: Layer<A, E, R>, key: string, value: unknown): Layer<A, E, R>
+  <A, E, R>(self: Layer<A, E, R>, values: Record<string, unknown>): Layer<A, E, R>
+} = internal.annotateLogs
+
+/**
+ * @since 3.3.0
+ * @category tracing
+ */
+export const annotateSpans: {
+  (key: string, value: unknown): <A, E, R>(self: Layer<A, E, R>) => Layer<A, E, R>
+  (values: Record<string, unknown>): <A, E, R>(self: Layer<A, E, R>) => Layer<A, E, R>
+  <A, E, R>(self: Layer<A, E, R>, key: string, value: unknown): Layer<A, E, R>
+  <A, E, R>(self: Layer<A, E, R>, values: Record<string, unknown>): Layer<A, E, R>
+} = internal.annotateSpans
+
 /**
  * Builds a layer into a scoped value.
  *
diff --git a/packages/effect/src/internal/layer.ts b/packages/effect/src/internal/layer.ts
index aec5500891..a1c61b6dba 100644
--- a/packages/effect/src/internal/layer.ts
+++ b/packages/effect/src/internal/layer.ts
@@ -8,6 +8,7 @@ import type { FiberRef } from "../FiberRef.js"
 import * as FiberRefsPatch from "../FiberRefsPatch.js"
 import type { LazyArg } from "../Function.js"
 import { dual, pipe } from "../Function.js"
+import * as HashMap from "../HashMap.js"
 import type * as Layer from "../Layer.js"
 import { pipeArguments } from "../Pipeable.js"
 import { hasProperty } from "../Predicate.js"
@@ -1108,10 +1109,72 @@ export const unwrapScoped = <A, E1, R1, E, R>(
   return flatMap(scoped(tag, self), (context) => Context.get(context, tag))
 }
 
+// -----------------------------------------------------------------------------
+// logging
+// -----------------------------------------------------------------------------
+
+export const annotateLogs = dual<
+  {
+    (key: string, value: unknown): <A, E, R>(self: Layer.Layer<A, E, R>) => Layer.Layer<A, E, R>
+    (
+      values: Record<string, unknown>
+    ): <A, E, R>(self: Layer.Layer<A, E, R>) => Layer.Layer<A, E, R>
+  },
+  {
+    <A, E, R>(self: Layer.Layer<A, E, R>, key: string, value: unknown): Layer.Layer<A, E, R>
+    <A, E, R>(self: Layer.Layer<A, E, R>, values: Record<string, unknown>): Layer.Layer<A, E, R>
+  }
+>(
+  (args) => isLayer(args[0]),
+  function<A, E, R>() {
+    const args = arguments
+    return fiberRefLocallyWith(
+      args[0] as Layer.Layer<A, E, R>,
+      core.currentLogAnnotations,
+      typeof args[1] === "string"
+        ? HashMap.set(args[1], args[2])
+        : (annotations) =>
+          Object.entries(args[1] as Record<string, unknown>).reduce(
+            (acc, [key, value]) => HashMap.set(acc, key, value),
+            annotations
+          )
+    )
+  }
+)
+
 // -----------------------------------------------------------------------------
 // tracing
 // -----------------------------------------------------------------------------
 
+export const annotateSpans = dual<
+  {
+    (key: string, value: unknown): <A, E, R>(self: Layer.Layer<A, E, R>) => Layer.Layer<A, E, R>
+    (
+      values: Record<string, unknown>
+    ): <A, E, R>(self: Layer.Layer<A, E, R>) => Layer.Layer<A, E, R>
+  },
+  {
+    <A, E, R>(self: Layer.Layer<A, E, R>, key: string, value: unknown): Layer.Layer<A, E, R>
+    <A, E, R>(self: Layer.Layer<A, E, R>, values: Record<string, unknown>): Layer.Layer<A, E, R>
+  }
+>(
+  (args) => isLayer(args[0]),
+  function<A, E, R>() {
+    const args = arguments
+    return fiberRefLocallyWith(
+      args[0] as Layer.Layer<A, E, R>,
+      core.currentTracerSpanAnnotations,
+      typeof args[1] === "string"
+        ? HashMap.set(args[1], args[2])
+        : (annotations) =>
+          Object.entries(args[1] as Record<string, unknown>).reduce(
+            (acc, [key, value]) => HashMap.set(acc, key, value),
+            annotations
+          )
+    )
+  }
+)
+
 /** @internal */
 export const withSpan: {
   (