From b7dea5c015ac80f9d5434c061de9af7c4990ec39 Mon Sep 17 00:00:00 2001
From: hasundue <hasundue@gmail.com>
Date: Fri, 22 Mar 2024 11:50:20 +0900
Subject: [PATCH] feat(std/relays): `WithPool` to create relay pools

---
 std/relays.ts      | 39 +++++++++++++++++++++++++++++++++++++++
 std/relays_test.ts | 36 +++++++++++++++++++++++++++++++++++-
 2 files changed, 74 insertions(+), 1 deletion(-)

diff --git a/std/relays.ts b/std/relays.ts
index 3481945..9ca0f56 100644
--- a/std/relays.ts
+++ b/std/relays.ts
@@ -4,12 +4,15 @@ import type {
   ClientToRelayMessage,
   EventKind,
   NostrEvent,
+  RelayUrl,
   SubscriptionFilter,
 } from "@lophus/core/protocol";
 import {
+  Relay,
   RelayLike,
   RelayLikeConfig,
   RelayLikeOptions,
+  RelayOptions,
   SubscriptionOptions,
 } from "@lophus/core/relays";
 
@@ -68,3 +71,39 @@ export class RelayGroup implements RelayLike {
     await Promise.resolve();
   }
 }
+
+export interface WithPool<
+  R extends typeof Relay,
+> {
+  pool: Map<RelayUrl, InstanceType<R>>;
+  new (url: RelayUrl, options?: RelayOptions): InstanceType<R>;
+}
+
+export function WithPool<
+  R extends typeof Relay,
+>(
+  BaseRelay: R,
+): WithPool<R> {
+  // @ts-ignore allow concrete arguments for constructor
+  return class Self extends BaseRelay {
+    static readonly pool = new Map<RelayUrl, InstanceType<R>>();
+
+    constructor(
+      url: RelayUrl,
+      options?: RelayOptions,
+    ) {
+      const pooled = Self.pool.get(url);
+      if (pooled) {
+        return pooled;
+      }
+      super(url, options);
+      Self.pool.set(url, this as InstanceType<R>);
+      return this;
+    }
+
+    override close() {
+      Self.pool.delete(this.config.url);
+      return super.close();
+    }
+  };
+}
diff --git a/std/relays_test.ts b/std/relays_test.ts
index 1f82c77..a67914e 100644
--- a/std/relays_test.ts
+++ b/std/relays_test.ts
@@ -6,7 +6,7 @@ import {
 import { afterAll, beforeAll, describe, it } from "@std/testing/bdd";
 import { MockWebSocket } from "@lophus/lib/testing";
 import { NostrEvent, Relay, SubscriptionId } from "@lophus/nips";
-import { RelayGroup } from "./relays.ts";
+import { RelayGroup, WithPool } from "./relays.ts";
 
 describe("RelayGroup", () => {
   let relays: Relay[];
@@ -113,3 +113,37 @@ describe("RelayGroup", () => {
     });
   });
 });
+
+describe("WithPool", () => {
+  let Pooled: WithPool<typeof Relay>;
+
+  beforeAll(() => {
+    globalThis.WebSocket = MockWebSocket;
+  });
+
+  it("should accept a NIP-enabled relay as an argument", () => {
+    Pooled = WithPool(Relay);
+  });
+
+  it("should have no relays in the pool initially", () => {
+    assertEquals(Pooled.pool.size, 0);
+  });
+
+  it("should add a relay to the pool and return it", () => {
+    const relay = new Pooled("ws://localhost:80");
+    assertEquals(Pooled.pool.size, 1);
+    assertEquals(Pooled.pool.has(relay.config.url), true);
+  });
+
+  it("should return the pooled relay if it exists", () => {
+    const relay = new Pooled("ws://localhost:80");
+    assertEquals(Pooled.pool.size, 1);
+    assertEquals(Pooled.pool.has(relay.config.url), true);
+  });
+
+  it("should remove a relay from the pool when it is closed", async () => {
+    const relay = new Pooled("ws://localhost:80");
+    await relay.close();
+    assertEquals(Pooled.pool.size, 0);
+  });
+});