rwsdk
Version:
Build fast, server-driven webapps on Cloudflare with SSR, RSC, and realtime
278 lines (277 loc) • 12.1 kB
JavaScript
import { loadCapnweb } from "./capnweb-loader.mjs";
import { DEFAULT_SYNCED_STATE_PATH } from "./constants.mjs";
// Converts a relative endpoint like "/__synced-state" to an absolute
// ws:// or wss:// URL so the same key is used by getSyncedStateClient,
// onStatusChange, and reconnect notifications.
function normalizeEndpoint(endpoint) {
if (endpoint.startsWith("/") && typeof window !== "undefined") {
const protocol = window.location.protocol === "https:" ? "wss:" : "ws:";
// Strip any trailing DNS root dot from the hostname so the WebSocket
// URL matches the normal origin Cloudflare routes on.
const host = window.location.host.replace(/\.$/, "");
return `${protocol}//${host}${endpoint}`;
}
return endpoint;
}
// Map of endpoint URLs to their respective clients
const clientCache = new Map();
// Tracks the promise of the underlying capnweb session per endpoint, exposed
// for tests so they can `await` the lazy load before making assertions.
const baseClientPromiseByEndpoint = new Map();
const activeSubscriptions = new Set();
// Status change listeners per endpoint. Uses an array rather than a Set so
// that two components passing the same callback reference (e.g. via
// createSyncedStateHook({ onStatusChange })) are tracked as two separate
// registrations — unsubscribing one must not cancel the other.
const statusListeners = new Map();
function notifyStatusChange(endpoint, status) {
const listeners = statusListeners.get(endpoint);
if (listeners) {
// Snapshot so unsubscribes fired by callbacks don't skip entries.
for (const cb of [...listeners]) {
cb(status);
}
}
}
/**
* Registers a callback that fires when the connection status changes for an endpoint.
* Returns an unsubscribe function.
*/
export const onStatusChange = (endpoint, callback) => {
// Normalize the endpoint so listeners registered with a relative URL
// (e.g. "/__synced-state") are notified using the same key as the
// cached client and the reconnect path.
const normalized = normalizeEndpoint(endpoint);
let listeners = statusListeners.get(normalized);
if (!listeners) {
listeners = [];
statusListeners.set(normalized, listeners);
}
listeners.push(callback);
return () => {
const idx = listeners.indexOf(callback);
if (idx !== -1) {
listeners.splice(idx, 1);
}
if (listeners.length === 0) {
statusListeners.delete(normalized);
}
};
};
// Tracks per-endpoint reconnection backoff state
const backoffState = new Map();
const BASE_BACKOFF_MS = 1000;
const MAX_BACKOFF_MS = 30000;
function getBackoffMs(attempt) {
const base = Math.min(BASE_BACKOFF_MS * 2 ** attempt, MAX_BACKOFF_MS);
// Add ±25% jitter to avoid thundering herd on server restart
const jittered = base * (0.75 + Math.random() * 0.5);
return Math.round(Math.min(jittered, MAX_BACKOFF_MS));
}
// Set up beforeunload handler to unsubscribe all active subscriptions
if (typeof window !== "undefined") {
const handleBeforeUnload = () => {
if (activeSubscriptions.size === 0) {
return;
}
// Unsubscribe all active subscriptions
// Use a synchronous approach where possible, but don't block page unload
const subscriptions = Array.from(activeSubscriptions);
activeSubscriptions.clear();
// Fire-and-forget unsubscribe calls - we can't await during beforeunload
for (const { key, handler, client } of subscriptions) {
void client.unsubscribe(key, handler).catch(() => {
// Ignore errors during page unload - the connection will be closed anyway
});
}
};
window.addEventListener("beforeunload", handleBeforeUnload);
}
function reconnect(endpoint, deadClient) {
// Don't schedule multiple reconnects for the same endpoint
const state = backoffState.get(endpoint) ?? { attempt: 0, timer: null };
if (state.timer !== null) {
return;
}
notifyStatusChange(endpoint, "disconnected");
const delayMs = getBackoffMs(state.attempt);
state.timer = setTimeout(() => {
state.timer = null;
state.attempt++;
backoffState.set(endpoint, state);
notifyStatusChange(endpoint, "reconnecting");
// Evict the dead client so getSyncedStateClient creates a fresh one
clientCache.delete(endpoint);
const newClient = getSyncedStateClient(endpoint);
// Re-subscribe everything that was on the dead client. Kick off both
// subscribe() and getState() synchronously so callers see the calls
// happen inside the timer tick, but only confirm "connected" once the
// subscribe promises resolve — otherwise a rejected resubscription
// would be masked as a successful reconnect.
const subscribePromises = [];
for (const sub of activeSubscriptions) {
if (sub.client === deadClient) {
sub.client = newClient;
subscribePromises.push(newClient.subscribe(sub.key, sub.handler));
void newClient.getState(sub.key).then((val) => {
if (val !== undefined) {
sub.handler(val);
}
});
}
}
Promise.all(subscribePromises).then(() => {
backoffState.set(endpoint, { attempt: 0, timer: null });
notifyStatusChange(endpoint, "connected");
}, () => {
// Resubscription failed — leave the attempt counter elevated so
// the next reconnect uses a longer backoff, and emit disconnected
// again. A subsequent onRpcBroken from the new (likely dead)
// client will drive the next retry.
notifyStatusChange(endpoint, "disconnected");
});
}, delayMs);
backoffState.set(endpoint, state);
}
/**
* Returns a cached client for the provided endpoint, creating it when necessary.
* The returned client is a proxy that loads `capnweb` lazily on first method
* call — consumers that never hit `use-synced-state` pay no import cost and
* don't need `capnweb` installed.
* @param endpoint Endpoint to connect to.
* @returns RPC client instance.
*/
export const getSyncedStateClient = (endpoint = DEFAULT_SYNCED_STATE_PATH) => {
// Convert relative endpoint to absolute URL for environments like WKWebView
endpoint = normalizeEndpoint(endpoint);
// Return existing client if already cached for this endpoint
const existingClient = clientCache.get(endpoint);
if (existingClient) {
return existingClient;
}
let baseClientPromise = null;
let wrappedClient;
const getBaseClient = () => {
if (!baseClientPromise) {
baseClientPromise = loadCapnweb().then((mod) => {
const session = mod.newWebSocketRpcSession(endpoint);
if (typeof session.onRpcBroken === "function") {
session.onRpcBroken(() => {
reconnect(endpoint, wrappedClient);
});
}
return session;
});
baseClientPromiseByEndpoint.set(endpoint, baseClientPromise);
}
return baseClientPromise;
};
wrappedClient = new Proxy({}, {
get(_target, prop) {
if (prop === "subscribe") {
return async (key, handler) => {
// Idempotent subscription registration. Reconnect mutates the
// existing entry's client in place and then re-subscribes on the
// new transport; without this guard we create duplicate registry
// entries that double on every reconnect and leak component
// closures after unmount.
const exists = [...activeSubscriptions].some((s) => s.key === key && s.handler === handler && s.client === wrappedClient);
if (!exists) {
activeSubscriptions.add({ key, handler, client: wrappedClient });
}
const base = await getBaseClient();
return base[prop](key, handler);
};
}
if (prop === "unsubscribe") {
return async (key, handler) => {
// Remove every matching registry entry. The reconnect path can
// create duplicates for the same (key, handler) pair, so a single
// component unmount must clean up all of them.
for (const sub of [...activeSubscriptions]) {
if (sub.key === key &&
sub.handler === handler &&
sub.client === wrappedClient) {
activeSubscriptions.delete(sub);
}
}
const base = await getBaseClient();
return base[prop](key, handler);
};
}
// Pass through all other properties/methods
return async (...args) => {
const base = await getBaseClient();
return base[prop](...args);
};
},
});
// Cache the client for this endpoint
clientCache.set(endpoint, wrappedClient);
// Eagerly kick off the capnweb load so the underlying session (and its
// onRpcBroken handler) is ready as soon as possible, and reconnect flows
// that don't call methods on the new client still create the replacement
// session. Errors are swallowed here to avoid unhandled rejections — they
// still surface through subsequent method calls because the rejected
// promise remains cached.
void getBaseClient().catch(() => { });
return wrappedClient;
};
/**
* Initializes and caches an RPC client instance for the sync state endpoint.
* The client is wrapped to track subscriptions for cleanup on page reload.
* @param options Optional endpoint override.
* @returns Cached client instance or `null` when running without `window`.
*/
export const initSyncedStateClient = (options = {}) => {
const endpoint = options.endpoint ?? DEFAULT_SYNCED_STATE_PATH;
if (typeof window === "undefined") {
return null;
}
// Use getSyncedStateClient which now handles caching via Map
return getSyncedStateClient(endpoint);
};
/**
* Injects a client instance for tests and updates the cached endpoint.
* Also clears the subscription registry for test isolation.
* @param client Stub client instance or `null` to clear the cache.
* @param endpoint Endpoint associated with the injected client.
*/
export const setSyncedStateClientForTesting = (client, endpoint = DEFAULT_SYNCED_STATE_PATH) => {
if (client) {
clientCache.set(endpoint, client);
}
else {
clientCache.delete(endpoint);
}
baseClientPromiseByEndpoint.delete(endpoint);
activeSubscriptions.clear();
statusListeners.clear();
// Clear any pending reconnection timers
for (const [, state] of backoffState) {
if (state.timer !== null) {
clearTimeout(state.timer);
}
}
backoffState.clear();
};
// Exported for testing only
export const __testing = {
activeSubscriptions,
clientCache,
backoffState,
statusListeners,
reconnect,
getBackoffMs,
// Awaits the eagerly-kicked-off capnweb load for a cached client. Tests
// should `await __testing.warmUp(endpoint)` after `getSyncedStateClient`
// (or after a reconnect) when they need the underlying session to exist
// before asserting on it.
async warmUp(endpoint = DEFAULT_SYNCED_STATE_PATH) {
const normalized = normalizeEndpoint(endpoint);
const promise = baseClientPromiseByEndpoint.get(normalized);
if (promise) {
await promise.catch(() => { });
}
},
};