diff --git a/app/lib/methods/sendMessage.test.ts b/app/lib/methods/sendMessage.test.ts new file mode 100644 index 0000000000..0a24bffb7a --- /dev/null +++ b/app/lib/methods/sendMessage.test.ts @@ -0,0 +1,262 @@ +import database from '../database'; +import log from './helpers/log'; +import { messagesStatus } from '../constants/messagesStatus'; +import { sendMessage } from './sendMessage'; + +type FakeRecord = Record; + +interface FakeCollection { + schema: Record; + records: Map; + find: (id: string) => Promise; + prepareCreate: (updater: (m: FakeRecord) => void) => FakeRecord; +} + +// Mirrors WatermelonDB's invariant: prepareUpdate on a record that already has a prepared change +// throws `Cannot update a record with pending changes` (Model/index.js). +const makeRecord = (debugName: string, fields: FakeRecord = {}): FakeRecord => { + const record: FakeRecord = { + ...fields, + __debugName: debugName, + _preparedState: null, + prepareUpdate(updater: (m: FakeRecord) => void) { + if (record._preparedState) { + throw new Error(`Cannot update a record with pending changes (${debugName})`); + } + updater(record); + record._preparedState = 'update'; + return record; + } + }; + return record; +}; + +const makeCollection = (name: string): FakeCollection => { + const collection: FakeCollection = { + schema: {}, + records: new Map(), + find: (id: string) => { + const existing = collection.records.get(id); + if (!existing) { + return Promise.reject(new Error(`Record ${name}#${id} not found`)); + } + return Promise.resolve(existing); + }, + prepareCreate: (updater: (m: FakeRecord) => void) => { + const record = makeRecord(`${name}#created`); + updater(record); + record._preparedState = 'create'; + // sanitizedRaw is mocked to identity below, so `_raw.id` is the client-generated id. + const id = record._raw?.id; + if (id) { + collection.records.set(id, record); + } + return record; + } + }; + return collection; +}; + +let collections: Record = {}; +const mockGetCollection = (name: string): FakeCollection => { + if (!collections[name]) { + collections[name] = makeCollection(name); + } + return collections[name]; +}; + +// db.batch commits prepared records, clearing their pending state (like the real writer). +const mockDbBatch = jest.fn((...args: any[]) => { + args.flat().forEach((item: FakeRecord) => { + if (item && typeof item === 'object' && '_preparedState' in item) { + item._preparedState = null; + } + }); + return Promise.resolve(undefined); +}); + +jest.mock('../database', () => { + let writerQueue: Promise = Promise.resolve(); + return { + __esModule: true, + default: { + active: { + get: (name: string) => mockGetCollection(name), + // Serialized writer lock, like WatermelonDB's. + write: (callback: () => Promise) => { + const run = writerQueue.then(() => callback()); + writerQueue = run.catch(() => undefined); + return run; + }, + batch: (...args: unknown[]) => mockDbBatch(...args) + } + } + }; +}); + +jest.mock('@nozbe/watermelondb/RawRecord', () => ({ + sanitizedRaw: (raw: unknown) => raw +})); + +const mockEncryptionGate: { promise: Promise | null } = { promise: null }; +jest.mock('../encryption', () => ({ + Encryption: { + encryptMessage: jest.fn(async (message: unknown) => { + if (mockEncryptionGate.promise) { + await mockEncryptionGate.promise; + } + return message; + }) + } +})); + +const mockPost = jest.fn, unknown[]>(() => Promise.resolve({ success: true, message: {} })); +jest.mock('../services/sdk', () => ({ + __esModule: true, + default: { + post: (...args: unknown[]) => mockPost(...args) + } +})); + +jest.mock('./helpers/log', () => ({ + __esModule: true, + default: jest.fn() +})); + +const db = (database as any).active; + +const deferred = () => { + let resolve: () => void = () => undefined; + const promise = new Promise(r => { + resolve = r; + }); + return { promise, resolve }; +}; + +// Let every already-queued microtask/promise chain settle. +const flush = () => new Promise(resolve => setImmediate(resolve)); + +const loggedPendingChanges = () => (log as jest.Mock).mock.calls.some(([error]) => /pending changes/.test(error?.message ?? '')); + +describe('sendMessage', () => { + const rid = 'GENERAL'; + const user = { id: 'userId', username: 'rocket.cat', name: 'Rocket Cat' }; + + beforeEach(() => { + jest.clearAllMocks(); + collections = {}; + mockEncryptionGate.promise = null; + mockPost.mockImplementation(() => Promise.resolve({ success: true, message: {} })); + }); + + describe('sendMessage', () => { + it('does not throw "pending changes" when a concurrent writer touches the subscription mid-send', async () => { + const subscriptions = mockGetCollection('subscriptions'); + const room = makeRecord(`subscriptions#${rid}`, { draftMessage: 'a draft' }); + subscriptions.records.set(rid, room); + + // Hold the writer lock while sendMessage is still encrypting, then update the very record + // sendMessage is about to prepare (the subscription whose draft gets cleared). + const concurrentGate = deferred(); + const concurrentWrite = db.write(async () => { + await concurrentGate.promise; + const record = await subscriptions.find(rid); + await db.batch([ + record.prepareUpdate((r: FakeRecord) => { + r.draftMessage = 'edited by another writer'; + }) + ]); + }); + + const encryption = deferred(); + mockEncryptionGate.promise = encryption.promise; + + const send = sendMessage(rid, 'hello', undefined, user); + + // Release encryption first so an unlocked implementation prepares its records now — before + // the concurrent writer runs — and holds them pending until its own batch. + encryption.resolve(); + await flush(); + concurrentGate.resolve(); + + await expect(Promise.all([concurrentWrite, send])).resolves.toBeDefined(); + + expect(loggedPendingChanges()).toBe(false); + const created = mockDbBatch.mock.calls + .flat(2) + .find((item: FakeRecord) => item?.status === messagesStatus.TEMP || item?.status === messagesStatus.SENT); + expect(created).toBeDefined(); + expect(created.msg).toBe('hello'); + expect(room.draftMessage).toBeFalsy(); + // Nothing was left prepared-but-uncommitted. + expect(created._preparedState).toBeNull(); + expect(room._preparedState).toBeNull(); + }); + }); + + describe('changeMessageStatus', () => { + it('does not throw "pending changes" when a concurrent writer touches the message mid-status-update', async () => { + const tmid = 'threadHeaderId'; + const messages = mockGetCollection('messages'); + messages.records.set( + tmid, + makeRecord(`messages#${tmid}`, { + msg: 'thread header', + ts: new Date(0), + tcount: 1, + u: { _id: 'other' }, + attachments: [] + }) + ); + mockGetCollection('threads'); + mockGetCollection('thread_messages'); + mockGetCollection('subscriptions'); + + // Block the server response so we can set up the race before changeMessageStatus runs. + const post = deferred(); + mockPost.mockImplementation(async () => { + await post.promise; + return { success: true, message: { mentions: [], channels: [] } }; + }); + + const send = sendMessage(rid, 'hi', tmid, user); + await flush(); + + // The message record created by the send — the one changeMessageStatus will update. + const messageId = [...messages.records.keys()].find(id => id !== tmid) as string; + expect(messageId).toBeDefined(); + const messageRecord = messages.records.get(messageId) as FakeRecord; + + const concurrentGate = deferred(); + const concurrentWrite = db.write(async () => { + await concurrentGate.promise; + await db.batch([ + messageRecord.prepareUpdate((m: FakeRecord) => { + m.msg = 'edited by another writer'; + }) + ]); + }); + + post.resolve(); + await flush(); + concurrentGate.resolve(); + + await expect(Promise.all([concurrentWrite, send])).resolves.toBeDefined(); + + expect(loggedPendingChanges()).toBe(false); + + const threadMessageRecord = mockGetCollection('thread_messages').records.get(messageId) as FakeRecord; + expect(messageRecord.status).toBe(messagesStatus.SENT); + expect(threadMessageRecord.status).toBe(messagesStatus.SENT); + + // The status update reached db.batch as one commit — it is the only batch holding both + // records — and neither was left prepared-but-uncommitted. + const statusBatch = mockDbBatch.mock.calls + .map(call => call.flat()) + .find(items => items.includes(messageRecord) && items.includes(threadMessageRecord)); + expect(statusBatch).toBeDefined(); + expect(messageRecord._preparedState).toBeNull(); + expect(threadMessageRecord._preparedState).toBeNull(); + }); + }); +}); diff --git a/app/lib/methods/sendMessage.ts b/app/lib/methods/sendMessage.ts index c7f4942f0d..f1015089ca 100644 --- a/app/lib/methods/sendMessage.ts +++ b/app/lib/methods/sendMessage.ts @@ -14,33 +14,34 @@ const changeMessageStatus = async (id: string, status: number, tmid?: string, me const db = database.active; const msgCollection = db.get('messages'); const threadMessagesCollection = db.get('thread_messages'); - const successBatch: Model[] = []; - const messageRecord = await msgCollection.find(id); - successBatch.push( - messageRecord.prepareUpdate(m => { - m.status = status; - if (message) { - m.mentions = message.mentions; - m.channels = message.channels; - } - }) - ); - - if (tmid) { - const threadMessageRecord = await threadMessagesCollection.find(id); - successBatch.push( - threadMessageRecord.prepareUpdate(tm => { - tm.status = status; - if (message) { - tm.mentions = message.mentions; - tm.channels = message.channels; - } - }) - ); - } try { await db.write(async () => { + const successBatch: Model[] = []; + const messageRecord = await msgCollection.find(id); + successBatch.push( + messageRecord.prepareUpdate(m => { + m.status = status; + if (message) { + m.mentions = message.mentions; + m.channels = message.channels; + } + }) + ); + + if (tmid) { + const threadMessageRecord = await threadMessagesCollection.find(id); + successBatch.push( + threadMessageRecord.prepareUpdate(tm => { + tm.status = status; + if (message) { + tm.mentions = message.mentions; + tm.channels = message.channels; + } + }) + ); + } + await db.batch(successBatch); }); } catch (error) { @@ -99,7 +100,6 @@ export async function sendMessage( const threadCollection = db.get('threads'); const threadMessagesCollection = db.get('thread_messages'); const messageId = random(17); - const batch: Model[] = []; const message = await Encryption.encryptMessage({ _id: messageId, @@ -110,120 +110,122 @@ export async function sendMessage( } as IMessage); const messageDate = new Date(); - let tMessageRecord: TMessageModel; - // If it's replying to a thread - if (tmid) { - try { - // Find thread message header in Messages collection - tMessageRecord = await msgCollection.find(tmid); - batch.push( - tMessageRecord.prepareUpdate(m => { - m.tlm = messageDate; - if (m.tcount) { - m.tcount += 1; + try { + await db.write(async () => { + const batch: Model[] = []; + let tMessageRecord: TMessageModel; + + // If it's replying to a thread + if (tmid) { + try { + // Find thread message header in Messages collection + tMessageRecord = await msgCollection.find(tmid); + batch.push( + tMessageRecord.prepareUpdate(m => { + m.tlm = messageDate; + if (m.tcount) { + m.tcount += 1; + } + }) + ); + + try { + // Find thread message header in Threads collection + await threadCollection.find(tmid); + } catch (error) { + // If there's no record, create one + batch.push( + threadCollection.prepareCreate(tm => { + tm._raw = sanitizedRaw({ id: tmid }, threadCollection.schema); + if (tm.subscription) { + tm.subscription.id = rid; + } + tm.tmid = tmid; + tm.msg = tMessageRecord.msg; + tm.ts = tMessageRecord.ts; + tm._updatedAt = messageDate; + tm.status = messagesStatus.SENT; // Original message was sent already + tm.u = tMessageRecord.u; + tm.t = message?.t as MessageType; + tm.attachments = tMessageRecord.attachments; + if (message?.t === E2E_MESSAGE_TYPE) { + tm.e2e = E2E_STATUS.DONE as E2EType; + } + }) + ); } - }) - ); - try { - // Find thread message header in Threads collection - await threadCollection.find(tmid); - } catch (error) { - // If there's no record, create one - batch.push( - threadCollection.prepareCreate(tm => { - tm._raw = sanitizedRaw({ id: tmid }, threadCollection.schema); - if (tm.subscription) { - tm.subscription.id = rid; - } - tm.tmid = tmid; - tm.msg = tMessageRecord.msg; - tm.ts = tMessageRecord.ts; - tm._updatedAt = messageDate; - tm.status = messagesStatus.SENT; // Original message was sent already - tm.u = tMessageRecord.u; - tm.t = message?.t as MessageType; - tm.attachments = tMessageRecord.attachments; - if (message?.t === E2E_MESSAGE_TYPE) { - tm.e2e = E2E_STATUS.DONE as E2EType; - } - }) - ); + // Create the message sent in ThreadMessages collection + batch.push( + threadMessagesCollection.prepareCreate(tm => { + tm._raw = sanitizedRaw({ id: messageId }, threadMessagesCollection.schema); + if (tm.subscription) { + tm.subscription.id = rid; + } + tm.rid = tmid; + tm.msg = msg; + tm.ts = messageDate; + tm._updatedAt = messageDate; + tm.status = messagesStatus.TEMP; + tm.u = { + _id: user.id || '1', + username: user.username, + name: user.name + }; + tm.t = message?.t as MessageType; + if (message?.t === E2E_MESSAGE_TYPE) { + tm.e2e = E2E_STATUS.DONE as E2EType; + } + }) + ); + } catch (e) { + log(e); + } } - // Create the message sent in ThreadMessages collection + // Create the message sent in Messages collection batch.push( - threadMessagesCollection.prepareCreate(tm => { - tm._raw = sanitizedRaw({ id: messageId }, threadMessagesCollection.schema); - if (tm.subscription) { - tm.subscription.id = rid; + msgCollection.prepareCreate(m => { + m._raw = sanitizedRaw({ id: messageId }, msgCollection.schema); + if (m.subscription) { + m.subscription.id = rid; } - tm.rid = tmid; - tm.msg = msg; - tm.ts = messageDate; - tm._updatedAt = messageDate; - tm.status = messagesStatus.TEMP; - tm.u = { + m.msg = msg; + m.ts = messageDate; + m._updatedAt = messageDate; + m.status = messagesStatus.TEMP; + m.u = { _id: user.id || '1', username: user.username, name: user.name }; - tm.t = message?.t as MessageType; + if (tmid && tMessageRecord) { + m.tmid = tmid; + // m.tlm = messageDate; // I don't think this is necessary... leaving it commented just in case... + m.tmsg = tMessageRecord.msg; + m.tshow = tshow; + } + m.t = message?.t as MessageType; if (message?.t === E2E_MESSAGE_TYPE) { - tm.e2e = E2E_STATUS.DONE as E2EType; + m.e2e = E2E_STATUS.DONE as E2EType; } }) ); - } catch (e) { - log(e); - } - } - // Create the message sent in Messages collection - batch.push( - msgCollection.prepareCreate(m => { - m._raw = sanitizedRaw({ id: messageId }, msgCollection.schema); - if (m.subscription) { - m.subscription.id = rid; - } - m.msg = msg; - m.ts = messageDate; - m._updatedAt = messageDate; - m.status = messagesStatus.TEMP; - m.u = { - _id: user.id || '1', - username: user.username, - name: user.name - }; - if (tmid && tMessageRecord) { - m.tmid = tmid; - // m.tlm = messageDate; // I don't think this is necessary... leaving it commented just in case... - m.tmsg = tMessageRecord.msg; - m.tshow = tshow; - } - m.t = message?.t as MessageType; - if (message?.t === E2E_MESSAGE_TYPE) { - m.e2e = E2E_STATUS.DONE as E2EType; + try { + const room = await subsCollection.find(rid); + if (room.draftMessage) { + batch.push( + room.prepareUpdate(r => { + r.draftMessage = null; + }) + ); + } + } catch (e) { + // Do nothing } - }) - ); - try { - const room = await subsCollection.find(rid); - if (room.draftMessage) { - batch.push( - room.prepareUpdate(r => { - r.draftMessage = null; - }) - ); - } - } catch (e) { - // Do nothing - } - - try { - await db.write(async () => { await db.batch(batch); }); } catch (e) {