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

src/Socket/chats.ts

@@ -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 {

src/Socket/index.ts

@@ -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

src/Types/State.ts

@@ -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 = {

src/Utils/event-buffer.ts

@@ -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.
 				}
 			}
 		},

src/__tests__/TestUtils/session.ts

@@ -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
+			}))
+		}
+	})
+}

src/__tests__/binary/connection-deadlock.test.ts

@@ -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()
+	})
+})