feat: allow worker thread persist

Author: stanley2058

src/api.js

@@ -118,6 +118,8 @@ export class Api {
     this.redisWorkerGroupName = this.prefix + ':worker'
     this.workerSetName = `${this.prefix}:worker:${this.consumername}:idset`
     this._destroyed = false
+    /** @type {import('worker_threads').Worker | null} */
+    this.persistWorker = null
 
     const addScript = WORKER_DISABLED
       ? 'redis.call("XADD", KEYS[1], "*", "m", ARGV[1])'

src/index.js

@@ -3,4 +3,5 @@ export * from './server.js'
 export * from './storage.js'
 export * from './api.js'
 export * from './subscriber.js'
+export * from './persist-worker-thread.js'
 export * from './y-socket-io/index.js'

src/persist-worker-thread.js

@@ -0,0 +1,43 @@
+import * as Y from 'yjs'
+import * as logging from 'lib0/logging'
+import { Worker, isMainThread, parentPort, workerData } from 'worker_threads'
+import path from 'path'
+
+export class PersistWorkerThread {
+  /**
+   * @private
+   * @readonly
+   */
+  log = logging.createModuleLogger('@y/persist-worker-thread')
+
+  /**
+   * @param {import('./storage.js').AbstractStorage} store
+   */
+  constructor(store) {
+    if (isMainThread) {
+      this.log('persist worker cannot run on main thread')
+      return
+    }
+    this.store = store
+    parentPort?.on('message', this.persist)
+  }
+
+  /**
+   * @param {{ room: string, docstate: SharedArrayBuffer }} props
+   */
+  persist = async ({ room, docstate }) => {
+    const state = new Uint8Array(docstate)
+    const doc = new Y.Doc()
+    Y.applyUpdateV2(doc, state)
+    await this.store?.persistDoc(room, 'index', doc)
+    doc.destroy()
+  }
+}
+
+/**
+ * @param {import('./storage.js').AbstractStorage} store
+ */
+export function createPersistWorkerThread(store) {
+  if (isMainThread) throw new Error('cannot create persist worker in main thread')
+  return new PersistWorkerThread(store)
+}

src/socketio.js

@@ -37,9 +37,10 @@ class YSocketIOServer {
  * @param {string} [conf.redisPrefix]
  * @param {string} [conf.redisUrl]
  * @param {import('./y-socket-io/y-socket-io.js').YSocketIOConfiguration['authenticate']} conf.authenticate
+ * @param {import('worker_threads').Worker=} [conf.persistWorker]
  */
-export const registerYSocketIOServer = async (io, store, { authenticate, redisUrl, redisPrefix }) => {
+export const registerYSocketIOServer = async (io, store, { authenticate, redisUrl, redisPrefix, persistWorker }) => {
   const app = new YSocketIO(io, { authenticate })
-  const { client, subscriber } = await app.initialize(store, { redisUrl, redisPrefix })
+  const { client, subscriber } = await app.initialize(store, { redisUrl, redisPrefix, persistWorker })
   return new YSocketIOServer(app, client, subscriber)
 }

src/y-socket-io/y-socket-io.js

@@ -136,6 +136,12 @@ export class YSocketIO {
    * @readonly
    */
   namespacePersistentMap = new Map()
+  /**
+   * @type {Map<string, () => void>}
+   * @private
+   * @readonly
+   */
+  awaitingPersistMap = new Map()
 
   /**
    * YSocketIO constructor.
@@ -156,16 +162,20 @@ export class YSocketIO {
    *
    *  It also starts socket connection listeners.
    * @param {import('../storage.js').AbstractStorage} store
-   * @param {{ redisPrefix?: string, redisUrl?: string }=} opts
+   * @param {{ redisPrefix?: string, redisUrl?: string, persistWorker?: import('worker_threads').Worker }=} opts
    * @public
    */
-  async initialize (store, { redisUrl, redisPrefix = 'y' } = {}) {
+  async initialize (store, { redisUrl, redisPrefix = 'y', persistWorker } = {}) {
     const [client, subscriber] = await promise.all([
       api.createApiClient(store, { redisUrl, redisPrefix }),
       createSubscriber(store, { redisUrl, redisPrefix })
     ])
     this.client = client
     this.subscriber = subscriber
+    if (persistWorker) {
+      this.client.persistWorker = persistWorker
+      this.registerPersistWorkerResolve()
+    }
 
     this.nsp = this.io.of(/^\/yjs\|.*$/)
 
@@ -475,7 +485,24 @@ export class YSocketIO {
           assert(this.client)
           const doc = this.debouncedPersistDocMap.get(namespace)
           if (!doc) return
-          await this.client.store.persistDoc(namespace, 'index', doc)
+          if (this.client.persistWorker) {
+            /** @type {Promise<void>} */
+            const promise = new Promise((res) => {
+              assert(this.client?.persistWorker)
+              this.awaitingPersistMap.set(namespace, res)
+
+              const docState = Y.encodeStateAsUpdateV2(doc)
+              const buf = new Uint8Array(new SharedArrayBuffer(docState.length))
+              buf.set(docState)
+              this.client.persistWorker.postMessage({
+                room: namespace,
+                docstate: buf
+              })
+            })
+            await promise
+          } else {
+            await this.client.store.persistDoc(namespace, 'index', doc)
+          }
           await this.client.trimRoomStream(namespace, 'index', true)
           this.debouncedPersistDocMap.delete(namespace)
           this.debouncedPersistMap.delete(namespace)
@@ -569,4 +596,11 @@ export class YSocketIO {
       console.error(e)
     }
   }
+
+  registerPersistWorkerResolve () {
+    if (!this.client?.persistWorker) return
+    this.client.persistWorker.on('message', ({ event, room }) => {
+      if (event === 'persisted') this.awaitingPersistMap.get(room)?.()
+    })
+  }
 }