Add timings to the setting up of the web socket (#59)

Author: bosschaert

src/edge.js

@@ -94,6 +94,9 @@ async function handleApiCall(url, request, env) {
 // This is where the requests for the worker come in. They can either be pure API requests or
 // requests to set up a session with a Durable Object through a Yjs WebSocket.
 export async function handleApiRequest(request, env) {
+  let timingDaAdminHeadDuration;
+  const timingStartTime = Date.now();
+
   // We've received a pure API request - handle it and return.
   const url = new URL(request.url);
   if (url.pathname.startsWith('/api/')) {
@@ -127,7 +130,9 @@ export async function handleApiRequest(request, env) {
       opts.headers = new Headers({ Authorization: auth });
     }
 
+    const timingBeforeDaAdminHead = Date.now();
     const initialReq = await env.daadmin.fetch(docName, opts);
+    timingDaAdminHeadDuration = Date.now() - timingBeforeDaAdminHead;
 
     if (!initialReq.ok && initialReq.status !== 404) {
       // eslint-disable-next-line no-console
@@ -140,6 +145,7 @@ export async function handleApiRequest(request, env) {
     return new Response('unable to get resource', { status: 500 });
   }
 
+  const timingBeforeDocRoomGet = Date.now();
   // Each Durable Object has a 256-bit unique ID. Route the request based on the path.
   const id = env.rooms.idFromName(docName);
 
@@ -150,11 +156,17 @@ export async function handleApiRequest(request, env) {
   // created on-demand when the ID is first used, there's nothing to wait for anyway; we know
   // an object will be available somewhere to receive our requests.
   const roomObject = env.rooms.get(id);
+  const timingDocRoomGetDuration = Date.now() - timingBeforeDocRoomGet;
 
   // eslint-disable-next-line no-console
   console.log(`FETCHING: ${docName} ${id}`);
 
-  const headers = [...request.headers, ['X-collab-room', docName]];
+  const headers = [...request.headers,
+    ['X-collab-room', docName],
+    ['X-timing-start', timingStartTime],
+    ['X-timing-da-admin-head-duration', timingDaAdminHeadDuration],
+    ['X-timing-docroom-get-duration', timingDocRoomGetDuration],
+  ];
   if (auth) {
     headers.push(['Authorization', auth]);
   }
@@ -248,6 +260,7 @@ export class DocRoom {
       return new Response('expected docName', { status: 400 });
     }
 
+    const timingBeforeSetupWebsocket = Date.now();
     // To accept the WebSocket request, we create a WebSocketPair (which is like a socketpair,
     // i.e. two WebSockets that talk to each other), we return one end of the pair in the
     // response, and we operate on the other end. Note that this API is not part of the
@@ -256,10 +269,20 @@ export class DocRoom {
     const pair = DocRoom.newWebSocketPair();
 
     // We're going to take pair[1] as our end, and return pair[0] to the client.
-    await this.handleSession(pair[1], docName, auth);
+    const timingData = await this.handleSession(pair[1], docName, auth);
+    const timingSetupWebSocketDuration = Date.now() - timingBeforeSetupWebsocket;
+
+    const reqHeaders = request.headers;
+    const respheaders = new Headers();
+    respheaders.set('X-1-timing-da-admin-head-duration', reqHeaders.get('X-timing-da-admin-head-duration'));
+    respheaders.set('X-2-timing-docroom-get-duration', reqHeaders.get('X-timing-docroom-get-duration'));
+    respheaders.set('X-4-timing-da-admin-get-duration', timingData.get('timingDaAdminGetDuration'));
+    respheaders.set('X-5-timing-read-state-duration', timingData.get('timingReadStateDuration'));
+    respheaders.set('X-7-timing-setup-websocket-duration', timingSetupWebSocketDuration);
+    respheaders.set('X-9-timing-full-duration', Date.now() - reqHeaders.get('X-timing-start'));
 
     // Now we return the other end of the pair to the client.
-    return new Response(null, { status: successCode, webSocket: pair[0] });
+    return new Response(null, { status: successCode, headers: respheaders, webSocket: pair[0] });
   }
 
   /**
@@ -277,6 +300,7 @@ export class DocRoom {
     // eslint-disable-next-line no-console
     console.log(`setupWSConnection ${docName} with auth(${webSocket.auth
       ? webSocket.auth.substring(0, webSocket.auth.indexOf(' ')) : 'none'})`);
-    await setupWSConnection(webSocket, docName, this.env, this.storage);
+    const timingData = await setupWSConnection(webSocket, docName, this.env, this.storage);
+    return timingData;
   }
 }

src/shareddoc.js

@@ -265,11 +265,16 @@ export const persistence = {
    * @param {TransactionalStorage} storage - the worker transactional storage object
    */
   bindState: async (docName, ydoc, conn, storage) => {
+    let timingReadStateDuration;
+    let timingDaAdminGetDuration;
+
     let current;
     let restored = false; // True if restored from worker storage
     try {
       let newDoc = false;
+      const timingBeforeDaAdminGet = Date.now();
       current = await persistence.get(docName, conn.auth, ydoc.daadmin);
+      timingDaAdminGetDuration = Date.now() - timingBeforeDaAdminGet;
       if (current === null) {
         // The document isn't there any more, clear the local storage
         await storage.deleteAll();
@@ -278,8 +283,11 @@ export const persistence = {
         newDoc = true;
       }
 
+      const timingBeforeReadState = Date.now();
       // Read the stored state from internal worker storage
       const stored = await readState(docName, storage);
+      timingReadStateDuration = Date.now() - timingBeforeReadState;
+
       if (stored && stored.length > 0) {
         Y.applyUpdate(ydoc, stored);
 
@@ -344,6 +352,11 @@ export const persistence = {
         current = await persistence.update(ydoc, current);
       }
     }, 2000, { maxWait: 10000 }));
+
+    const timingMap = new Map();
+    timingMap.set('timingReadStateDuration', timingReadStateDuration);
+    timingMap.set('timingDaAdminGetDuration', timingDaAdminGetDuration);
+    return timingMap;
   },
 };
 
@@ -403,7 +416,7 @@ export class WSSharedDoc extends Y.Doc {
  * @param {boolean} gc - whether garbage collection is enabled
  * @returns The Yjs document object, which may be shared across multiple sockets.
  */
-export const getYDoc = async (docname, conn, env, storage, gc = true) => {
+export const getYDoc = async (docname, conn, env, storage, timingData, gc = true) => {
   let doc = docs.get(docname);
   if (doc === undefined) {
     // The doc is not yet in the cache, create a new one.
@@ -426,7 +439,10 @@ export const getYDoc = async (docname, conn, env, storage, gc = true) => {
 
   // We wait for the promise, for second and subsequent connections to the same doc, this will
   // already be resolved.
-  await doc.promise;
+  const timings = await doc.promise;
+  if (timingData) {
+    timings.forEach((v, k) => timingData.set(k, v));
+  }
   return doc;
 };
 
@@ -498,10 +514,12 @@ export const invalidateFromAdmin = async (docName) => {
  * @returns {Promise<void>} - The return value of this
  */
 export const setupWSConnection = async (conn, docName, env, storage) => {
+  const timingData = new Map();
+
   // eslint-disable-next-line no-param-reassign
   conn.binaryType = 'arraybuffer';
   // get doc, initialize if it does not exist yet
-  const doc = await getYDoc(docName, conn, env, storage, true);
+  const doc = await getYDoc(docName, conn, env, storage, timingData, true);
 
   // listen and reply to events
   conn.addEventListener('message', (message) => messageListener(conn, doc, new Uint8Array(message.data)));
@@ -527,4 +545,6 @@ export const setupWSConnection = async (conn, docName, env, storage) => {
       send(doc, conn, encoding.toUint8Array(encoder));
     }
   }
+
+  return timingData;
 };

test/edge.test.js

@@ -220,7 +220,10 @@ describe('Worker test suite', () => {
 
     try {
       const bindCalled = [];
-      persistence.bindState = async (nm, d, c) => bindCalled.push({nm, d, c});
+      persistence.bindState = async (nm, d, c) => {
+        bindCalled.push({nm, d, c});
+        return new Map();
+      }
 
       const wspCalled = [];
       const wsp0 = {};

test/shareddoc.test.js

@@ -781,7 +781,10 @@ describe('Collab Test Suite', () => {
 
     try {
       const bindCalls = [];
-      persistence.bindState = async (nm, d, c, s) => bindCalls.push({nm, d, c, s});
+      persistence.bindState = async (nm, d, c, s) => {
+        bindCalls.push({nm, d, c, s});
+        return new Map();
+      }
 
       const docName = 'https://somewhere.com/somedoc.html';
       const eventListeners = new Map();
@@ -826,7 +829,7 @@ describe('Collab Test Suite', () => {
     const savedBind = persistence.bindState;
 
     try {
-      persistence.bindState = async (nm, d, c, s) => {};
+      persistence.bindState = async (nm, d, c, s) => new Map();
 
       const docName = 'https://somewhere.com/myotherdoc.html';
       const closeCalls = [];