import { Injectable } from '@nestjs/common'; import { ModuleRef, Reflector } from '@nestjs/core'; import { ClassConstructor } from 'class-transformer'; import { makeWorkerUtils, run, Runner, TaskSpec, WorkerUtils } from 'graphile-worker'; import { Kysely } from 'kysely'; import { DateTime, Duration } from 'luxon'; import { InjectKysely } from 'nestjs-kysely'; import pg, { PoolConfig } from 'pg'; import { DB } from 'src/db'; import { GenerateSql, JobConfig } from 'src/decorators'; import { JobName, JobStatus, MetadataKey, QueueName, SystemMetadataKey } from 'src/enum'; import { ConfigRepository } from 'src/repositories/config.repository'; import { EventRepository } from 'src/repositories/event.repository'; import { LoggingRepository } from 'src/repositories/logging.repository'; import { SystemMetadataRepository } from 'src/repositories/system-metadata.repository'; import { JobCounts, JobItem, JobOf, QueueStatus } from 'src/types'; import { asPostgresConnectionConfig } from 'src/utils/database'; import { getKeyByValue, getMethodNames, ImmichStartupError } from 'src/utils/misc'; type JobMapItem = { jobName: JobName; queueName: QueueName; handler: (job: JobOf) => Promise; label: string; }; type QueueConfiguration = { paused: boolean; concurrency: number; }; @Injectable() export class JobRepository { private handlers: Partial> = {}; // todo inject the pg pool private pool?: pg.Pool; // todo inject worker utils? private workerUtils?: WorkerUtils; private queueConfig: Record = {}; private runners: Record = {}; constructor( @InjectKysely() private db: Kysely, private logger: LoggingRepository, private moduleRef: ModuleRef, private eventRepository: EventRepository, private configRepository: ConfigRepository, private systemMetadataRepository: SystemMetadataRepository, ) { logger.setContext(JobRepository.name); } async setup(services: ClassConstructor[]) { const reflector = this.moduleRef.get(Reflector, { strict: false }); for (const service of services) { const instance = this.moduleRef.get(service); for (const methodName of getMethodNames(instance)) { const handler = instance[methodName]; const config = reflector.get(MetadataKey.JOB_CONFIG, handler); if (!config) { continue; } const { name: jobName, queue: queueName } = config; const label = `${service.name}.${handler.name}`; // one handler per job if (this.handlers[jobName]) { const jobKey = getKeyByValue(JobName, jobName); const errorMessage = `Failed to add job handler for ${label}`; this.logger.error( `${errorMessage}. JobName.${jobKey} is already handled by ${this.handlers[jobName].label}.`, ); throw new ImmichStartupError(errorMessage); } this.handlers[jobName] = { label, jobName, queueName, handler: handler.bind(instance), }; this.logger.verbose(`Added job handler: ${jobName} => ${label}`); } } const { database } = this.configRepository.getEnv(); const pool = new pg.Pool({ ...asPostgresConnectionConfig(database.config), max: 100, } as PoolConfig); // todo: remove debug info setInterval(() => { this.logger.log(`connections: total: ${pool.totalCount} idle: ${pool.idleCount} waiting: ${pool.waitingCount}`); }, 5000); pool.on('connect', (client) => { client.setMaxListeners(200); }); this.pool = pool; this.workerUtils = await makeWorkerUtils({ pgPool: pool }); } async start(queueName: QueueName, concurrency?: number): Promise { if (concurrency) { this.queueConfig[queueName] = { ...this.queueConfig[queueName], concurrency, }; } else { concurrency = this.queueConfig[queueName].concurrency; } if (this.queueConfig[queueName].paused) { return; } await this.stop(queueName); this.runners[queueName] = await run({ concurrency, taskList: { [queueName]: async (payload: unknown): Promise => { this.logger.log(`Job ${queueName} started with payload: ${JSON.stringify(payload)}`); await this.eventRepository.emit('job.start', queueName, payload as JobItem); }, }, pgPool: this.pool, }); } async stop(queueName: QueueName): Promise { const runner = this.runners[queueName]; if (runner) { await runner.stop(); delete this.runners[queueName]; } } async pause(queueName: QueueName): Promise { await this.setState(queueName, true); await this.stop(queueName); } async resume(queueName: QueueName): Promise { await this.setState(queueName, false); await this.start(queueName); } private async setState(queueName: QueueName, paused: boolean): Promise { const state = await this.systemMetadataRepository.get(SystemMetadataKey.QUEUES_STATE); await this.systemMetadataRepository.set(SystemMetadataKey.QUEUES_STATE, { ...state, [queueName]: { paused }, }); this.queueConfig[queueName] = { ...this.queueConfig[queueName], paused, }; } // todo: we should consolidate queue and job names and have queues be // homogenous. // // the reason there are multiple kinds of jobs per queue is so that // concurrency settings apply to all of them. We could instead create a // concept of "queue" groups, such that workers will run for groups of queues // rather than just a single queue and achieve the same outcome. private getQueueName(name: JobName) { return (this.handlers[name] as JobMapItem).queueName; } async run({ name, data }: JobItem): Promise { const item = this.handlers[name as JobName]; if (!item) { this.logger.warn(`Skipping unknown job: "${name}"`); return JobStatus.SKIPPED; } return item.handler(data); } async queue(item: JobItem): Promise { this.logger.log(`Queueing job: ${this.getQueueName(item.name)}, data: ${JSON.stringify(item)}`); await this.workerUtils!.addJob(this.getQueueName(item.name), item, this.getJobOptions(item)); } async queueAll(items: JobItem[]): Promise { for (const item of items) { await this.queue(item); } } // todo: are we actually generating sql async clear(name: QueueName): Promise { await this.db .deleteFrom('graphile_worker._private_jobs') .where(({ eb, selectFrom }) => eb('task_id', 'in', selectFrom('graphile_worker._private_tasks').select('id').where('identifier', '=', name)), ) .execute(); const workers = await this.db .selectFrom('graphile_worker.jobs') .select('locked_by') .where('locked_by', 'is not', null) .distinct() .execute(); // Potentially dangerous? It helps if jobs get stuck active though. The // documentation says that stuck jobs will be unlocked automatically after 4 // hours. Though, it can be strange to click "clear" in the UI and see // nothing happen. Especially as the UI is binary, such that new jobs cannot // usually be scheduled unless both active and waiting are zero. await this.workerUtils!.forceUnlockWorkers(workers.map((worker) => worker.locked_by!)); } async clearFailed(name: QueueName): Promise { await this.db .deleteFrom('graphile_worker._private_jobs') .where(({ eb, selectFrom }) => eb( 'task_id', 'in', selectFrom('graphile_worker._private_tasks') .select('id') .where((eb) => eb.and([eb('identifier', '=', name), eb('attempts', '>=', eb.ref('max_attempts'))])), ), ) .execute(); } // todo: are we actually generating sql @GenerateSql({ params: [] }) async getJobCounts(name: QueueName): Promise { return await this.db .selectFrom('graphile_worker.jobs') .select((eb) => [ eb.fn .countAll() .filterWhere((eb) => eb.and([eb('task_identifier', '=', name), eb('locked_by', 'is not', null)])) .as('active'), eb.fn .countAll() .filterWhere((eb) => eb.and([ eb('task_identifier', '=', name), eb('locked_by', 'is', null), eb('run_at', '<=', eb.fn('now')), ]), ) .as('waiting'), eb.fn .countAll() .filterWhere((eb) => eb.and([ eb('task_identifier', '=', name), eb('locked_by', 'is', null), eb('run_at', '>', eb.fn('now')), ]), ) .as('delayed'), eb.fn .countAll() .filterWhere((eb) => eb.and([eb('task_identifier', '=', name), eb('attempts', '>=', eb.ref('max_attempts'))])) .as('failed'), ]) .executeTakeFirstOrThrow(); } async getQueueStatus(queueName: QueueName): Promise { const state = await this.systemMetadataRepository.get(SystemMetadataKey.QUEUES_STATE); return { paused: state?.[queueName]?.paused ?? false }; } private getJobOptions(item: JobItem): TaskSpec | undefined { switch (item.name) { case JobName.NOTIFY_ALBUM_UPDATE: { let runAt: Date | undefined; if (item.data?.delay) { runAt = DateTime.now().plus(Duration.fromMillis(item.data.delay)).toJSDate(); } return { jobKey: item.data.id, runAt }; } case JobName.STORAGE_TEMPLATE_MIGRATION_SINGLE: { return { jobKey: QueueName.STORAGE_TEMPLATE_MIGRATION }; } case JobName.GENERATE_PERSON_THUMBNAIL: { return { priority: 1 }; } case JobName.QUEUE_FACIAL_RECOGNITION: { return { jobKey: JobName.QUEUE_FACIAL_RECOGNITION }; } } } }