From ffe987e1bcffb9ce15d161eb9dd609bbfbd8ed88 Mon Sep 17 00:00:00 2001
From: Tim <hello@timsmart.co>
Date: Fri, 26 Apr 2024 11:32:33 +1200
Subject: [PATCH] add Effect.annotateLogsScoped (#2618)

---
 .changeset/fifty-readers-battle.md           | 17 +++++++++++++++
 packages/effect/src/Effect.ts                | 22 +++++++++++++++++++
 packages/effect/src/internal/fiberRuntime.ts | 23 ++++++++++++++++++++
 packages/effect/test/FiberRefs.test.ts       | 14 ++++++++++++
 4 files changed, 76 insertions(+)
 create mode 100644 .changeset/fifty-readers-battle.md

diff --git a/.changeset/fifty-readers-battle.md b/.changeset/fifty-readers-battle.md
new file mode 100644
index 0000000000..c7bdc54903
--- /dev/null
+++ b/.changeset/fifty-readers-battle.md
@@ -0,0 +1,17 @@
+---
+"effect": minor
+---
+
+add Effect.annotateLogsScoped
+
+This api allows you to annotate logs until the Scope has been closed.
+
+```ts
+import { Effect } from "effect"
+
+Effect.gen(function* () {
+  yield* Effect.log("no annotations")
+  yield* Effect.annotateLogsScoped({ foo: "bar" })
+  yield* Effect.log("annotated with foo=bar")
+}).pipe(Effect.scoped, Effect.andThen(Effect.log("no annotations again")))
+```
diff --git a/packages/effect/src/Effect.ts b/packages/effect/src/Effect.ts
index 8c6a67d8ac..90c9c42a39 100644
--- a/packages/effect/src/Effect.ts
+++ b/packages/effect/src/Effect.ts
@@ -4452,6 +4452,28 @@ export const annotateLogs: {
   <A, E, R>(effect: Effect<A, E, R>, values: Record<string, unknown>): Effect<A, E, R>
 } = effect.annotateLogs
 
+/**
+ * Annotates each log with the specified log annotation(s), until the Scope is closed.
+ *
+ * @since 3.1.0
+ * @category logging
+ * @example
+ * import { Effect } from "effect"
+ *
+ * Effect.gen(function*() {
+ *   yield* Effect.log("no annotations")
+ *   yield* Effect.annotateLogsScoped({ foo: "bar" })
+ *   yield* Effect.log("annotated with foo=bar")
+ * }).pipe(
+ *   Effect.scoped,
+ *   Effect.andThen(Effect.log("no annotations again"))
+ * )
+ */
+export const annotateLogsScoped: {
+  (key: string, value: unknown): Effect<void, never, Scope.Scope>
+  (values: Record<string, unknown>): Effect<void, never, Scope.Scope>
+} = fiberRuntime.annotateLogsScoped
+
 /**
  * Retrieves the log annotations associated with the current scope.
  *
diff --git a/packages/effect/src/internal/fiberRuntime.ts b/packages/effect/src/internal/fiberRuntime.ts
index 0288bef698..15fd069f90 100644
--- a/packages/effect/src/internal/fiberRuntime.ts
+++ b/packages/effect/src/internal/fiberRuntime.ts
@@ -1502,6 +1502,29 @@ export const batchedLogger = dual<
     )
   }))
 
+export const annotateLogsScoped: {
+  (key: string, value: unknown): Effect.Effect<void, never, Scope.Scope>
+  (values: Record<string, unknown>): Effect.Effect<void, never, Scope.Scope>
+} = function() {
+  if (typeof arguments[0] === "string") {
+    return fiberRefLocallyScopedWith(
+      core.currentLogAnnotations,
+      HashMap.set(arguments[0], arguments[1])
+    )
+  }
+  const entries = Object.entries(arguments[0])
+  return fiberRefLocallyScopedWith(
+    core.currentLogAnnotations,
+    HashMap.mutate((annotations) => {
+      for (let i = 0; i < entries.length; i++) {
+        const [key, value] = entries[i]
+        HashMap.set(annotations, key, value)
+      }
+      return annotations
+    })
+  )
+}
+
 // circular with Effect
 
 /* @internal */
diff --git a/packages/effect/test/FiberRefs.test.ts b/packages/effect/test/FiberRefs.test.ts
index 3c958cf8af..4f8cff861b 100644
--- a/packages/effect/test/FiberRefs.test.ts
+++ b/packages/effect/test/FiberRefs.test.ts
@@ -1,6 +1,7 @@
 import * as it from "effect-test/utils/extend"
 import * as Cause from "effect/Cause"
 import * as Effect from "effect/Effect"
+import * as Exit from "effect/Exit"
 import * as Fiber from "effect/Fiber"
 import * as FiberId from "effect/FiberId"
 import * as FiberRef from "effect/FiberRef"
@@ -8,6 +9,7 @@ import * as FiberRefs from "effect/FiberRefs"
 import * as HashMap from "effect/HashMap"
 import * as Option from "effect/Option"
 import * as Queue from "effect/Queue"
+import * as Scope from "effect/Scope"
 import { assert, describe, expect } from "vitest"
 
 describe("FiberRefs", () => {
@@ -48,5 +50,17 @@ describe("FiberRefs", () => {
       Effect.void.pipe(Effect.annotateLogs("test", "abc"), Effect.runSync)
       expect(FiberRef.currentLogAnnotations.pipe(FiberRef.get, Effect.map(HashMap.size), Effect.runSync)).toBe(0)
     })
+
+    it.effect("annotateLogsScoped", () =>
+      Effect.gen(function*() {
+        const scope = yield* Scope.make()
+        assert.strictEqual(HashMap.size(yield* FiberRef.get(FiberRef.currentLogAnnotations)), 0)
+        yield Effect.annotateLogsScoped({
+          test: 123
+        }).pipe(Scope.extend(scope))
+        assert.strictEqual(HashMap.size(yield* FiberRef.get(FiberRef.currentLogAnnotations)), 1)
+        yield Scope.close(scope, Exit.void)
+        assert.strictEqual(HashMap.size(yield* FiberRef.get(FiberRef.currentLogAnnotations)), 0)
+      }))
   })
 })