mirror of
https://github.com/immich-app/immich.git
synced 2025-03-25 01:15:54 +01:00
294 lines
9.2 KiB
TypeScript
294 lines
9.2 KiB
TypeScript
import { getQueueToken } from '@nestjs/bullmq';
|
|
import { Inject, Injectable } from '@nestjs/common';
|
|
import { ModuleRef } from '@nestjs/core';
|
|
import { SchedulerRegistry } from '@nestjs/schedule';
|
|
import { Job, JobsOptions, Processor, Queue, Worker, WorkerOptions } from 'bullmq';
|
|
import { CronJob, CronTime } from 'cron';
|
|
import { setTimeout } from 'node:timers/promises';
|
|
import { IConfigRepository } from 'src/interfaces/config.interface';
|
|
import {
|
|
IEntityJob,
|
|
IJobRepository,
|
|
JobCounts,
|
|
JobItem,
|
|
JobName,
|
|
QueueCleanType,
|
|
QueueName,
|
|
QueueStatus,
|
|
} from 'src/interfaces/job.interface';
|
|
import { ILoggerRepository } from 'src/interfaces/logger.interface';
|
|
|
|
export const JOBS_TO_QUEUE: Record<JobName, QueueName> = {
|
|
// misc
|
|
[JobName.ASSET_DELETION]: QueueName.BACKGROUND_TASK,
|
|
[JobName.ASSET_DELETION_CHECK]: QueueName.BACKGROUND_TASK,
|
|
[JobName.USER_DELETE_CHECK]: QueueName.BACKGROUND_TASK,
|
|
[JobName.USER_DELETION]: QueueName.BACKGROUND_TASK,
|
|
[JobName.DELETE_FILES]: QueueName.BACKGROUND_TASK,
|
|
[JobName.CLEAN_OLD_AUDIT_LOGS]: QueueName.BACKGROUND_TASK,
|
|
[JobName.CLEAN_OLD_SESSION_TOKENS]: QueueName.BACKGROUND_TASK,
|
|
[JobName.PERSON_CLEANUP]: QueueName.BACKGROUND_TASK,
|
|
[JobName.USER_SYNC_USAGE]: QueueName.BACKGROUND_TASK,
|
|
|
|
// backups
|
|
[JobName.BACKUP_DATABASE]: QueueName.BACKUP_DATABASE,
|
|
|
|
// conversion
|
|
[JobName.QUEUE_VIDEO_CONVERSION]: QueueName.VIDEO_CONVERSION,
|
|
[JobName.VIDEO_CONVERSION]: QueueName.VIDEO_CONVERSION,
|
|
|
|
// thumbnails
|
|
[JobName.QUEUE_GENERATE_THUMBNAILS]: QueueName.THUMBNAIL_GENERATION,
|
|
[JobName.GENERATE_THUMBNAILS]: QueueName.THUMBNAIL_GENERATION,
|
|
[JobName.GENERATE_PERSON_THUMBNAIL]: QueueName.THUMBNAIL_GENERATION,
|
|
|
|
// tags
|
|
[JobName.TAG_CLEANUP]: QueueName.BACKGROUND_TASK,
|
|
|
|
// metadata
|
|
[JobName.QUEUE_METADATA_EXTRACTION]: QueueName.METADATA_EXTRACTION,
|
|
[JobName.METADATA_EXTRACTION]: QueueName.METADATA_EXTRACTION,
|
|
[JobName.LINK_LIVE_PHOTOS]: QueueName.METADATA_EXTRACTION,
|
|
|
|
// storage template
|
|
[JobName.STORAGE_TEMPLATE_MIGRATION]: QueueName.STORAGE_TEMPLATE_MIGRATION,
|
|
[JobName.STORAGE_TEMPLATE_MIGRATION_SINGLE]: QueueName.STORAGE_TEMPLATE_MIGRATION,
|
|
|
|
// migration
|
|
[JobName.QUEUE_MIGRATION]: QueueName.MIGRATION,
|
|
[JobName.MIGRATE_ASSET]: QueueName.MIGRATION,
|
|
[JobName.MIGRATE_PERSON]: QueueName.MIGRATION,
|
|
|
|
// facial recognition
|
|
[JobName.QUEUE_FACE_DETECTION]: QueueName.FACE_DETECTION,
|
|
[JobName.FACE_DETECTION]: QueueName.FACE_DETECTION,
|
|
[JobName.QUEUE_FACIAL_RECOGNITION]: QueueName.FACIAL_RECOGNITION,
|
|
[JobName.FACIAL_RECOGNITION]: QueueName.FACIAL_RECOGNITION,
|
|
|
|
// smart search
|
|
[JobName.QUEUE_SMART_SEARCH]: QueueName.SMART_SEARCH,
|
|
[JobName.SMART_SEARCH]: QueueName.SMART_SEARCH,
|
|
|
|
// duplicate detection
|
|
[JobName.QUEUE_DUPLICATE_DETECTION]: QueueName.DUPLICATE_DETECTION,
|
|
[JobName.DUPLICATE_DETECTION]: QueueName.DUPLICATE_DETECTION,
|
|
|
|
// XMP sidecars
|
|
[JobName.QUEUE_SIDECAR]: QueueName.SIDECAR,
|
|
[JobName.SIDECAR_DISCOVERY]: QueueName.SIDECAR,
|
|
[JobName.SIDECAR_SYNC]: QueueName.SIDECAR,
|
|
[JobName.SIDECAR_WRITE]: QueueName.SIDECAR,
|
|
|
|
// Library management
|
|
[JobName.LIBRARY_SYNC_FILE]: QueueName.LIBRARY,
|
|
[JobName.LIBRARY_QUEUE_SYNC_FILES]: QueueName.LIBRARY,
|
|
[JobName.LIBRARY_QUEUE_SYNC_ASSETS]: QueueName.LIBRARY,
|
|
[JobName.LIBRARY_DELETE]: QueueName.LIBRARY,
|
|
[JobName.LIBRARY_SYNC_ASSET]: QueueName.LIBRARY,
|
|
[JobName.LIBRARY_QUEUE_SYNC_ALL]: QueueName.LIBRARY,
|
|
[JobName.LIBRARY_QUEUE_CLEANUP]: QueueName.LIBRARY,
|
|
|
|
// Notification
|
|
[JobName.SEND_EMAIL]: QueueName.NOTIFICATION,
|
|
[JobName.NOTIFY_ALBUM_INVITE]: QueueName.NOTIFICATION,
|
|
[JobName.NOTIFY_ALBUM_UPDATE]: QueueName.NOTIFICATION,
|
|
[JobName.NOTIFY_SIGNUP]: QueueName.NOTIFICATION,
|
|
|
|
// Version check
|
|
[JobName.VERSION_CHECK]: QueueName.BACKGROUND_TASK,
|
|
|
|
// Trash
|
|
[JobName.QUEUE_TRASH_EMPTY]: QueueName.BACKGROUND_TASK,
|
|
};
|
|
|
|
@Injectable()
|
|
export class JobRepository implements IJobRepository {
|
|
private workers: Partial<Record<QueueName, Worker>> = {};
|
|
|
|
constructor(
|
|
private moduleReference: ModuleRef,
|
|
private schedulerReqistry: SchedulerRegistry,
|
|
@Inject(IConfigRepository) private configRepository: IConfigRepository,
|
|
@Inject(ILoggerRepository) private logger: ILoggerRepository,
|
|
) {
|
|
this.logger.setContext(JobRepository.name);
|
|
}
|
|
|
|
addHandler(queueName: QueueName, concurrency: number, handler: (item: JobItem) => Promise<void>) {
|
|
const { bull } = this.configRepository.getEnv();
|
|
const workerHandler: Processor = async (job: Job) => handler(job as JobItem);
|
|
const workerOptions: WorkerOptions = { ...bull.config, concurrency };
|
|
this.workers[queueName] = new Worker(queueName, workerHandler, workerOptions);
|
|
}
|
|
|
|
addCronJob(name: string, expression: string, onTick: () => void, start = true): void {
|
|
const job = new CronJob<null, null>(
|
|
expression,
|
|
onTick,
|
|
// function to run onComplete
|
|
undefined,
|
|
// whether it should start directly
|
|
start,
|
|
// timezone
|
|
undefined,
|
|
// context
|
|
undefined,
|
|
// runOnInit
|
|
undefined,
|
|
// utcOffset
|
|
undefined,
|
|
// prevents memory leaking by automatically stopping when the node process finishes
|
|
true,
|
|
);
|
|
|
|
this.schedulerReqistry.addCronJob(name, job);
|
|
}
|
|
|
|
updateCronJob(name: string, expression?: string, start?: boolean): void {
|
|
const job = this.schedulerReqistry.getCronJob(name);
|
|
if (expression) {
|
|
job.setTime(new CronTime(expression));
|
|
}
|
|
if (start !== undefined) {
|
|
if (start) {
|
|
job.start();
|
|
} else {
|
|
job.stop();
|
|
}
|
|
}
|
|
}
|
|
|
|
setConcurrency(queueName: QueueName, concurrency: number) {
|
|
const worker = this.workers[queueName];
|
|
if (!worker) {
|
|
this.logger.warn(`Unable to set queue concurrency, worker not found: '${queueName}'`);
|
|
return;
|
|
}
|
|
|
|
worker.concurrency = concurrency;
|
|
}
|
|
|
|
async getQueueStatus(name: QueueName): Promise<QueueStatus> {
|
|
const queue = this.getQueue(name);
|
|
|
|
return {
|
|
isActive: !!(await queue.getActiveCount()),
|
|
isPaused: await queue.isPaused(),
|
|
};
|
|
}
|
|
|
|
pause(name: QueueName) {
|
|
return this.getQueue(name).pause();
|
|
}
|
|
|
|
resume(name: QueueName) {
|
|
return this.getQueue(name).resume();
|
|
}
|
|
|
|
empty(name: QueueName) {
|
|
return this.getQueue(name).drain();
|
|
}
|
|
|
|
clear(name: QueueName, type: QueueCleanType) {
|
|
return this.getQueue(name).clean(0, 1000, type);
|
|
}
|
|
|
|
getJobCounts(name: QueueName): Promise<JobCounts> {
|
|
return this.getQueue(name).getJobCounts(
|
|
'active',
|
|
'completed',
|
|
'failed',
|
|
'delayed',
|
|
'waiting',
|
|
'paused',
|
|
) as unknown as Promise<JobCounts>;
|
|
}
|
|
|
|
async queueAll(items: JobItem[]): Promise<void> {
|
|
if (items.length === 0) {
|
|
return;
|
|
}
|
|
|
|
const promises = [];
|
|
const itemsByQueue = {} as Record<string, (JobItem & { data: any; options: JobsOptions | undefined })[]>;
|
|
for (const item of items) {
|
|
const queueName = JOBS_TO_QUEUE[item.name];
|
|
const job = {
|
|
name: item.name,
|
|
data: item.data || {},
|
|
options: this.getJobOptions(item) || undefined,
|
|
} as JobItem & { data: any; options: JobsOptions | undefined };
|
|
|
|
if (job.options?.jobId) {
|
|
// need to use add() instead of addBulk() for jobId deduplication
|
|
promises.push(this.getQueue(queueName).add(item.name, item.data, job.options));
|
|
} else {
|
|
itemsByQueue[queueName] = itemsByQueue[queueName] || [];
|
|
itemsByQueue[queueName].push(job);
|
|
}
|
|
}
|
|
|
|
for (const [queueName, jobs] of Object.entries(itemsByQueue)) {
|
|
const queue = this.getQueue(queueName as QueueName);
|
|
promises.push(queue.addBulk(jobs));
|
|
}
|
|
|
|
await Promise.all(promises);
|
|
}
|
|
|
|
async queue(item: JobItem): Promise<void> {
|
|
return this.queueAll([item]);
|
|
}
|
|
|
|
async waitForQueueCompletion(...queues: QueueName[]): Promise<void> {
|
|
let activeQueue: QueueStatus | undefined;
|
|
do {
|
|
const statuses = await Promise.all(queues.map((name) => this.getQueueStatus(name)));
|
|
activeQueue = statuses.find((status) => status.isActive);
|
|
} while (activeQueue);
|
|
{
|
|
this.logger.verbose(`Waiting for ${activeQueue} queue to stop...`);
|
|
await setTimeout(1000);
|
|
}
|
|
}
|
|
|
|
private getJobOptions(item: JobItem): JobsOptions | null {
|
|
switch (item.name) {
|
|
case JobName.NOTIFY_ALBUM_UPDATE: {
|
|
return { jobId: item.data.id, delay: item.data?.delay };
|
|
}
|
|
case JobName.STORAGE_TEMPLATE_MIGRATION_SINGLE: {
|
|
return { jobId: item.data.id };
|
|
}
|
|
case JobName.GENERATE_PERSON_THUMBNAIL: {
|
|
return { priority: 1 };
|
|
}
|
|
case JobName.QUEUE_FACIAL_RECOGNITION: {
|
|
return { jobId: JobName.QUEUE_FACIAL_RECOGNITION };
|
|
}
|
|
default: {
|
|
return null;
|
|
}
|
|
}
|
|
}
|
|
|
|
private getQueue(queue: QueueName): Queue {
|
|
return this.moduleReference.get<Queue>(getQueueToken(queue), { strict: false });
|
|
}
|
|
|
|
public async removeJob(jobId: string, name: JobName): Promise<IEntityJob | undefined> {
|
|
const existingJob = await this.getQueue(JOBS_TO_QUEUE[name]).getJob(jobId);
|
|
if (!existingJob) {
|
|
return;
|
|
}
|
|
try {
|
|
await existingJob.remove();
|
|
} catch (error: any) {
|
|
if (error.message?.includes('Missing key for job')) {
|
|
return;
|
|
}
|
|
throw error;
|
|
}
|
|
return existingJob.data;
|
|
}
|
|
}
|