fix(maintenance): prevent enable/disable maintenance CLI hanging on occasion (#24713)

* fix(maintenance): prevent CLI hanging on occassion
fix(maintenance): always ack messages
fix(maintenance): ensure Redis is connected first

* chore(maintenance): validate app restart responses

* chore: mock the app restart callback

* fix: ack may not exist depending on caller

* refactor: move one shot into app.repository

* fix: send correct state in one shot

* chore: log restart event

server/src/maintenance/maintenance-websocket.repository.ts

@@ -37,7 +37,13 @@ export class MaintenanceWebsocketRepository implements OnGatewayConnection, OnGa
 
   afterInit(websocketServer: Server) {
     this.logger.log('Initialized websocket server');
-    websocketServer.on('AppRestart', () => this.appRepository.exitApp());
+
+    websocketServer.on('AppRestart', (event: ArgsOf<'AppRestart'>, ack?: (ok: 'ok') => void) => {
+      this.logger.log(`Restarting due to event... ${JSON.stringify(event)}`);
+
+      ack?.('ok');
+      this.appRepository.exitApp();
+    });
   }
 
   clientBroadcast<T extends keyof ClientEventMap>(event: T, ...data: ClientEventMap[T]) {

server/src/repositories/app.repository.ts

@@ -1,5 +1,10 @@
 import { Injectable } from '@nestjs/common';
+import { createAdapter } from '@socket.io/redis-adapter';
+import Redis from 'ioredis';
+import { Server as SocketIO } from 'socket.io';
 import { ExitCode } from 'src/enum';
+import { ConfigRepository } from 'src/repositories/config.repository';
+import { AppRestartEvent } from 'src/repositories/event.repository';
 
 @Injectable()
 export class AppRepository {
@@ -17,4 +22,26 @@ export class AppRepository {
   setCloseFn(fn: () => Promise<void>) {
     this.closeFn = fn;
   }
+
+  async sendOneShotAppRestart(state: AppRestartEvent): Promise<void> {
+    const server = new SocketIO();
+    const { redis } = new ConfigRepository().getEnv();
+    const pubClient = new Redis({ ...redis, lazyConnect: true });
+    const subClient = pubClient.duplicate();
+
+    await Promise.all([pubClient.connect(), subClient.connect()]);
+
+    server.adapter(createAdapter(pubClient, subClient));
+
+    // => corresponds to notification.service.ts#onAppRestart
+    server.emit('AppRestartV1', state, async () => {
+      const responses = await server.serverSideEmitWithAck('AppRestart', state);
+      if (responses.some((response) => response !== 'ok')) {
+        throw new Error("One or more node(s) returned a non-'ok' response to our restart request!");
+      }
+
+      pubClient.disconnect();
+      subClient.disconnect();
+    });
+  }
 }

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

@@ -89,6 +89,7 @@ describe(CliService.name, () => {
         alreadyDisabled: true,
       });
 
+      expect(mocks.app.sendOneShotAppRestart).toHaveBeenCalledTimes(0);
       expect(mocks.systemMetadata.set).toHaveBeenCalledTimes(0);
       expect(mocks.event.emit).toHaveBeenCalledTimes(0);
     });
@@ -99,6 +100,7 @@ describe(CliService.name, () => {
         alreadyDisabled: false,
       });
 
+      expect(mocks.app.sendOneShotAppRestart).toHaveBeenCalled();
       expect(mocks.systemMetadata.set).toHaveBeenCalledWith(SystemMetadataKey.MaintenanceMode, {
         isMaintenanceMode: false,
       });
@@ -114,6 +116,7 @@ describe(CliService.name, () => {
         }),
       );
 
+      expect(mocks.app.sendOneShotAppRestart).toHaveBeenCalledTimes(0);
       expect(mocks.systemMetadata.set).toHaveBeenCalledTimes(0);
       expect(mocks.event.emit).toHaveBeenCalledTimes(0);
     });
@@ -126,6 +129,7 @@ describe(CliService.name, () => {
         }),
       );
 
