chore(server): refactor locks (#5953)

* lock refactor

* add mocks

* add await

* move database repo injection to service

* update tests

* add mock implementation

* remove unused imports

* this

server/src/infra/database-locks.ts

@@ -1,41 +0,0 @@
-import { dataSource } from '@app/infra';
-import AsyncLock from 'async-lock';
-export enum DatabaseLock {
-  GeodataImport = 100,
-  CLIPDimSize = 512,
-}
-
-export async function acquireLock(lock: DatabaseLock): Promise<void> {
-  return dataSource.query('SELECT pg_advisory_lock($1)', [lock]);
-}
-
-export async function releaseLock(lock: DatabaseLock): Promise<void> {
-  return dataSource.query('SELECT pg_advisory_unlock($1)', [lock]);
-}
-
-export const asyncLock = new AsyncLock();
-
-export function RequireLock<T>(
-  lock: DatabaseLock,
-): (target: any, propertyKey: string, descriptor: PropertyDescriptor) => void {
-  return function (target: any, propertyKey: string, descriptor: PropertyDescriptor): void {
-    const originalMethod = descriptor.value;
-    descriptor.value = async function (...args: any[]): Promise<T> {
-      if (!dataSource.isInitialized) {
-        await dataSource.initialize();
-      }
-
-      let res;
-      await asyncLock.acquire(DatabaseLock[lock], async () => {
-        try {
-          await acquireLock(lock);
-          res = await originalMethod.apply(this, args);
-        } finally {
-          await releaseLock(lock);
-        }
-      });
-
-      return res as any;
-    };
-  };
-}

server/src/infra/index.ts

@@ -1,4 +1,3 @@
-export * from './database-locks';
 export * from './database.config';
 export * from './infra.config';
 export * from './infra.module';

server/src/infra/repositories/database.repository.ts

@@ -1,10 +1,13 @@
-import { DatabaseExtension, IDatabaseRepository, Version } from '@app/domain';
+import { DatabaseExtension, DatabaseLock, IDatabaseRepository, Version } from '@app/domain';
 import { Injectable } from '@nestjs/common';
 import { InjectDataSource } from '@nestjs/typeorm';
+import AsyncLock from 'async-lock';
 import { DataSource } from 'typeorm';
 
 @Injectable()
 export class DatabaseRepository implements IDatabaseRepository {
+  readonly asyncLock = new AsyncLock();
+
   constructor(@InjectDataSource() private dataSource: DataSource) {}
 
   async getExtensionVersion(extension: DatabaseExtension): Promise<Version | null> {
@@ -25,4 +28,34 @@ export class DatabaseRepository implements IDatabaseRepository {
   async runMigrations(options?: { transaction?: 'all' | 'none' | 'each' }): Promise<void> {
     await this.dataSource.runMigrations(options);
   }
+
+  async withLock<R>(lock: DatabaseLock, callback: () => Promise<R>): Promise<R> {
+    let res;
+    await this.asyncLock.acquire(DatabaseLock[lock], async () => {
+      try {
+        await this.acquireLock(lock);
+        res = await callback();
+      } finally {
+        await this.releaseLock(lock);
+      }
+    });
+
+    return res as R;
+  }
+
+  isBusy(lock: DatabaseLock): boolean {
+    return this.asyncLock.isBusy(DatabaseLock[lock]);
+  }
+
+  async wait(lock: DatabaseLock): Promise<void> {
+    await this.asyncLock.acquire(DatabaseLock[lock], () => {});
+  }
+
+  private async acquireLock(lock: DatabaseLock): Promise<void> {
+    return this.dataSource.query('SELECT pg_advisory_lock($1)', [lock]);
+  }
+
+  private async releaseLock(lock: DatabaseLock): Promise<void> {
+    return this.dataSource.query('SELECT pg_advisory_unlock($1)', [lock]);
+  }
 }

server/src/infra/repositories/metadata.repository.ts

@@ -5,7 +5,6 @@ import {
   ISystemMetadataRepository,
   ReverseGeocodeResult,
 } from '@app/domain';
-import { DatabaseLock, RequireLock } from '@app/infra';
 import { GeodataAdmin1Entity, GeodataAdmin2Entity, GeodataPlacesEntity, SystemMetadataKey } from '@app/infra/entities';
 import { ImmichLogger } from '@app/infra/logger';
 import { Inject } from '@nestjs/common';
@@ -34,7 +33,6 @@ export class MetadataRepository implements IMetadataRepository {
 
   private logger = new ImmichLogger(MetadataRepository.name);
 
-  @RequireLock(DatabaseLock.GeodataImport)
   async init(): Promise<void> {
     this.logger.log('Initializing metadata repository');
     const geodataDate = await readFile('/usr/src/resources/geodata-date.txt', 'utf8');
@@ -46,7 +44,17 @@ export class MetadataRepository implements IMetadataRepository {
     }
 
     this.logger.log('Importing geodata to database from file');
+    await this.importGeodata();
 
+    await this.systemMetadataRepository.set(SystemMetadataKey.REVERSE_GEOCODING_STATE, {
+      lastUpdate: geodataDate,
+      lastImportFileName: CITIES_FILE,
+    });
+
+    this.logger.log('Geodata import completed');
+  }
+
+  private async importGeodata() {
     const queryRunner = this.dataSource.createQueryRunner();
     await queryRunner.connect();
 
@@ -65,13 +73,6 @@ export class MetadataRepository implements IMetadataRepository {
     } finally {
       await queryRunner.release();
     }
-
-    await this.systemMetadataRepository.set(SystemMetadataKey.REVERSE_GEOCODING_STATE, {
-      lastUpdate: geodataDate,
-      lastImportFileName: CITIES_FILE,
-    });
-
-    this.logger.log('Geodata import completed');
   }
 
   private async loadGeodataToTableFromFile<T extends GeoEntity>(

server/src/infra/repositories/smart-info.repository.ts

@@ -1,6 +1,5 @@
 import { Embedding, EmbeddingSearch, ISmartInfoRepository } from '@app/domain';
 import { getCLIPModelInfo } from '@app/domain/smart-info/smart-info.constant';
-import { DatabaseLock, RequireLock, asyncLock } from '@app/infra';
 import { AssetEntity, AssetFaceEntity, SmartInfoEntity, SmartSearchEntity } from '@app/infra/entities';
 import { ImmichLogger } from '@app/infra/logger';
 import { Injectable } from '@nestjs/common';
@@ -121,18 +120,12 @@ export class SmartInfoRepository implements ISmartInfoRepository {
   }
 
   private async upsertEmbedding(assetId: string, embedding: number[]): Promise<void> {
-    if (asyncLock.isBusy(DatabaseLock[DatabaseLock.CLIPDimSize])) {
-      this.logger.verbose(`Waiting for CLIP dimension size to be updated`);
-      await asyncLock.acquire(DatabaseLock[DatabaseLock.CLIPDimSize], () => {});
-    }
-
     await this.smartSearchRepository.upsert(
       { assetId, embedding: () => asVector(embedding, true) },
       { conflictPaths: ['assetId'] },
     );
   }
 
-  @RequireLock(DatabaseLock.CLIPDimSize)
   private async updateDimSize(dimSize: number): Promise<void> {
     if (!isValidInteger(dimSize, { min: 1, max: 2 ** 16 })) {
       throw new Error(`Invalid CLIP dimension size: ${dimSize}`);