From 8d2de5727a81cd0d7715100200a5d2b982cf7d57 Mon Sep 17 00:00:00 2001
From: Sophia Xu <sophia.xu@mail.mcgill.ca>
Date: Mon, 28 Feb 2022 19:34:13 -0500
Subject: [PATCH] perf: tidy up subscription on client side

---
 packages/unigraph-dev-backend/src/server.ts        |  5 +++--
 packages/unigraph-dev-backend/src/subscriptions.ts | 12 ++++--------
 packages/unigraph-dev-common/src/api/unigraph.ts   |  7 +++++++
 packages/unigraph-dev-common/src/utils/utils.ts    | 10 ++++++++--
 4 files changed, 22 insertions(+), 12 deletions(-)

diff --git a/packages/unigraph-dev-backend/src/server.ts b/packages/unigraph-dev-backend/src/server.ts
index 826917c3..a6244ef4 100644
--- a/packages/unigraph-dev-backend/src/server.ts
+++ b/packages/unigraph-dev-backend/src/server.ts
@@ -106,7 +106,7 @@ export default async function startServer(client: DgraphClient) {
     await checkOrCreateDefaultDataModel(client);
 
     // Initialize subscriptions
-    const pollCallback: MsgCallbackFn = (newdata, sub, ofUpdate) => {
+    const pollCallback: MsgCallbackFn = (newdata, sub, ofUpdate, supplementary?) => {
         if (sub?.callbackType === 'messageid') {
             const msgPort = sub.msgPort!;
             if (msgPort?.readyState === 1) {
@@ -118,8 +118,9 @@ export default async function startServer(client: DgraphClient) {
                             id: sub.id,
                             result: newdata,
                             ofUpdate,
+                            supplementary,
                         },
-                        { replacer: getCircularReplacer() },
+                        supplementary ? {} : { replacer: getCircularReplacer() },
                     ),
                 );
             }
diff --git a/packages/unigraph-dev-backend/src/subscriptions.ts b/packages/unigraph-dev-backend/src/subscriptions.ts
index 410bdb40..3e9df418 100644
--- a/packages/unigraph-dev-backend/src/subscriptions.ts
+++ b/packages/unigraph-dev-backend/src/subscriptions.ts
@@ -90,7 +90,7 @@ export function buildPollingQuery(subs: { id: any; query: any }[], states: any)
     }, '{')} }`;
 }
 
-export type MsgCallbackFn = (updated: any, sub: Subscription, ofUpdate?: number | string) => any;
+export type MsgCallbackFn = (updated: any, sub: Subscription, ofUpdate?: number | string, supplementary?: any) => any;
 
 export type MergedSubscription = {
     subscriptions: Subscription[];
@@ -104,13 +104,13 @@ export function mergeSubscriptions(
     ids?: any[],
     states?: any,
 ): MergedSubscription[] {
-    function callbackIfChanged(updated: any, sub: Subscription, ofUpdate: any) {
+    function callbackIfChanged(updated: any, sub: Subscription, ofUpdate: any, supplementary?: any) {
         if (
             stringify(updated, { replacer: getCircularReplacer() }) !==
             stringify(sub.data, { replacer: getCircularReplacer() })
         ) {
             sub.data = updated;
-            msgCallback(updated, sub, ofUpdate);
+            msgCallback(updated, sub, ofUpdate, supplementary);
         }
     }
 
@@ -204,10 +204,6 @@ export function mergeSubscriptions(
                 subscriptions: subs,
                 aggregateQuery: query,
                 resolver: (updated: any, ofUpdate: any) => {
-                    const startTime = new Date().getTime();
-                    buildGraph(updated, true);
-                    const graphTime = new Date().getTime() - startTime;
-                    if (graphTime > 5) console.log(`Build graph took ${graphTime}ms, which is a bit slow`);
                     subs.forEach((el) => {
                         const uidResolver = (uu: string) => (uu.startsWith('$/') ? states.namespaceMap[uu].uid : uu);
                         const allUids = (el.query as QueryObject).uid;
@@ -220,7 +216,7 @@ export function mergeSubscriptions(
                             (Array.isArray((el.query as QueryObject).uid) ? (el.query as QueryObject).uid.length : 1)
                         )
                             return;
-                        callbackIfChanged(updatedIts, el, ofUpdate);
+                        callbackIfChanged(updatedIts, el, ofUpdate, updated);
                     });
                 },
             });
diff --git a/packages/unigraph-dev-common/src/api/unigraph.ts b/packages/unigraph-dev-common/src/api/unigraph.ts
index 37355469..eeed3e9c 100644
--- a/packages/unigraph-dev-common/src/api/unigraph.ts
+++ b/packages/unigraph-dev-common/src/api/unigraph.ts
@@ -209,6 +209,13 @@ export default function unigraph(url: string, browserId: string): Unigraph<WebSo
                 }
                 subFakeUpdates[parsed.id] = [];
                 subResults[parsed.id] = JSON.parse(JSON.stringify(parsed.result));
+                if (parsed.supplementary) {
+                    buildGraph(
+                        _.uniqBy([...parsed.result, ...parsed.supplementary], 'uid'),
+                        true,
+                        parsed.result.length,
+                    );
+                }
                 subscriptions[parsed.id](parsed.result);
             }
             if (parsed.type === 'open_url' && window) window.open(parsed.url, '_blank');
diff --git a/packages/unigraph-dev-common/src/utils/utils.ts b/packages/unigraph-dev-common/src/utils/utils.ts
index 593e1d72..5a89ee53 100644
--- a/packages/unigraph-dev-common/src/utils/utils.ts
+++ b/packages/unigraph-dev-common/src/utils/utils.ts
@@ -296,7 +296,11 @@ export class UnigraphObject extends Object {
  *
  * @param objects Objects with uid references
  */
-export function buildGraph(objects: UnigraphObject[], topLevelOnly?: boolean): UnigraphObject[] {
+export function buildGraph(
+    objects: UnigraphObject[],
+    topLevelOnly?: boolean,
+    onlyLookAtFirstN?: number,
+): UnigraphObject[] {
     const objs: any[] = [...objects].map((el: any) => (Array.isArray(el) ? el : new UnigraphObject(el)));
     const dict: any = {};
 
@@ -352,7 +356,9 @@ export function buildGraph(objects: UnigraphObject[], topLevelOnly?: boolean): U
 
     // console.log(dict);
 
-    objs.forEach((object) => buildGraphRecurse(object));
+    objs.filter((el, index) => (onlyLookAtFirstN ? index < onlyLookAtFirstN : true)).forEach((object) =>
+        buildGraphRecurse(object),
+    );
 
     const graphTime = new Date().getTime() - startTime;
     if (graphTime > 15) console.log(`Build graph took ${graphTime}ms, which is a bit slow`);