Refactor: Fix Event Buffer Deadlock with State Machine (#1633)

* feat: implement state machine for chat synchronization and buffer management

* test: Add Jest configuration and tests for connection deadlock handling

- Created a Jest configuration file to set up testing environment.
- Implemented utility functions for creating isolated authentication sessions and mocking WebSocket connections.
- Added a test case to ensure that the connection does not deadlock when history sync is disabled, verifying the correct handling of message events.

* feat: add GitHub Actions workflow for running tests

* chore: sort import lint

* fix: implement timeout handling for AwaitingInitialSync state in chat socket, maybe fix memory leak

* feat: enhance chat synchronization by refining AwaitingInitialSync handling and adding history sync checks
This commit is contained in:
João Lucas de Oliveira Lopes
2025-08-06 19:15:41 -03:00
committed by GitHub
parent 1a721bb242
commit 812e53e4eb
13 changed files with 1137 additions and 717 deletions
+41
View File
@@ -0,0 +1,41 @@
name: Run Tests
on:
push:
branches:
- master
pull_request:
jobs:
run-tests:
runs-on: ubuntu-latest
timeout-minutes: 10
steps:
- name: Checkout code
uses: actions/checkout@v3
- name: Setup Node.js and Corepack
uses: actions/setup-node@v3
with:
node-version: 20.x
- name: Enable Corepack and Set Yarn Version
run: |
corepack enable
corepack prepare yarn@4.x --activate
- name: Restore Yarn Cache
uses: actions/cache@v3
id: yarn-cache
with:
path: .yarn/cache
key: ${{ runner.os }}-yarn-${{ hashFiles('**/yarn.lock') }}
restore-keys: |
${{ runner.os }}-yarn-
- name: Install dependencies
run: yarn install --immutable
- name: Run tests
run: yarn test
-14
View File
@@ -1,14 +0,0 @@
module.exports = {
'roots': [
'<rootDir>/src'
],
'testMatch': [
'**/Tests/test.*.+(ts|tsx|js)',
],
'transform': {
'^.+\\.(ts|tsx)$': 'ts-jest'
},
moduleNameMapper: {
'^axios$': require.resolve('axios'),
},
}
+13
View File
@@ -0,0 +1,13 @@
/** @jest-config-loader esbuild-register */
import type { Config } from 'jest'
const config: Config = {
roots: ['<rootDir>/src'],
testMatch: ['**.test.ts'],
extensionsToTreatAsEsm: ['.ts'],
transform: {
'^.+\\.(ts|tsx)$': ['ts-jest', { useESM: true }]
},
}
export default config
+4 -3
View File
@@ -36,7 +36,7 @@
"prepare": "npm run build",
"preinstall": "node ./engine-requirements.js",
"release": "release-it",
"test": "jest"
"test": "node --experimental-vm-modules ./node_modules/.bin/jest"
},
"dependencies": {
"@cacheable/node-cache": "^1.4.0",
@@ -59,10 +59,11 @@
"@typescript-eslint/parser": "^8.32.0",
"@whiskeysockets/eslint-config": "github:whiskeysockets/eslint-config",
"conventional-changelog-cli": "^2.2.2",
"esbuild-register": "^3.6.0",
"eslint": "^9.31.0",
"eslint-config-prettier": "^10.1.2",
"eslint-plugin-prettier": "^5.4.0",
"jest": "^29.7.0",
"jest": "^30.0.5",
"jimp": "^1.6.0",
"jiti": "^2.4.2",
"json": "^11.0.0",
@@ -71,7 +72,7 @@
"prettier": "^3.5.3",
"protobufjs-cli": "^1.1.3",
"release-it": "^15.10.3",
"ts-jest": "^29.3.2",
"ts-jest": "^29.4.0",
"tsc-esm-fix": "^3.1.2",
"tsx": "^4.20.3",
"typedoc": "^0.27.9",
+74 -33
View File
@@ -27,6 +27,7 @@ import type {
} from '../Types'
import { ALL_WA_PATCH_NAMES } from '../Types'
import type { LabelActionBody } from '../Types/Label'
import { SyncState } from '../Types/State'
import {
chatModificationToAppPatch,
type ChatMutationMap,
@@ -67,11 +68,14 @@ export const makeChatsSocket = (config: SocketConfig) => {
const { ev, ws, authState, generateMessageTag, sendNode, query, onUnexpectedError } = sock
let privacySettings: { [_: string]: string } | undefined
let needToFlushWithAppStateSync = false
let pendingAppStateSync = false
let syncState: SyncState = SyncState.Connecting
/** this mutex ensures that the notifications (receipts, messages etc.) are processed in order */
const processingMutex = makeMutex()
// Timeout for AwaitingInitialSync state
let awaitingSyncTimeout: NodeJS.Timeout | undefined
const placeholderResendCache =
config.placeholderResendCache ||
(new NodeCache<number>({
@@ -989,15 +993,42 @@ export const makeChatsSocket = (config: SocketConfig) => {
? shouldSyncHistoryMessage(historyMsg) && PROCESSABLE_HISTORY_TYPES.includes(historyMsg.syncType!)
: false
if (historyMsg && !authState.creds.myAppStateKeyId) {
logger.warn('skipping app state sync, as myAppStateKeyId is not set')
pendingAppStateSync = true
// State machine: decide on sync and flush
if (historyMsg && syncState === SyncState.AwaitingInitialSync) {
if (awaitingSyncTimeout) {
clearTimeout(awaitingSyncTimeout)
awaitingSyncTimeout = undefined
}
if (shouldProcessHistoryMsg) {
syncState = SyncState.Syncing
logger.info('Transitioned to Syncing state')
// Let doAppStateSync handle the final flush after it's done
} else {
syncState = SyncState.Online
logger.info('History sync skipped, transitioning to Online state and flushing buffer')
ev.flush()
}
}
const doAppStateSync = async () => {
if (syncState === SyncState.Syncing) {
logger.info('Doing app state sync')
await resyncAppState(ALL_WA_PATCH_NAMES, true)
// Sync is complete, go online and flush everything
syncState = SyncState.Online
logger.info('App state sync complete, transitioning to Online state and flushing buffer')
ev.flush()
const accountSyncCounter = (authState.creds.accountSyncCounter || 0) + 1
ev.emit('creds.update', { accountSyncCounter })
}
}
await Promise.all([
(async () => {
if (historyMsg && authState.creds.myAppStateKeyId) {
pendingAppStateSync = false
if (shouldProcessHistoryMsg) {
await doAppStateSync()
}
})(),
@@ -1012,24 +1043,10 @@ export const makeChatsSocket = (config: SocketConfig) => {
})
])
if (msg.message?.protocolMessage?.appStateSyncKeyShare && pendingAppStateSync) {
// If the app state key arrives and we are waiting to sync, trigger the sync now.
if (msg.message?.protocolMessage?.appStateSyncKeyShare && syncState === SyncState.Syncing) {
logger.info('App state sync key arrived, triggering app state sync')
await doAppStateSync()
pendingAppStateSync = false
}
async function doAppStateSync() {
if (!authState.creds.accountSyncCounter) {
logger.info('doing initial app state sync')
await resyncAppState(ALL_WA_PATCH_NAMES, true)
const accountSyncCounter = (authState.creds.accountSyncCounter || 0) + 1
ev.emit('creds.update', { accountSyncCounter })
if (needToFlushWithAppStateSync) {
logger.debug('flushing with app state sync')
ev.flush()
}
}
}
})
@@ -1072,16 +1089,40 @@ export const makeChatsSocket = (config: SocketConfig) => {
)
}
if (
receivedPendingNotifications && // if we don't have the app state key
// we keep buffering events until we finally have
// the key and can sync the messages
// todo scrutinize
!authState.creds?.myAppStateKeyId
) {
ev.buffer()
needToFlushWithAppStateSync = true
if (!receivedPendingNotifications || syncState !== SyncState.Connecting) {
return
}
syncState = SyncState.AwaitingInitialSync
logger.info('Connection is now AwaitingInitialSync, buffering events')
ev.buffer()
const willSyncHistory = shouldSyncHistoryMessage(
proto.Message.HistorySyncNotification.fromObject({
syncType: proto.HistorySync.HistorySyncType.RECENT
})
)
if (!willSyncHistory) {
logger.info('History sync is disabled by config, not waiting for notification. Transitioning to Online.')
syncState = SyncState.Online
setTimeout(() => ev.flush(), 0)
return
}
logger.info('History sync is enabled, awaiting notification with a 20s timeout.')
if (awaitingSyncTimeout) {
clearTimeout(awaitingSyncTimeout)
}
awaitingSyncTimeout = setTimeout(() => {
if (syncState === SyncState.AwaitingInitialSync) {
logger.warn('Timeout in AwaitingInitialSync, forcing state to Online and flushing buffer')
syncState = SyncState.Online
ev.flush()
}
}, 20_000)
})
return {
+12 -3
View File
@@ -3,10 +3,19 @@ import type { UserFacingSocketConfig } from '../Types'
import { makeCommunitiesSocket } from './communities'
// export the last socket layer
const makeWASocket = (config: UserFacingSocketConfig) =>
makeCommunitiesSocket({
const makeWASocket = (config: UserFacingSocketConfig) => {
const newConfig = {
...DEFAULT_CONNECTION_CONFIG,
...config
})
}
// If the user hasn't provided their own history sync function,
// let's create a default one that respects the syncFullHistory flag.
if (config.shouldSyncHistoryMessage === undefined) {
newConfig.shouldSyncHistoryMessage = () => !!newConfig.syncFullHistory
}
return makeCommunitiesSocket(newConfig)
}
export default makeWASocket
+11
View File
@@ -1,6 +1,17 @@
import { Boom } from '@hapi/boom'
import type { Contact } from './Contact'
export enum SyncState {
/** The socket is connecting, but we haven't received pending notifications yet. */
Connecting,
/** Pending notifications received. Buffering events until we decide whether to sync or not. */
AwaitingInitialSync,
/** The initial app state sync (history, etc.) is in progress. Buffering continues. */
Syncing,
/** Initial sync is complete, or was skipped. The socket is fully operational and events are processed in real-time. */
Online
}
export type WAConnectionState = 'open' | 'connecting' | 'close'
export type ConnectionState = {
+14 -22
View File
@@ -56,10 +56,9 @@ type BaileysBufferableEventEmitter = BaileysEventEmitter & {
createBufferedFunction<A extends any[], T>(work: (...args: A) => Promise<T>): (...args: A) => Promise<T>
/**
* flushes all buffered events
* @param force if true, will flush all data regardless of any pending buffers
* @returns returns true if the flush actually happened, otherwise false
*/
flush(force?: boolean): boolean
flush(): boolean
/** is there an ongoing buffer */
isBuffering(): boolean
}
@@ -74,7 +73,7 @@ export const makeEventBuffer = (logger: ILogger): BaileysBufferableEventEmitter
const historyCache = new Set<string>()
let data = makeBufferData()
let buffersInProgress = 0
let isBuffering = false
// take the generic event and fire it as a baileys event
ev.on('event', (map: BaileysEventData) => {
@@ -84,28 +83,22 @@ export const makeEventBuffer = (logger: ILogger): BaileysBufferableEventEmitter
})
function buffer() {
buffersInProgress += 1
if (!isBuffering) {
logger.info('Event buffer activated')
isBuffering = true
}
}
function flush(force = false) {
// no buffer going on
if (!buffersInProgress) {
function flush() {
if (!isBuffering) {
return false
}
if (!force) {
// reduce the number of buffers in progress
buffersInProgress -= 1
// if there are still some buffers going on
// then we don't flush now
if (buffersInProgress) {
return false
}
}
logger.info('Flushing event buffer')
isBuffering = false
const newData = makeBufferData()
const chatUpdates = Object.values(data.chatUpdates)
// gather the remaining conditional events so we re-queue them
let conditionalChatUpdatesLeft = 0
for (const update of chatUpdates) {
if (update.conditional) {
@@ -139,7 +132,7 @@ export const makeEventBuffer = (logger: ILogger): BaileysBufferableEventEmitter
}
},
emit<T extends BaileysEvent>(event: BaileysEvent, evData: BaileysEventMap[T]) {
if (buffersInProgress && BUFFERABLE_EVENT_SET.has(event)) {
if (isBuffering && BUFFERABLE_EVENT_SET.has(event)) {
append(data, historyCache, event as BufferableEvent, evData, logger)
return true
}
@@ -147,7 +140,7 @@ export const makeEventBuffer = (logger: ILogger): BaileysBufferableEventEmitter
return ev.emit('event', { [event]: evData })
},
isBuffering() {
return buffersInProgress > 0
return isBuffering
},
buffer,
flush,
@@ -155,10 +148,9 @@ export const makeEventBuffer = (logger: ILogger): BaileysBufferableEventEmitter
return async (...args) => {
buffer()
try {
const result = await work(...args)
return result
return await work(...args)
} finally {
flush()
// Flushing is now controlled centrally by the state machine.
}
}
},
+47
View File
@@ -0,0 +1,47 @@
import { jest } from '@jest/globals'
import { promises as fs } from 'fs'
import { tmpdir } from 'os'
import { join } from 'path'
import { useMultiFileAuthState } from '../..'
/**
* Creates a temporary, isolated authentication state for tests.
* This prevents tests from interfering with each other or with a real session.
* @returns An object with the authentication state and a cleanup function.
*/
export const makeSession = async () => {
// Create a temporary directory for the session files
const dir = join(tmpdir(), `baileys-test-session-${Date.now()}`)
await fs.mkdir(dir, { recursive: true })
// Use the multi-file auth state with the temporary directory
const { state, saveCreds } = await useMultiFileAuthState(dir)
return {
state,
saveCreds,
/**
* Cleans up the temporary session files.
* Call this at the end of your test.
*/
clear: async () => {
await fs.rm(dir, { recursive: true, force: true })
}
}
}
export const mockWebSocket = () => {
jest.mock('../../Socket/Client/websocket', () => {
return {
WebSocketClient: jest.fn().mockImplementation(() => ({
connect: jest.fn(() => Promise.resolve()),
close: jest.fn(),
on: jest.fn(),
off: jest.fn(),
emit: jest.fn(),
send: jest.fn(),
isOpen: true
}))
}
})
}
@@ -0,0 +1,54 @@
import { jest } from '@jest/globals'
import { proto } from '../..'
import { DEFAULT_CONNECTION_CONFIG } from '../../Defaults'
import makeWASocket from '../../Socket'
import { makeSession, mockWebSocket } from '../TestUtils/session'
mockWebSocket()
describe('Connection Deadlock Test', () => {
it('should not deadlock when history sync is disabled', async () => {
const { state, clear } = await makeSession()
state.creds.me = { id: '1234567890:1@s.whatsapp.net', name: 'Test User' }
const sock = makeWASocket({
...DEFAULT_CONNECTION_CONFIG,
auth: state,
shouldSyncHistoryMessage: () => false
})
const regularMessageListener = jest.fn()
sock.ev.on('messages.upsert', regularMessageListener)
// 1. Simulate receiving pending notifications. This activates the buffer.
sock.ev.emit('connection.update', { receivedPendingNotifications: true })
// 2. Now, emit a regular message. Because the previous step should have
// flushed the buffer, this message should be processed immediately.
const regularMessage = proto.WebMessageInfo.fromObject({
key: { remoteJid: '1234567890@s.whatsapp.net', fromMe: false, id: 'REGULAR_MSG_1' },
messageTimestamp: Date.now() / 1000,
message: { conversation: 'Hello, world!' }
})
sock.ev.emit('messages.upsert', { messages: [regularMessage], type: 'notify' })
// Wait for the event loop to process any final events.
await new Promise(resolve => setTimeout(resolve, 50))
// 3. Check if the regular message listener was called.
expect(regularMessageListener).toHaveBeenCalledTimes(1)
expect(regularMessageListener).toHaveBeenCalledWith(
expect.objectContaining({
messages: expect.arrayContaining([
expect.objectContaining({
key: regularMessage.key
})
]),
type: 'notify'
})
)
sock.end(new Error('Test completed'))
await clear()
})
})
+2 -1
View File
@@ -7,5 +7,6 @@
"declaration": true,
"declarationMap": true,
"sourceMap": true
}
},
"exclude": ["node_modules","src/__tests__"]
}
+1 -1
View File
@@ -32,7 +32,7 @@
// DX Boosters
"allowJs": false,
"checkJs": false
"checkJs": false,
},
"include": ["src/**/*.ts"],
"exclude": ["node_modules"]
+864 -640
View File
File diff suppressed because it is too large Load Diff