batch wikilinks redis requests on article parsing together

openzim · Jan 7, 2023 · af568c6 · af568c6
1 parent 6fb1e71
commit af568c6
Show file tree

Hide file tree

Showing 4 changed files with 325 additions and 196 deletions.
diff --git a/src/util/RedisKvs.ts b/src/util/RedisKvs.ts
@@ -1,5 +1,4 @@
 import {cpus} from 'os';
-import pmap from 'p-map';
 import type {RedisClient} from 'redis';
 
 
@@ -59,22 +58,41 @@ export class RedisKvs<T> {
     });
   }
 
-  public exists(prop: string[]): Promise<{ [key: string]: number }> {
-    return pmap(
-      prop,
-      (key: string) => {
-        return new Promise((resolve, reject) => {
-          this.redisClient.hexists(this.dbName, key, (err, val) => {
+  public exists(
+    prop: string | string[],
+  ): Promise<{ [key: string]: number } | number> {
+    return new Promise((resolve, reject) => {
+      if (typeof prop === 'string') {
+        // single key
+        this.redisClient.hexists(this.dbName, prop, (err, val) => {
+          if (err) {
+            reject(err);
+          } else {
+            resolve(val);
+          }
+        });
+      } else {
+        // array of keys
+        const batch = this.redisClient.batch();
+        prop.forEach((index) => {
+          batch.hexists(this.dbName, index);
+        });
+        batch.exec((err, replies) => {
+          try {
             if (err) {
-              reject(err);
-            } else {
-              resolve({[key]: val});
+              throw err;
             }
-          });
+            const result: { [key: string]: number } = {};
+            for (let u = 0; u < prop.length; u += 1) {
+              result[prop[u]] = replies[u];
+            }
+            resolve(result);
+          } catch(err) {
+            reject(err);
+          }
         });
-      },
-      {concurrency: cpus().length * 4}
-    ).then((vals: any[]) => vals.reduce((acc, val) => Object.assign(acc, val), {}));
+      }
+    });
   }
 
   public set(prop: string, val: T) {

diff --git a/src/util/categories.ts b/src/util/categories.ts
@@ -115,9 +115,9 @@ export async function trimUnmirroredPages(downloader: Downloader) {
                         subCategoriesExist,
                         pagesExist,
                     ] = await Promise.all([
-                        categoryIds.length ? articleDetailXId.exists(categoryIds) : Promise.resolve({}),
-                        subCategoryIds.length ? articleDetailXId.exists(subCategoryIds) : Promise.resolve({}),
-                        pageIds.length ? articleDetailXId.exists(pageIds) : Promise.resolve({}),
+                        categoryIds.length ? articleDetailXId.exists(categoryIds) as Promise<{ [key: string]: number; }> : Promise.resolve({}),
+                        subCategoryIds.length ? articleDetailXId.exists(subCategoryIds) as Promise<{ [key: string]: number; }> : Promise.resolve({}),
+                        pageIds.length ? articleDetailXId.exists(pageIds) as Promise<{ [key: string]: number; }> : Promise.resolve({}),
                     ]);
 
                     const existingCategories = Object.keys(categoriesExist).filter((key) => !!categoriesExist[key]);