feat(pubsub): Add shared context among Redis instances (#1110)

* feat(pubsub): add shared context among redis instances, fix #773 * fix: command transformers Co-authored-by: Sibelius Seraphini <sibeliusseraphini@gmail.com>
stipsan · Jan 22, 2022 · 3da06de · 3da06de
1 parent cb17f5a
commit 3da06de
Show file tree

Hide file tree

Showing 13 changed files with 363 additions and 128 deletions.
diff --git a/jest.config.js b/jest.config.js
@@ -3,4 +3,5 @@ module.exports = {
   testMatch: ['**/test/**/*.js'],
   coverageDirectory: 'coverage',
   setupFiles: ['<rootDir>/testSetupBabel.js'],
+  setupFilesAfterEnv: ['<rootDir>/testSetupAfterEnv.js'],
 };
diff --git a/src/commands/flushall.js b/src/commands/flushall.js
@@ -1,4 +1,20 @@
+import contextMap, { createContext } from '../context';
+import { createData } from '../data';
+import { createExpires } from '../expires';
+
 export function flushall() {
-  this.data.clear();
+  const oldContext = contextMap.get(this.keyData);
+  const newContext = createContext(oldContext.keyPrefix);
+
+  contextMap.set(this.keyData, newContext);
+
+  this.expires = createExpires(newContext.expires, newContext.keyPrefix);
+  this.data = createData(
+    newContext.data,
+    this.expires,
+    {},
+    newContext.keyPrefix
+  );
+
   return 'OK';
 }
diff --git a/src/commands/flushdb.js b/src/commands/flushdb.js
@@ -1,4 +1,4 @@
 export function flushdb() {
-  this.data.clear();
+  this.flushall();
   return 'OK';
 }
diff --git a/src/context.js b/src/context.js
@@ -0,0 +1,19 @@
+import { EventEmitter } from 'events';
+import { createSharedData } from './data';
+import { createSharedExpires } from './expires';
+
+const contextMap = new Map();
+
+export default contextMap;
+
+export function createContext(keyPrefix) {
+  const expires = createSharedExpires();
+
+  return {
+    channels: new EventEmitter(),
+    expires,
+    data: createSharedData(expires),
+    patternChannels: new EventEmitter(),
+    keyPrefix,
+  };
+}
diff --git a/src/data.js b/src/data.js
@@ -1,79 +1,95 @@
 import { assign } from 'lodash';
 
-export default function createData(
-  expiresInstance,
-  initial = {},
-  keyPrefix = ''
-) {
+export function createSharedData(sharedExpires) {
   let raw = {};
 
-  function createInstance(prefix, expires) {
-    return Object.freeze({
-      clear() {
-        raw = {};
-      },
-      delete(key) {
-        if (expires.has(key)) {
-          expires.delete(key);
-        }
-        delete raw[`${prefix}${key}`];
-      },
-      get(key) {
-        if (expires.has(key) && expires.isExpired(key)) {
-          this.delete(key);
-        }
+  return Object.freeze({
+    clear() {
+      raw = {};
+    },
+    delete(key) {
+      if (sharedExpires.has(key)) {
+        sharedExpires.delete(key);
+      }
+      delete raw[key];
+    },
+    get(key) {
+      if (sharedExpires.has(key) && sharedExpires.isExpired(key)) {
+        this.delete(key);
+      }
 
-        const value = raw[`${prefix}${key}`];
+      const value = raw[key];
 
-        if (Array.isArray(value)) {
-          return value.slice();
-        }
+      if (Array.isArray(value)) {
+        return value.slice();
+      }
 
-        if (Buffer.isBuffer(value)) {
-          return Buffer.from(value);
-        }
+      if (Buffer.isBuffer(value)) {
+        return Buffer.from(value);
+      }
 
-        if (value instanceof Set) {
-          return new Set(value);
-        }
+      if (value instanceof Set) {
+        return new Set(value);
+      }
 
-        if (value instanceof Map) {
-          return new Map(value);
-        }
+      if (value instanceof Map) {
+        return new Map(value);
+      }
 
-        if (typeof value === 'object' && value) {
-          return assign({}, value);
-        }
+      if (typeof value === 'object' && value) {
+        return assign({}, value);
+      }
 
-        return value;
-      },
-      has(key) {
-        if (expires.has(key) && expires.isExpired(key)) {
-          this.delete(key);
-        }
+      return value;
+    },
+    has(key) {
+      if (sharedExpires.has(key) && sharedExpires.isExpired(key)) {
+        this.delete(key);
+      }
 
-        return {}.hasOwnProperty.call(raw, `${prefix}${key}`);
-      },
-      keys() {
-        return Object.keys(raw);
-      },
-      set(key, val) {
-        let item = val;
-
-        if (Array.isArray(val)) {
-          item = val.slice();
-        } else if (Buffer.isBuffer(val)) {
-          item = Buffer.from(val);
-        } else if (val instanceof Set) {
-          item = new Set(val);
-        } else if (val instanceof Map) {
-          item = new Map(val);
-        } else if (typeof val === 'object' && val) {
-          item = assign({}, val);
-        }
-
-        raw[`${prefix}${key}`] = item;
-      },
+      return {}.hasOwnProperty.call(raw, key);
+    },
+    keys(prefix) {
+      const keys = Object.keys(raw);
+
+      if (!prefix) return keys;
+
+      return keys.filter((key) => key.startsWith(prefix));
+    },
+    set(key, val) {
+      let item = val;
+
+      if (Array.isArray(val)) {
+        item = val.slice();
+      } else if (Buffer.isBuffer(val)) {
+        item = Buffer.from(val);
+      } else if (val instanceof Set) {
+        item = new Set(val);
+      } else if (val instanceof Map) {
+        item = new Map(val);
+      } else if (typeof val === 'object' && val) {
+        item = assign({}, val);
+      }
+
+      raw[key] = item;
+    },
+  });
+}
+
+export function createData(
+  sharedData,
+  expiresInstance,
+  initial = {},
+  keyPrefix = ''
+) {
+  function createInstance(prefix, expires) {
+    return Object.freeze({
+      clear: () => sharedData.clear(),
+      delete: (key) => sharedData.delete(`${prefix}${key}`),
+      get: (key) => sharedData.get(`${prefix}${key}`),
+      has: (key) => sharedData.has(`${prefix}${key}`),
+      keys: () => sharedData.keys(prefix),
+      set: (key, val) => sharedData.set(`${prefix}${key}`, val),
       withKeyPrefix(newKeyPrefix) {
         if (newKeyPrefix === prefix) return this;
         return createInstance(

diff --git a/src/expires.js b/src/expires.js
@@ -1,27 +1,39 @@
-export default function createExpires(keyPrefix = '') {
+export function createSharedExpires() {
   const expires = {};
+
+  return Object.freeze({
+    get(key) {
+      return expires[key];
+    },
+    set(key, timestamp) {
+      expires[key] = +timestamp;
+    },
+    has(key) {
+      return {}.hasOwnProperty.call(expires, key);
+    },
+    isExpired(key) {
+      return expires[key] <= Date.now();
+    },
+    delete(key) {
+      delete expires[key];
+    },
+  });
+}
+
+export function createExpires(sharedExpires, keyPrefix = '') {
   function createInstance(prefix) {
     return {
-      get(key) {
-        return expires[`${prefix}${key}`];
-      },
-      set(key, timestamp) {
-        expires[`${prefix}${key}`] = +timestamp;
-      },
-      has(key) {
-        return {}.hasOwnProperty.call(expires, `${prefix}${key}`);
-      },
-      isExpired(key) {
-        return expires[`${prefix}${key}`] <= Date.now();
-      },
-      delete(key) {
-        delete expires[`${prefix}${key}`];
-      },
+      get: (key) => sharedExpires.get(`${prefix}${key}`),
+      set: (key, timestamp) => sharedExpires.set(`${prefix}${key}`, timestamp),
+      has: (key) => sharedExpires.has(`${prefix}${key}`),
+      isExpired: (key) => sharedExpires.isExpired(`${prefix}${key}`),
+      delete: (key) => sharedExpires.delete(`${prefix}${key}`),
       withKeyPrefix(newPrefix) {
         if (newPrefix === prefix) return this;
         return createInstance(newPrefix);
       },
     };
   }
+
   return createInstance(keyPrefix);
 }
diff --git a/src/index.js b/src/index.js
@@ -4,18 +4,21 @@ import redisCommands from 'redis-commands';
 import * as commands from './commands';
 import * as commandsStream from './commands-stream';
 import createCommand, { Command } from './command';
-import createData from './data';
-import createExpires from './expires';
 import emitConnectEvent from './commands-utils/emitConnectEvent';
 import Pipeline from './pipeline';
 import promiseContainer from './promise-container';
 import parseKeyspaceEvents from './keyspace-notifications';
+import contextMap, { createContext } from './context';
+import { createExpires } from './expires';
+import { createData } from './data';
 
 const defaultOptions = {
   data: {},
   keyPrefix: '',
   lazyConnect: false,
   notifyKeyspaceEvents: '', // string pattern as specified in https://redis.io/topics/notifications#configuration e.g. 'gxK'
+  host: 'localhost',
+  port: '6379',
 };
 
 class RedisMock extends EventEmitter {
@@ -29,8 +32,7 @@ class RedisMock extends EventEmitter {
 
   constructor(options = {}) {
     super();
-    this.channels = new EventEmitter();
-    this.patternChannels = new EventEmitter();
+
     this.batch = undefined;
     this.connected = false;
     this.subscriberMode = false;
@@ -41,9 +43,19 @@ class RedisMock extends EventEmitter {
     // eslint-disable-next-line prefer-object-spread
     const optionsWithDefault = Object.assign({}, defaultOptions, options);
 
-    this.expires = createExpires(optionsWithDefault.keyPrefix);
+    this.keyData = `${optionsWithDefault.host}:${optionsWithDefault.port}`;
+
+    if (!contextMap.get(this.keyData)) {
+      const context = createContext(optionsWithDefault.keyPrefix);
+
+      contextMap.set(this.keyData, context);
+    }
+
+    const context = contextMap.get(this.keyData);
 
+    this.expires = createExpires(context.expires, optionsWithDefault.keyPrefix);
     this.data = createData(
+      context.data,
       this.expires,
       optionsWithDefault.data,
       optionsWithDefault.keyPrefix
@@ -61,6 +73,36 @@ class RedisMock extends EventEmitter {
     }
   }
 
+  get channels() {
+    return contextMap.get(this.keyData).channels;
+  }
+
+  set channels(channels) {
+    const oldContext = contextMap.get(this.keyData);
+
+    const newContext = {
+      ...oldContext,
+      channels,
+    };
+
+    contextMap.set(this.keyData, newContext);
+  }
+
+  get patternChannels() {
+    return contextMap.get(this.keyData).patternChannels;
+  }
+
+  set patternChannels(patternChannels) {
+    const oldContext = contextMap.get(this.keyData);
+
+    const newContext = {
+      ...oldContext,
+      patternChannels,
+    };
+
+    contextMap.set(this.keyData, newContext);
+  }
+
   multi(batch = []) {
     this.batch = new Pipeline(this);
     // eslint-disable-next-line no-underscore-dangle
@@ -106,7 +148,7 @@ class RedisMock extends EventEmitter {
   }
 
   duplicate() {
-    return this.createConnectedClient()
+    return this.createConnectedClient();
   }
 
   // eslint-disable-next-line class-methods-use-this

diff --git a/test/command.js b/test/command.js
@@ -2,18 +2,6 @@ import _ from 'lodash';
 import Redis from 'ioredis';
 import command from '../src/command';
 
-// Ensure that we're getting the correct instance of Command when running in test:jest.js, as jest.js isn't designed to test code directly imported private functions like src/command
-jest.mock('ioredis', () => {
-  const { Command } = jest.requireActual('ioredis');
-  const RedisMock = jest.requireActual('../src/index');
-
-  return {
-    __esModule: true,
-    Command,
-    default: RedisMock,
-  };
-});
-
 describe('basic command', () => {
   const stub = command((...args) => args, 'testCommandName', {
     Command: { transformers: { argument: {}, reply: {} } },