fix: sync client disconnect (#31461)

server/src/services/sync.service.spec.ts

@@ -1,6 +1,7 @@
 import { Writable } from 'node:stream';
 import { SyncEntityType } from 'src/enum.js';
 import { send } from 'src/services/sync.service.js';
+import { ClientDisconnectedError } from 'src/utils/response.js';
 import { serialize } from 'src/utils/sync.js';
 
 type TestStream = {
@@ -37,7 +38,7 @@ describe('send', () => {
     ids: ['now-id'] as [string],
   };
 
-  it('resolves immediately when the stream has capacity', async () => {
+  it('should resolve immediately when the stream has capacity', async () => {
     // A large highWaterMark means write() never signals backpressure for a
     // single small item.
     const { stream, chunks, flushNext } = createTestStream(1024 * 1024);
@@ -49,7 +50,7 @@ describe('send', () => {
     expect(chunks).toEqual([serialize(item)]);
   });
 
-  it('waits for the drain event before resolving when the stream signals backpressure', async () => {
+  it('should wait for the drain event when the stream signals backpressure', async () => {
     // A tiny highWaterMark means the very first write already exceeds
     // capacity, so write() returns false and send() must wait for 'drain'.
     const { stream, chunks, flushNext, pendingCount } = createTestStream(1);
@@ -75,4 +76,36 @@ describe('send', () => {
     expect(resolved).toBe(true);
     expect(chunks).toEqual([serialize(item)]);
   });
+
+  it('should throw a disconnect error when the stream destroyed', async () => {
+    const { stream, chunks } = createTestStream(1024 * 1024);
+
+    stream.destroy();
+
+    await expect(send(stream, item)).rejects.toBeInstanceOf(ClientDisconnectedError);
+    expect(chunks).toEqual([]);
+  });
+
+  it('should throw a disconnect error when the stream is destroyed after writing some data', async () => {
+    const { stream } = createTestStream(1);
+
+    const sendPromise = send(stream, item);
+
+    await Promise.resolve();
+    stream.destroy();
+
+    await expect(sendPromise).rejects.toBeInstanceOf(ClientDisconnectedError);
+  });
+
+  it('should handle a stream error', async () => {
+    const { stream } = createTestStream(1);
+    const error = new Error('socket hang up');
+
+    const sendPromise = send(stream, item);
+
+    await Promise.resolve();
+    stream.emit('error', error);
+
+    await expect(sendPromise).rejects.toBe(error);
+  });
 });

server/src/services/sync.service.ts

@@ -1,7 +1,6 @@
 import { BadRequestException, ForbiddenException, Injectable } from '@nestjs/common';
 import { Insertable } from 'kysely';
 import { DateTime, Duration } from 'luxon';
-import { once } from 'node:events';
 import { Writable } from 'node:stream';
 import { OnJob } from 'src/decorators.js';
 import { AuthDto } from 'src/dtos/auth.dto.js';
@@ -19,6 +18,7 @@ import { SessionSyncCheckpointTable } from 'src/schema/tables/sync-checkpoint.ta
 import { BaseService } from 'src/services/base.service.js';
 import type { SyncAck } from 'src/types.js';
 import { hexOrBufferToBase64 } from 'src/utils/bytes.js';
+import { ClientDisconnectedError, waitForDrain } from 'src/utils/response.js';
 import { fromAck, serialize, SerializeOptions, toAck } from 'src/utils/sync.js';
 
 type CheckpointMap = Partial<Record<SyncEntityType, SyncAck>>;
@@ -47,8 +47,13 @@ export const send = async <T extends keyof SyncItem, D extends SyncItem[T]>(
   response: Writable,
   item: SerializeOptions<T, D>,
 ) => {
+  if (response.destroyed || response.writableEnded) {
+    throw new ClientDisconnectedError();
+  }
+
+  // indicates back pressure, so we wait for 'drain' event
   if (!response.write(serialize(item))) {
-    await once(response, 'drain');
+    await waitForDrain(response);
   }
 };
 
@@ -136,6 +141,19 @@ export class SyncService extends BaseService {
   }
 
   async stream(auth: AuthDto, response: Writable, dto: SyncStreamDto) {
+    try {
+      await this.streamInternal(auth, response, dto);
+    } catch (error) {
+      if (error instanceof ClientDisconnectedError) {
+        this.logger.debug('Client closed the connection');
+        return;
+      }
+
+      throw error;
+    }
+  }
+
+  private async streamInternal(auth: AuthDto, response: Writable, dto: SyncStreamDto) {
     const session = auth.session;
     if (!session) {
       return throwSessionRequired();

server/src/utils/response.ts

@@ -1,8 +1,44 @@
 import { CookieOptions, Response } from 'express';
 import { Duration } from 'luxon';
+import { Writable } from 'node:stream';
 import { CookieResponse } from 'src/dtos/auth.dto.js';
 import { ImmichCookie } from 'src/enum.js';
 
+export class ClientDisconnectedError extends Error {}
+
+export const waitForDrain = (response: Writable) =>
+  new Promise<void>((resolve, reject) => {
+    const cleanup = () => {
+      response.off('drain', onDrain);
+      response.off('close', onClose);
+      response.off('error', onError);
+    };
+
+    const onDrain = () => {
+      cleanup();
+      resolve();
+    };
+
+    const onClose = () => {
+      cleanup();
+      reject(new ClientDisconnectedError());
+    };
+
+    const onError = (error: Error) => {
+      cleanup();
+      reject(error);
+    };
+
+    response.once('drain', onDrain);
+    response.once('close', onClose);
+    response.once('error', onError);
+
+    // 'close' may already have fired
+    if (response.destroyed) {
+      onClose();
+    }
+  });
+
 export const respondWithCookie = <T>(res: Response, body: T, { isSecure, values }: CookieResponse) => {
   const defaults: CookieOptions = {
     path: '/',

server/test/medium.factory.ts

@@ -812,7 +812,6 @@ const tagInsert = (tag: Partial<Insertable<TagTable>>) => {
 class CustomWritable extends Writable {
   private data = '';
 
-  // determined by Writable interface
   // eslint-disable-next-line unicorn/prefer-private-class-fields
   _write(chunk: any, encoding: string, callback: () => void) {
     this.data += chunk.toString();

server/test/medium/specs/sync/sync-client-disconnect.spec.ts

@@ -0,0 +1,69 @@
+import { Writable } from 'node:stream';
+import { SyncEntityType, SyncRequestType } from 'src/enum.js';
+import { SyncTestContext } from 'test/medium.factory.js';
+import { getKyselyDB } from 'test/utils.js';
+
+const ASSET_COUNT = 100;
+
+const setup = async () => {
+  const db = await getKyselyDB();
+  const ctx = new SyncTestContext(db);
+  const { auth, user, session } = await ctx.newSyncAuthUser();
+
+  for (let i = 0; i < ASSET_COUNT; i++) {
+    await ctx.newAsset({ ownerId: auth.user.id });
+  }
+
+  return { auth, user, session, ctx };
+};
+
+// automatically destroy the stream after X lines are written
+class LineAutoClosingStream extends Writable {
+  private count: number;
+  lines: string[] = [];
+
+  constructor({ after }: { after: number }) {
+    super({ highWaterMark: 1 });
+    this.count = after;
+  }
+
+  // eslint-disable-next-line unicorn/prefer-private-class-fields
+  _write(chunk: any, _encoding: string, callback: () => void) {
+    this.lines.push(chunk.toString());
+
+    if (this.lines.length >= this.count) {
+      this.destroy();
+      return;
+    }
+
+    setImmediate(callback);
+  }
+}
+
+describe('client disconnect', () => {
+  it('should stop streaming when the client disconnects mid-stream', async () => {
+    const { auth, ctx } = await setup();
+
+    const stream = new LineAutoClosingStream({ after: 5 });
+
+    await expect(ctx.sut.stream(auth, stream, { types: [SyncRequestType.AssetsV2] })).resolves.toBeUndefined();
+
+    expect(stream.lines).toHaveLength(5);
+    expect(stream.lines.length).toBeLessThan(ASSET_COUNT);
+  });
+
+  it('should release the database connection when the client disconnects', async () => {
+    const { auth, ctx } = await setup();
+
+    // more than the pool size
+    for (let i = 0; i < 12; i++) {
+      const stream = new LineAutoClosingStream({ after: 5 });
+      await ctx.sut.stream(auth, stream, { types: [SyncRequestType.AssetsV2] });
+    }
+
+    const response = await ctx.syncStream(auth, [SyncRequestType.UsersV1]);
+    expect(response).toEqual(
+      expect.arrayContaining([expect.objectContaining({ type: SyncEntityType.SyncCompleteV1 })]),
+    );
+  });
+});