+      expect(mocks.app.sendOneShotAppRestart).toHaveBeenCalled();
       expect(mocks.systemMetadata.set).toHaveBeenCalledWith(SystemMetadataKey.MaintenanceMode, {
         isMaintenanceMode: true,
         secret: expect.stringMatching(/^\w{128}$/),

server/src/services/cli.service.ts

@@ -5,7 +5,7 @@ import { MaintenanceAuthDto } from 'src/dtos/maintenance.dto';
 import { UserAdminResponseDto, mapUserAdmin } from 'src/dtos/user.dto';
 import { SystemMetadataKey } from 'src/enum';
 import { BaseService } from 'src/services/base.service';
-import { createMaintenanceLoginUrl, generateMaintenanceSecret, sendOneShotAppRestart } from 'src/utils/maintenance';
+import { createMaintenanceLoginUrl, generateMaintenanceSecret } from 'src/utils/maintenance';
 import { getExternalDomain } from 'src/utils/misc';
 
 @Injectable()
@@ -55,8 +55,7 @@ export class CliService extends BaseService {
 
     const state = { isMaintenanceMode: false as const };
     await this.systemMetadataRepository.set(SystemMetadataKey.MaintenanceMode, state);
-
+    await this.appRepository.sendOneShotAppRestart(state);
-    sendOneShotAppRestart(state);
 
     return {
       alreadyDisabled: false,
@@ -89,7 +88,7 @@ export class CliService extends BaseService {
       secret,
     });
 
-    sendOneShotAppRestart({
+    await this.appRepository.sendOneShotAppRestart({
       isMaintenanceMode: true,
     });
 

server/src/services/maintenance.service.ts

@@ -2,6 +2,7 @@ import { Injectable } from '@nestjs/common';
 import { OnEvent } from 'src/decorators';
 import { MaintenanceAuthDto } from 'src/dtos/maintenance.dto';
 import { SystemMetadataKey } from 'src/enum';
+import { ArgOf } from 'src/repositories/event.repository';
 import { BaseService } from 'src/services/base.service';
 import { MaintenanceModeState } from 'src/types';
 import { createMaintenanceLoginUrl, generateMaintenanceSecret, signMaintenanceJwt } from 'src/utils/maintenance';
@@ -31,7 +32,10 @@ export class MaintenanceService extends BaseService {
   }
 
   @OnEvent({ name: 'AppRestart', server: true })
-  onRestart(): void {
+  onRestart(event: ArgOf<'AppRestart'>, ack?: (ok: 'ok') => void): void {
+    this.logger.log(`Restarting due to event... ${JSON.stringify(event)}`);
+
+    ack?.('ok');
     this.appRepository.exitApp();
   }
 

server/src/utils/maintenance.ts

@@ -1,55 +1,6 @@
-import { createAdapter } from '@socket.io/redis-adapter';
-import Redis from 'ioredis';
 import { SignJWT } from 'jose';
 import { randomBytes } from 'node:crypto';
-import { Server as SocketIO } from 'socket.io';
 import { MaintenanceAuthDto } from 'src/dtos/maintenance.dto';
-import { ConfigRepository } from 'src/repositories/config.repository';
-import { AppRestartEvent } from 'src/repositories/event.repository';
-
-export function sendOneShotAppRestart(state: AppRestartEvent): void {
-  const server = new SocketIO();
-  const { redis } = new ConfigRepository().getEnv();
-  const pubClient = new Redis(redis);
-  const subClient = pubClient.duplicate();
-  server.adapter(createAdapter(pubClient, subClient));
-
-  /**
-   * Keep trying until we manage to stop Immich
-   *
-   * Sometimes there appear to be communication
-   * issues between to the other servers.
-   *
-   * This issue only occurs with this method.
-   */
-  async function tryTerminate() {
-    while (true) {
-      try {
-        const responses = await server.serverSideEmitWithAck('AppRestart', state);
-        if (responses.length > 0) {
-          return;
-        }
-      } catch (error) {
-        console.error(error);
-        console.error('Encountered an error while telling Immich to stop.');
-      }
-
-      console.info(
-        "\nIt doesn't appear that Immich stopped, trying again in a moment.\nIf Immich is already not running, you can ignore this error.",
-      );
-
-      await new Promise((r) => setTimeout(r, 1e3));
-    }
-  }
-
-  // => corresponds to notification.service.ts#onAppRestart
-  server.emit('AppRestartV1', state, () => {
-    void tryTerminate().finally(() => {
-      pubClient.disconnect();
-      subClient.disconnect();
-    });
-  });
-}
 
 export async function createMaintenanceLoginUrl(
   baseUrl: string,