From 9bc7539f2b3f4908d13b91f5a17463ec71e214ce Mon Sep 17 00:00:00 2001
From: kirillgroshkov <kirill.groshkov@naturalcycles.com>
Date: Sun, 27 Oct 2024 12:51:21 +0100
Subject: [PATCH] chore: @_SwarmSafe (experimental)

---
 src/decorators/swarmSafe.decorator.test.ts | 49 ++++++++++++++++++++++
 src/decorators/swarmSafe.decorator.ts      | 47 +++++++++++++++++++++
 2 files changed, 96 insertions(+)
 create mode 100644 src/decorators/swarmSafe.decorator.test.ts
 create mode 100644 src/decorators/swarmSafe.decorator.ts

diff --git a/src/decorators/swarmSafe.decorator.test.ts b/src/decorators/swarmSafe.decorator.test.ts
new file mode 100644
index 00000000..cbff44aa
--- /dev/null
+++ b/src/decorators/swarmSafe.decorator.test.ts
@@ -0,0 +1,49 @@
+import { _range } from '../array/range'
+import { pDelay } from '../promise/pDelay'
+import { _SwarmSafe } from './swarmSafe.decorator'
+
+class C {
+  ranTimes = 0
+
+  @_SwarmSafe()
+  async run(n: number): Promise<number> {
+    this.ranTimes++
+    await pDelay(1)
+    return n * 2
+  }
+}
+
+test('swarmSafe sequential', async () => {
+  const c = new C()
+
+  for (const _ of _range(10)) {
+    const r = await c.run(1)
+    expect(r).toBe(2)
+  }
+
+  expect(c.ranTimes).toBe(10)
+})
+
+test('swarmSafe parallel', async () => {
+  const c = new C()
+
+  const [r1, r2, r3] = await Promise.all([
+    c.run(1),
+    c.run(2),
+    c.run(3),
+    c.run(4),
+    c.run(5),
+    c.run(6),
+    c.run(7),
+  ])
+
+  expect(c.ranTimes).toBe(1)
+  expect(r1).toBe(2)
+  expect(r2).toBe(2)
+  expect(r3).toBe(2)
+
+  // Then, run again
+  const r = await c.run(2)
+  expect(r).toBe(4)
+  expect(c.ranTimes).toBe(2)
+})
diff --git a/src/decorators/swarmSafe.decorator.ts b/src/decorators/swarmSafe.decorator.ts
new file mode 100644
index 00000000..5d1ac966
--- /dev/null
+++ b/src/decorators/swarmSafe.decorator.ts
@@ -0,0 +1,47 @@
+import { AnyObject } from '../types'
+import { _getTargetMethodSignature } from './decorator.util'
+
+/**
+ * Prevents "swarm" of async calls to the same method.
+ * Allows max 1 in-flight promise to exist.
+ * If more calls appear, while Promise is not resolved yet - same Promise is returned.
+ *
+ * Does not support `cacheKey`.
+ * So, the same Promise is returned, regardless of the arguments.
+ */
+// eslint-disable-next-line @typescript-eslint/naming-convention
+export const _SwarmSafe = (): MethodDecorator => (target, key, descriptor) => {
+  if (typeof descriptor.value !== 'function') {
+    throw new TypeError('@_SwarmSafe can be applied only to methods')
+  }
+
+  const originalFn = descriptor.value
+  const keyStr = String(key)
+  const methodSignature = _getTargetMethodSignature(target, keyStr)
+  const instanceCache = new Map<AnyObject, Promise<any>>()
+
+  console.log('SwarmSafe constructor called', { key, methodSignature })
+
+  // eslint-disable-next-line @typescript-eslint/promise-function-async
+  descriptor.value = function (this: typeof target, ...args: any[]): Promise<any> {
+    console.log('SwarmSafe method called', { key, methodSignature, args })
+    const ctx = this
+
+    let inFlightPromise = instanceCache.get(ctx)
+    if (inFlightPromise) {
+      console.log(`SwarmSafe: returning in-flight promise`)
+      return inFlightPromise
+    }
+
+    console.log(`SwarmSafe: first-time call, creating in-flight promise`)
+
+    inFlightPromise = originalFn.apply(ctx, args) as Promise<any>
+    instanceCache.set(ctx, inFlightPromise)
+    void inFlightPromise.finally(() => {
+      console.log(`SwarmSafe: in-flight promise resolved`)
+      instanceCache.delete(ctx)
+    })
+
+    return inFlightPromise
+  } as any
+}