@risingwave/wavelet-server 0.2.1 → 0.2.4

package/package.json

@@ -1,7 +1,16 @@
 {
   "name": "@risingwave/wavelet-server",
-  "version": "0.2.1",
+  "version": "0.2.4",
   "description": "Wavelet server - WebSocket fanout layer for RisingWave",
+  "homepage": "https://github.com/risingwavelabs/wavelet",
+  "repository": {
+    "type": "git",
+    "url": "https://github.com/risingwavelabs/wavelet.git",
+    "directory": "packages/server"
+  },
+  "bugs": {
+    "url": "https://github.com/risingwavelabs/wavelet/issues"
+  },
   "main": "./dist/index.js",
   "types": "./dist/index.d.ts",
   "scripts": {
@@ -13,7 +22,7 @@
     "pg": "^8.13.0",
     "ws": "^8.18.0",
     "jose": "^6.0.0",
-    "@risingwave/wavelet": "0.2.1"
+    "@risingwave/wavelet": "0.2.4"
   },
   "devDependencies": {
     "@types/pg": "^8.11.0",

package/src/__tests__/integration.test.ts

@@ -188,10 +188,12 @@ describe.runIf(process.env.WAVELET_INTEGRATION === '1')('Integration: Full Serve
   it('pushes diffs via WebSocket', async () => {
     const diff = await new Promise<any>((resolve, reject) => {
       const ws = new WebSocket(`ws://localhost:${port}/subscribe/${VIEW_NAME}`)
+      let sawSnapshot = false
 
       ws.on('message', (data: Buffer) => {
         const msg = JSON.parse(data.toString())
-        if (msg.type === 'connected') {
+        if (msg.type === 'snapshot') {
+          sawSnapshot = true
           // Write event to trigger diff
           fetch(`http://localhost:${port}/v1/events/${STREAM_NAME}`, {
             method: 'POST',
@@ -200,6 +202,7 @@ describe.runIf(process.env.WAVELET_INTEGRATION === '1')('Integration: Full Serve
           })
         }
         if (msg.type === 'diff') {
+          expect(sawSnapshot).toBe(true)
           ws.close()
           resolve(msg)
         }
@@ -214,4 +217,33 @@ describe.runIf(process.env.WAVELET_INTEGRATION === '1')('Integration: Full Serve
     // Should have at least one insert or update
     expect(diff.inserted.length + diff.updated.length).toBeGreaterThan(0)
   }, 20000)
+
+  it('pushes a snapshot immediately on connect', async () => {
+    await fetch(`http://localhost:${port}/v1/events/${STREAM_NAME}`, {
+      method: 'POST',
+      headers: { 'Content-Type': 'application/json' },
+      body: JSON.stringify({ user_id: 'snapshot_user', value: 11 }),
+    })
+
+    await new Promise(r => setTimeout(r, 2000))
+
+    const snapshot = await new Promise<any>((resolve, reject) => {
+      const ws = new WebSocket(`ws://localhost:${port}/subscribe/${VIEW_NAME}`)
+
+      ws.on('message', (data: Buffer) => {
+        const msg = JSON.parse(data.toString())
+        if (msg.type === 'snapshot') {
+          ws.close()
+          resolve(msg)
+        }
+      })
+
+      ws.on('error', reject)
+      setTimeout(() => { ws.close(); reject(new Error('WebSocket snapshot timeout')) }, 15000)
+    })
+
+    expect(snapshot.type).toBe('snapshot')
+    expect(Array.isArray(snapshot.rows)).toBe(true)
+    expect(snapshot.rows.some((row: any) => row.user_id === 'snapshot_user')).toBe(true)
+  }, 20000)
 })

package/src/__tests__/ws-fanout.test.ts

@@ -0,0 +1,143 @@
+import { EventEmitter } from 'node:events'
+import { describe, it, expect, vi } from 'vitest'
+import { WebSocket } from 'ws'
+import { WebSocketFanout } from '../ws-fanout.js'
+import type { BootstrapResult, ViewDiff } from '../cursor-manager.js'
+
+class MockSocket extends EventEmitter {
+  readyState: number = WebSocket.OPEN
+  sent: string[] = []
+
+  send(message: string): void {
+    this.sent.push(message)
+  }
+
+  ping(): void {}
+
+  close(): void {
+    this.readyState = WebSocket.CLOSED
+    this.emit('close')
+  }
+}
+
+function deferred<T>() {
+  let resolve!: (value: T) => void
+  let reject!: (reason?: unknown) => void
+  const promise = new Promise<T>((res, rej) => {
+    resolve = res
+    reject = rej
+  })
+  return { promise, resolve, reject }
+}
+
+describe('WebSocketFanout', () => {
+  it('drops queued shared diffs that are already covered by bootstrap', async () => {
+    const bootstrapDeferred = deferred<BootstrapResult>()
+    const cursorManager = {
+      bootstrap: vi.fn().mockReturnValue(bootstrapDeferred.promise),
+    } as any
+    const jwt = {
+      isConfigured: () => false,
+    } as any
+
+    const fanout = new WebSocketFanout(cursorManager, jwt, {
+      leaderboard: {} as any,
+    })
+
+    const ws = new MockSocket()
+    const req = {
+      url: '/subscribe/leaderboard',
+      headers: { host: 'localhost' },
+    } as any
+
+    const connectionPromise = (fanout as any).handleConnection(ws, req)
+
+    fanout.broadcast('leaderboard', {
+      cursor: '150',
+      inserted: [{ player_id: 'alice', score: 10 }],
+      updated: [],
+      deleted: [],
+    })
+    fanout.broadcast('leaderboard', {
+      cursor: '300',
+      inserted: [{ player_id: 'bob', score: 20 }],
+      updated: [],
+      deleted: [],
+    })
+
+    bootstrapDeferred.resolve({
+      snapshotRows: [{ player_id: 'alice', score: 10 }],
+      diffs: [{
+        cursor: '200',
+        inserted: [],
+        updated: [{ player_id: 'alice', score: 15 }],
+        deleted: [],
+      }],
+      lastCursor: '200',
+    })
+
+    await connectionPromise
+
+    const messages = ws.sent.map((message) => JSON.parse(message))
+    expect(messages).toEqual([
+      { type: 'connected', query: 'leaderboard' },
+      {
+        type: 'snapshot',
+        query: 'leaderboard',
+        rows: [{ player_id: 'alice', score: 10 }],
+      },
+      {
+        type: 'diff',
+        query: 'leaderboard',
+        cursor: '200',
+        inserted: [],
+        updated: [{ player_id: 'alice', score: 15 }],
+        deleted: [],
+      },
+      {
+        type: 'diff',
+        query: 'leaderboard',
+        cursor: '300',
+        inserted: [{ player_id: 'bob', score: 20 }],
+        updated: [],
+        deleted: [],
+      },
+    ])
+  })
+
+  it('filters snapshot rows with the same claim rule as diffs', async () => {
+    const cursorManager = {
+      bootstrap: vi.fn().mockResolvedValue({
+        snapshotRows: [
+          { user_id: 'u1', total: 10 },
+          { user_id: 'u2', total: 20 },
+        ],
+        diffs: [] as ViewDiff[],
+        lastCursor: null,
+      }),
+    } as any
+    const jwt = {
+      isConfigured: () => true,
+      verify: vi.fn().mockResolvedValue({ user_id: 'u1' }),
+    } as any
+
+    const fanout = new WebSocketFanout(cursorManager, jwt, {
+      totals: { filterBy: 'user_id' } as any,
+    })
+
+    const ws = new MockSocket()
+    const req = {
+      url: '/subscribe/totals?token=test-token',
+      headers: { host: 'localhost' },
+    } as any
+
+    await (fanout as any).handleConnection(ws, req)
+
+    const snapshotMessage = JSON.parse(ws.sent[1])
+    expect(snapshotMessage).toEqual({
+      type: 'snapshot',
+      query: 'totals',
+      rows: [{ user_id: 'u1', total: 10 }],
+    })
+  })
+})

package/src/cursor-manager.ts

@@ -16,6 +16,12 @@ export interface ViewDiff {
   deleted: Record<string, unknown>[]
 }
 
+export interface BootstrapResult {
+  snapshotRows: Record<string, unknown>[]
+  diffs: ViewDiff[]
+  lastCursor: string | null
+}
+
 type DiffCallback = (queryName: string, diff: ViewDiff) => void
 
 /**
@@ -122,9 +128,12 @@ export class CursorManager {
           }
         }
 
-        const diff = this.parseDiffs(allRows)
+        const diffs = this.parseDiffBatches(allRows)
 
-        if (diff.inserted.length > 0 || diff.updated.length > 0 || diff.deleted.length > 0) {
+        for (const diff of diffs) {
+          if (diff.inserted.length === 0 && diff.updated.length === 0 && diff.deleted.length === 0) {
+            continue
+          }
           callback(queryName, diff)
         }
       } catch (err: any) {
@@ -173,6 +182,88 @@ export class CursorManager {
     return diff
   }
 
+  parseDiffBatches(rows: any[]): ViewDiff[] {
+    const diffs: ViewDiff[] = []
+    let currentRows: any[] = []
+    let currentCursor: string | null = null
+
+    for (const row of rows) {
+      const cursor = this.normalizeCursor(row.rw_timestamp)
+      if (!cursor) continue
+
+      if (currentCursor !== null && cursor !== currentCursor) {
+        diffs.push(this.parseDiffs(currentRows))
+        currentRows = []
+      }
+
+      currentCursor = cursor
+      currentRows.push(row)
+    }
+
+    if (currentRows.length > 0) {
+      diffs.push(this.parseDiffs(currentRows))
+    }
+
+    return diffs
+  }
+
+  async bootstrap(queryName: string): Promise<BootstrapResult> {
+    const subName = this.subscriptions.get(queryName)
+    if (!subName) {
+      throw new Error(
+        `Subscription for query '${queryName}' is not initialized. Start the server before accepting WebSocket clients.`
+      )
+    }
+
+    const conn = new Client({ connectionString: this.connectionString })
+    await conn.connect()
+
+    const cursorName = `wavelet_boot_${Date.now()}_${Math.random().toString(36).slice(2, 8)}`
+
+    try {
+      await conn.query(`DECLARE ${cursorName} SUBSCRIPTION CURSOR FOR ${subName} FULL`)
+
+      const snapshotRows: Record<string, unknown>[] = []
+      const incrementalRows: any[] = []
+      let readingSnapshot = true
+
+      while (readingSnapshot) {
+        const result = await conn.query(`FETCH 1000 FROM ${cursorName}`)
+        if (result.rows.length === 0) break
+
+        let firstIncrementalIndex = result.rows.findIndex((row) => this.normalizeCursor(row.rw_timestamp) !== null)
+        if (firstIncrementalIndex === -1) firstIncrementalIndex = result.rows.length
+
+        for (const row of result.rows.slice(0, firstIncrementalIndex)) {
+          snapshotRows.push(this.stripSubscriptionMetadata(row))
+        }
+
+        if (firstIncrementalIndex < result.rows.length) {
+          incrementalRows.push(...result.rows.slice(firstIncrementalIndex))
+          readingSnapshot = false
+        }
+      }
+
+      while (true) {
+        const result = await conn.query(`FETCH 1000 FROM ${cursorName}`)
+        if (result.rows.length === 0) break
+        incrementalRows.push(...result.rows)
+      }
+
+      const diffs = this.parseDiffBatches(incrementalRows)
+      const lastCursor = diffs.length > 0 ? diffs[diffs.length - 1].cursor : null
+
+      return { snapshotRows, diffs, lastCursor }
+    } finally {
+      try {
+        await conn.query(`CLOSE ${cursorName}`)
+      } catch {}
+      try {
+        await conn.end()
+      } catch {}
+    }
+  }
+
   async query(sql: string): Promise<any[]> {
     if (!this.client) throw new Error('Not connected')
     const result = await this.client.query(sql)
@@ -206,4 +297,15 @@ export class CursorManager {
     } catch {}
     this.client = null
   }
+
+  private stripSubscriptionMetadata(row: Record<string, unknown>): Record<string, unknown> {
+    const { op: _op, rw_timestamp: _rwTimestamp, ...data } = row
+    return data
+  }
+
+  private normalizeCursor(value: unknown): string | null {
+    if (value === null || value === undefined) return null
+    const cursor = String(value).trim()
+    return cursor === '' ? null : cursor
+  }
 }

package/src/ws-fanout.ts

@@ -8,6 +8,8 @@ interface Subscriber {
   ws: WebSocket
   queryName: string
   claims: JwtClaims | null
+  ready: boolean
+  pendingDiffs: ViewDiff[]
 }
 
 export class WebSocketFanout {
@@ -80,7 +82,13 @@ export class WebSocketFanout {
       claims = await this.jwt.verify(token)
     }
 
-    const subscriber: Subscriber = { ws, queryName, claims }
+    const subscriber: Subscriber = {
+      ws,
+      queryName,
+      claims,
+      ready: false,
+      pendingDiffs: [],
+    }
 
     if (!this.subscribers.has(queryName)) {
       this.subscribers.set(queryName, new Set())
@@ -101,40 +109,84 @@ export class WebSocketFanout {
     ws.on('pong', () => { /* connection alive */ })
 
     ws.send(JSON.stringify({ type: 'connected', query: queryName }))
+
+    const bootstrap = await this.cursorManager.bootstrap(queryName)
+    const snapshotRows = this.filterSnapshotRows(queryName, bootstrap.snapshotRows, claims)
+    ws.send(JSON.stringify({
+      type: 'snapshot',
+      query: queryName,
+      rows: snapshotRows,
+    }))
+
+    for (const diff of bootstrap.diffs) {
+      const filteredDiff = this.filterDiffForSubscriber(queryName, diff, claims)
+      if (this.isEmptyDiff(filteredDiff)) continue
+      if (ws.readyState !== WebSocket.OPEN) break
+      ws.send(this.serializeDiffMessage(queryName, filteredDiff))
+    }
+
+    const handoffCursor = bootstrap.lastCursor
+    subscriber.ready = true
+    for (const diff of subscriber.pendingDiffs) {
+      if (ws.readyState !== WebSocket.OPEN) break
+      if (handoffCursor && this.compareCursor(diff.cursor, handoffCursor) <= 0) {
+        continue
+      }
+      ws.send(this.serializeDiffMessage(queryName, diff))
+    }
+    subscriber.pendingDiffs = []
   }
 
   broadcast(queryName: string, diff: ViewDiff): void {
     const subs = this.subscribers.get(queryName)
     if (!subs || subs.size === 0) return
 
-    const queryDef = this.queries[queryName]
-    const filterBy = this.getFilterBy(queryDef)
-
     for (const sub of subs) {
       if (sub.ws.readyState !== WebSocket.OPEN) continue
 
-      let filteredDiff = diff
-      if (filterBy && sub.claims) {
-        filteredDiff = this.filterDiff(diff, filterBy, sub.claims)
-      }
+      const filteredDiff = this.filterDiffForSubscriber(queryName, diff, sub.claims)
+      if (this.isEmptyDiff(filteredDiff)) continue
 
-      if (
-        filteredDiff.inserted.length === 0 &&
-        filteredDiff.updated.length === 0 &&
-        filteredDiff.deleted.length === 0
-      ) {
+      if (!sub.ready) {
+        sub.pendingDiffs.push(filteredDiff)
         continue
       }
 
-      sub.ws.send(JSON.stringify({
-        type: 'diff',
-        query: queryName,
-        cursor: filteredDiff.cursor,
-        inserted: filteredDiff.inserted,
-        updated: filteredDiff.updated,
-        deleted: filteredDiff.deleted,
-      }))
+      sub.ws.send(this.serializeDiffMessage(queryName, filteredDiff))
+    }
+  }
+
+  private filterSnapshotRows(
+    queryName: string,
+    rows: Record<string, unknown>[],
+    claims: JwtClaims | null
+  ): Record<string, unknown>[] {
+    const queryDef = this.queries[queryName]
+    const filterBy = this.getFilterBy(queryDef)
+
+    if (filterBy && claims) {
+      const claimValue = claims[filterBy]
+      if (claimValue === undefined) return []
+
+      return rows.filter((row) => String(row[filterBy]) === String(claimValue))
     }
+
+    return rows
+  }
+
+  private filterDiffForSubscriber(
+    queryName: string,
+    diff: ViewDiff,
+    claims: JwtClaims | null
+  ): ViewDiff {
+    const queryDef = this.queries[queryName]
+    const filterBy = this.getFilterBy(queryDef)
+
+    if (filterBy && claims) {
+      return this.filterDiff(diff, filterBy, claims)
+    }
+
+    return diff
   }
 
   private filterDiff(diff: ViewDiff, filterBy: string, claims: JwtClaims): ViewDiff {
@@ -160,6 +212,28 @@ export class WebSocketFanout {
     return (queryDef as QueryDef).filterBy
   }
 
+  private isEmptyDiff(diff: ViewDiff): boolean {
+    return diff.inserted.length === 0 && diff.updated.length === 0 && diff.deleted.length === 0
+  }
+
+  private serializeDiffMessage(queryName: string, diff: ViewDiff): string {
+    return JSON.stringify({
+      type: 'diff',
+      query: queryName,
+      cursor: diff.cursor,
+      inserted: diff.inserted,
+      updated: diff.updated,
+      deleted: diff.deleted,
+    })
+  }
+
+  private compareCursor(left: string, right: string): number {
+    const leftValue = BigInt(left)
+    const rightValue = BigInt(right)
+    if (leftValue === rightValue) return 0
+    return leftValue < rightValue ? -1 : 1
+  }
+
   closeAll(): void {
     for (const [, subs] of this.subscribers) {
       for (const sub of subs) {