@pluv/io 0.40.2 → 0.41.1

package/src/IORoom.ts

@@ -2,6 +2,7 @@ import type { AbstractCrdtDoc, AbstractCrdtDocFactory } from "@pluv/crdt";
 import { noop } from "@pluv/crdt";
 import type {
     BaseIOEventRecord,
+    BaseUser,
     EventMessage,
     IOEventMessage,
     IOLike,
@@ -25,7 +26,7 @@ import type { AbstractCloseEvent, AbstractErrorEvent, AbstractMessageEvent } fro
 import { AbstractWebSocket } from "./AbstractWebSocket";
 import type { PluvRouter, PluvRouterEventConfig } from "./PluvRouter";
 import { authorize } from "./authorize";
-import { PING_TIMEOUT_MS } from "./constants";
+import { GARBAGE_COLLECT_INTERVAL_MS, PING_TIMEOUT_MS } from "./constants";
 import type {
     EventResolverContext,
     EventResolverKind,
@@ -87,6 +88,12 @@ interface SendMessageSender {
     user: JsonObject | null;
 }
 
+interface PatchPresenceParams {
+    presence: JsonObject | null;
+    sessionId: string;
+    timer?: number | null;
+}
+
 export type WebSocketRegisterConfig<TPlatform extends AbstractPlatform<any> = AbstractPlatform<any>> = {
     token?: string | null;
 } & InferInitContextType<TPlatform>;
@@ -101,6 +108,7 @@ export class IORoom<
     public readonly id: string;
 
     private _doc: Promise<AbstractCrdtDoc<any>>;
+    private _lastGarbageCollectMs: number = -1 * (GARBAGE_COLLECT_INTERVAL_MS + 1);
     private _uninitialize: Promise<() => Promise<void>> | null = null;
 
     private readonly _authorize: TAuthorize = null as TAuthorize;
@@ -110,7 +118,8 @@ export class IORoom<
     private readonly _listeners: IORoomListeners<TPlatform, TAuthorize, TContext, TEvents>;
     private readonly _platform: TPlatform;
     private readonly _router: PluvRouter<TPlatform, TAuthorize, TContext, TEvents>;
-    private readonly _sessions = new Map<string, AbstractWebSocket>();
+    private readonly _sessions = new Map<[sessionId: string][0], AbstractWebSocket<any, TAuthorize>>();
+    private readonly _userSessionss = new Map<[userId: string][0], Set<[sessionId: string][0]>>();
 
     /**
      * @ignore
@@ -185,6 +194,11 @@ export class IORoom<
 
         if (authorize) this._authorize = authorize;
 
+        /**
+         * @description Everything below this line relates to waking up from hibernation.
+         * @see https://developers.cloudflare.com/durable-objects/best-practices/websockets/#websocket-hibernation-api
+         * @date April 17, 2025
+         */
         const webSockets = this._platform.getWebSockets() as readonly InferPlatformWebSocketSource<TPlatform>[];
 
         webSockets.forEach((webSocket) => {
@@ -197,6 +211,10 @@ export class IORoom<
             const pluvWs = this._platform.convertWebSocket(webSocket, { room: this.id });
 
             this._sessions.set(sessionId, pluvWs);
+
+            const userId = pluvWs.session.user?.id;
+
+            if (!!userId) this._addUserSession(userId, sessionId);
         });
 
         const { doc, uninitialize } = this._initialize();
@@ -231,6 +249,15 @@ export class IORoom<
         await this._emitQuitters();
     }
 
+    /**
+     * @description The IORoom will garbage collect occasionally as connections ping/pong the
+     * server
+     */
+    public async garbageCollect(): Promise<void> {
+        this._lastGarbageCollectMs = Date.now();
+        this._emitQuitters();
+    }
+
     public getSize(): number {
         const currentTime = new Date().getTime();
 
@@ -286,42 +313,54 @@ export class IORoom<
     ): Promise<void> {
         const _options = (options[0] ?? {}) as WebSocketRegisterConfig<TPlatform>;
         const token = _options.token ?? null;
+
+        /**
+         * TODO
+         * @description Update the sessionId strategy so that the userId can be used to retrieve
+         * the session in O(1) time if authorized.
+         * @date April 15, 2025
+         */
         const sessionId = this._platform.getSessionId(webSocket);
         const sessionExists = typeof sessionId === "string" && this._sessions.has(sessionId);
 
         if (sessionExists) return;
+
         if (!(await this._initialized)) {
             this._initialize();
-
             await this._initialized;
         }
 
-        const user = await this._getAuthorizedUser(token, _options);
+        const user: BaseUser | null = await this._getAuthorizedUser(token, _options);
         const ioAuthorize = this._getIOAuthorize(_options);
         const isUnauthorized = !!ioAuthorize && !user;
-
-        // There is an attempt for multiple of the same identities. Probably malicious.
-        const isTokenInUse: boolean =
-            !!user &&
-            (await this._getSessions().then((sessions) => sessions.some((session) => session.user?.id === user.id)));
-
         const pluvWs = this._platform.convertWebSocket(webSocket, { room: this.id });
 
-        this._logDebug(`${colors.blue(`Registering connection for room ${this.id}:`)} ${pluvWs.sessionId}`);
-
-        await this._platform.acceptWebSocket(pluvWs);
-
-        if (isUnauthorized || isTokenInUse) {
-            this._logDebug(`${colors.blue("Authorization failed for connection:")} ${pluvWs.sessionId}`);
-
+        if (isUnauthorized) {
+            this._logDebug(`${colors.blue("Authorization failed for connection")}`);
             pluvWs.handleError({ error: new Error("Not authorized"), room: this.id });
             pluvWs.close(3000, "WebSocket unauthorized.");
 
             return;
         }
 
-        this._sessions.set(pluvWs.sessionId, pluvWs);
+        if (!!user) {
+            const latest = this._getLatestPresence(user.id);
+            const prevState = pluvWs.state;
+
+            pluvWs.user = user;
+
+            this._platform.setSerializedState(pluvWs, {
+                ...prevState,
+                presence: latest.presence,
+                timers: { ...prevState.timers, presence: latest.timer },
+            });
+            this._addUserSession(user.id, pluvWs.sessionId);
+        }
+
+        this._logDebug(`${colors.blue(`Registering connection for room ${this.id}:`)} ${pluvWs.sessionId}`);
 
+        await this._platform.acceptWebSocket(pluvWs);
+        this._sessions.set(pluvWs.sessionId, pluvWs);
         await this._platform.persistence.addUser(this.id, pluvWs.sessionId, user ?? {});
 
         if (this._platform._config.registrationMode === "attached") {
@@ -335,18 +374,26 @@ export class IORoom<
 
         await this._emitRegistered(pluvWs);
 
-        this._logDebug(`${colors.blue(`Registered connection for room ${this.id}:`)} ${pluvWs.sessionId}`);
-
         const size = this.getSize();
 
+        this._logDebug(`${colors.blue(`Registered connection for room ${this.id}:`)} ${pluvWs.sessionId}`);
         this._logDebug(`${colors.blue(`Room ${this.id} size:`)} ${size}`);
     }
 
+    private _addUserSession(userId: string, sessionId: string): Set<[sessionId: string][0]> {
+        const set = this._userSessionss.get(userId) ?? new Set<string>();
+        const updated = set.add(sessionId);
+
+        this._userSessionss.set(userId, updated);
+
+        return updated;
+    }
+
     private async _broadcast(params: BroadcastParams<this>): Promise<void> {
         const { message, senderId } = params;
 
         const sender = senderId ? (this._sessions.get(senderId) ?? null) : null;
-        const session = (await sender?.getSession()) ?? null;
+        const session = sender?.session ?? null;
         const user = session?.user ?? null;
 
         if (typeof senderId !== "string") return;
@@ -359,8 +406,8 @@ export class IORoom<
         });
     }
 
-    private async _closeWebSockets(webSockets: readonly AbstractWebSocket[]): Promise<void> {
-        const closeWebSocket = async (webSocket: AbstractWebSocket): Promise<void> => {
+    private async _closeWebSockets(webSockets: readonly AbstractWebSocket<any, TAuthorize>[]): Promise<void> {
+        const closeWebSocket = async (webSocket: AbstractWebSocket<any, TAuthorize>): Promise<void> => {
             webSocket.state = { ...webSocket.state, quit: true };
 
             const sessionId = webSocket.sessionId;
@@ -374,8 +421,13 @@ export class IORoom<
                 senderId: sessionId,
             });
 
+            const session = webSocket.session;
+            const user = session.user;
+
+            if (!!user) this._removeUserSession(user.id, sessionId);
+
             try {
-                const [session, doc] = await Promise.all([webSocket.getSession(), this._doc]);
+                const doc = await this._doc;
                 const encodedState = doc.getEncodedState();
 
                 await Promise.resolve(
@@ -383,7 +435,7 @@ export class IORoom<
                         context: this._context,
                         encodedState,
                         room: this.id,
-                        user: session.user,
+                        user,
                     }),
                 );
             } catch (error) {
@@ -419,12 +471,23 @@ export class IORoom<
         this._closeWebSockets(quitters);
     }
 
-    private async _emitRegistered(pluvWs: AbstractWebSocket): Promise<void> {
-        const session = await pluvWs.getSession();
+    private async _emitRegistered(pluvWs: AbstractWebSocket<any, TAuthorize>): Promise<void> {
+        const session = pluvWs.session;
         const sessionId = session.id;
+        const presence = session.presence;
         const user = session.user;
 
-        await this._sendSelfMessage({ type: "$registered", data: { sessionId } }, { sessionId, user });
+        await this._sendSelfMessage(
+            {
+                type: "$registered",
+                data: {
+                    presence,
+                    sessionId,
+                    timers: { presence: session.timers.presence },
+                },
+            },
+            { sessionId, user },
+        );
 
         try {
             const doc = await this._doc;
@@ -450,7 +513,7 @@ export class IORoom<
         throw new Error("Platform must use detached mode");
     }
 
-    private _getAbstractWs(webSocket: WebSocketType<TPlatform>): AbstractWebSocket | null {
+    private _getAbstractWs(webSocket: WebSocketType<TPlatform>): AbstractWebSocket<any, TAuthorize> | null {
         if ((webSocket as unknown as any) instanceof AbstractWebSocket) return webSocket;
 
         const sessionId = this._platform.getSessionId(webSocket);
@@ -469,7 +532,7 @@ export class IORoom<
     private async _getAuthorizedUser(
         token: Maybe<string>,
         options: WebSocketRegisterConfig<TPlatform>,
-    ): Promise<InferIOAuthorizeUser<InferIOAuthorize<this>>> {
+    ): Promise<InferIOAuthorizeUser<TAuthorize>> {
         const ioAuthorize = this._getIOAuthorize(options);
 
         if (!ioAuthorize) return null as InferIOAuthorizeUser<InferIOAuthorize<this>>;
@@ -521,6 +584,36 @@ export class IORoom<
         return this._authorize as ResolvedPluvIOAuthorize<any, any> | null;
     }
 
+    private _getLatestPresence(userId: string): { timer: number | null; presence: JsonObject | null } {
+        if (!this._authorize) return { timer: null, presence: null };
+
+        const sessionIds = Array.from(this._userSessionss.get(userId)?.values() ?? []);
+
+        if (!sessionIds.length) return { timer: null, presence: null };
+
+        return sessionIds.reduce(
+            (state, sessionId) => {
+                const pluvWs = this._sessions.get(sessionId) ?? null;
+
+                if (!pluvWs) return state;
+
+                const session = pluvWs.session;
+                const presence = session.presence;
+                const timer = session.timers.presence;
+
+                if (session.user?.id !== userId) return state;
+                if (typeof state.timer !== "number") return { presence, timer };
+                if (typeof timer !== "number") return state;
+
+                return timer > state.timer ? { presence, timer } : state;
+            },
+            { presence: null, timer: null } as {
+                presence: JsonObject | null;
+                timer: number | null;
+            },
+        );
+    }
+
     private _getProcedure(
         message: EventMessage<string, any>,
     ): (typeof this._router)["_defs"]["events"][keyof (typeof this._router)["_defs"]["events"]] | null {
@@ -535,18 +628,18 @@ export class IORoom<
         return procedure.config.input ? procedure.config.input.parse(message.data) : message.data;
     }
 
-    private async _getSession(webSocket: WebSocketType<TPlatform>): Promise<WebSocketSession<TAuthorize>> {
+    private _getSession(webSocket: WebSocketType<TPlatform>): WebSocketSession<TAuthorize> {
         const pluvWs = this._getAbstractWs(webSocket);
 
         if (!pluvWs) throw new Error("Session could not be found");
 
-        return await pluvWs.getSession<TAuthorize>();
+        return pluvWs.session;
     }
 
-    private async _getSessions() {
-        const promises = Array.from(this._sessions.values()).map(async (pluvWs) => await pluvWs.getSession());
+    private _getSessions() {
+        const sessions = Array.from(this._sessions.values()).map((pluvWs) => pluvWs.session);
 
-        return await Promise.all(promises);
+        return sessions;
     }
 
     private _initialize() {
@@ -624,14 +717,14 @@ export class IORoom<
         if (this._debug) console.log(...data);
     }
 
-    private _onClose(webSocket: AbstractWebSocket): () => Promise<void> {
+    private _onClose(webSocket: AbstractWebSocket<any, TAuthorize>): () => Promise<void> {
         return async (): Promise<void> => {
             if (!(await this._initialized)) return;
             await this._closeWebSockets([webSocket]);
         };
     }
 
-    private _onMessage(webSocket: AbstractWebSocket): (event: AbstractMessageEvent) => Promise<void> {
+    private _onMessage(webSocket: AbstractWebSocket<any, TAuthorize>): (event: AbstractMessageEvent) => Promise<void> {
         return async (event: AbstractMessageEvent): Promise<void> => {
             if (!(await this._initialized)) return;
 
@@ -639,18 +732,34 @@ export class IORoom<
 
             if (!pluvWs) throw new Error("Could not get session");
 
-            const [session, sessions] = await Promise.all([
-                this._getSession(pluvWs as WebSocketType<TPlatform>),
-                this._getSessions(),
-            ]);
+            const session = this._getSession(pluvWs as WebSocketType<TPlatform>);
+            const sessions = this._getSessions();
+
+            const setPresence = (params: PatchPresenceParams): void => {
+                this._setPresence.bind(this)(params);
+            };
 
             const doc = await this._doc;
             const eventContext: EventResolverContext<EventResolverKind, TPlatform, TAuthorize, TContext> = {
                 context: this._context,
                 doc,
+                garbageCollect: async () => {
+                    await this.garbageCollect();
+                },
+                get presence() {
+                    return session.presence as JsonObject | null;
+                },
+                set presence(presence: JsonObject | null) {
+                    const sessionId = this.session?.id;
+
+                    if (!sessionId) return;
+
+                    setPresence({ presence, sessionId, timer: this.time });
+                },
                 room: this.id,
                 session,
                 sessions,
+                time: new Date().getTime(),
             };
 
             if (pluvWs.state.quit) {
@@ -711,24 +820,24 @@ export class IORoom<
                 const handleBroadcast = async () => {
                     if (!broadcast) return;
 
-                    await Promise.all(
-                        Object.entries(broadcast).map(async ([type, data]: any) => {
-                            await this._broadcast({
-                                message: { data, type },
-                                senderId: sessionId,
-                            });
-                        }),
-                    );
+                    const messages = Object.entries(broadcast).map(async ([type, data]: any) => {
+                        await this._broadcast({
+                            message: { data, type },
+                            senderId: sessionId,
+                        });
+                    });
+
+                    await Promise.all(messages);
                 };
 
                 const handleSelf = async () => {
                     if (!self) return;
 
-                    await Promise.all(
-                        Object.entries(self).map(async ([type, data]) => {
-                            await this._sendSelfMessage({ data, type }, { sessionId, user });
-                        }),
-                    );
+                    const messages = Object.entries(self).map(async ([type, data]) => {
+                        await this._sendSelfMessage({ data, type }, { sessionId, user });
+                    });
+
+                    await Promise.all(messages);
                 };
 
                 const handleSync = async () => {
@@ -742,11 +851,11 @@ export class IORoom<
                         ...message,
                     });
 
-                    await Promise.all(
-                        Object.entries(sync).map(async ([type, data]) => {
-                            await this._sendSelfMessage({ data, type }, { sessionId, user });
-                        }),
-                    );
+                    const messages = Object.entries(sync).map(async ([type, data]) => {
+                        await this._sendSelfMessage({ data, type }, { sessionId, user });
+                    });
+
+                    await Promise.all(messages);
                 };
 
                 await Promise.all([handleBroadcast(), handleSelf(), handleSync()]);
@@ -768,7 +877,58 @@ export class IORoom<
         }
     }
 
-    private async _sendMessage(pluvWs: AbstractWebSocket<any>, message: IOEventMessage<any>): Promise<void> {
+    private _setPresence(params: PatchPresenceParams): void {
+        const { presence, sessionId, timer: _timer } = params;
+
+        const timer = _timer ?? new Date().getTime();
+        const pluvWs = this._sessions.get(sessionId) ?? null;
+
+        if (!pluvWs) return;
+
+        const wsSession = pluvWs.session;
+        const user = wsSession.user as BaseUser | null;
+
+        const sessionIds = user
+            ? new Set<string>([...(this._userSessionss.get(user.id) ?? []), sessionId])
+            : new Set([sessionId]);
+
+        sessionIds.forEach((sId) => {
+            const pWs = this._sessions.get(sId);
+            const session = pWs?.session;
+
+            if (!session) return;
+
+            const prevState = session.webSocket.state;
+
+            this._platform.setSerializedState(session.webSocket, {
+                ...prevState,
+                presence,
+                timers: {
+                    ...prevState.timers,
+                    presence: timer,
+                },
+            });
+        });
+    }
+
+    private _removeUserSession(userId: string, sessionId: string): Set<[sessionId: string][0]> | null {
+        const set = this._userSessionss.get(userId);
+
+        if (!set) return null;
+
+        set.delete(sessionId);
+
+        if (!!set.size) return set;
+
+        this._userSessionss.delete(userId);
+
+        return set;
+    }
+
+    private async _sendMessage(
+        pluvWs: AbstractWebSocket<any, TAuthorize>,
+        message: IOEventMessage<any>,
+    ): Promise<void> {
         if (!(await this._initialized)) return;
 
         await Promise.resolve(pluvWs.sendMessage(message));
@@ -790,10 +950,10 @@ export class IORoom<
                 const pluvWs = this._sessions.get(id);
 
                 return pluvWs ? dict.set(id, pluvWs) : dict;
-            }, new Map<string, AbstractWebSocket>()) ?? this._sessions;
+            }, new Map<string, AbstractWebSocket<any, TAuthorize>>()) ?? this._sessions;
 
         const promises = Array.from(webSockets.values()).map(async (pluvWs) => {
-            const session = await pluvWs.getSession();
+            const session = pluvWs.session;
             const user = session.user ?? null;
 
             this._sendMessage(pluvWs, { connectionId, data, room, type, user });
@@ -813,9 +973,22 @@ export class IORoom<
         const pluvWs = this._sessions.get(senderId);
 
         if (!pluvWs) return;
-        if (message.type === "$pong") await this._emitQuitters();
 
-        const session = await pluvWs.getSession();
+        /**
+         * @description Note that this will not fire for Cloudflare's hibernatable websockets
+         * because the server will auto-respond with the $pong event without calling any
+         * functions.
+         * @date April 18, 2025
+         */
+        if (message.type === "$pong") {
+            const elapsedMs = Date.now() - this._lastGarbageCollectMs;
+
+            if (elapsedMs > GARBAGE_COLLECT_INTERVAL_MS) {
+                await this.garbageCollect();
+            }
+        }
+
+        const session = pluvWs.session;
         const user = session.user ?? null;
 
         this._sendMessage(pluvWs, {
@@ -840,9 +1013,19 @@ export class IORoom<
         const context: EventResolverContext<"sync", TPlatform, TAuthorize, TContext> = {
             context: this._context,
             doc,
+            garbageCollect: async () => {
+                await this.garbageCollect();
+            },
+            get presence() {
+                return null;
+            },
+            set presence(presence: JsonObject | null) {
+                return;
+            },
             room: this.id,
             session: null,
-            sessions: await this._getSessions(),
+            sessions: this._getSessions(),
+            time: new Date().getTime(),
         };
 
         const resolver = this._getProcedure(message)?.config.sync;