Abstract away rollback management (#2889)

Author: CarsonF

src/components/file/file.service.ts

@@ -4,7 +4,6 @@ import {
 } from '@aws-sdk/client-s3';
 import { Injectable } from '@nestjs/common';
 import { bufferFromStream, cleanJoin, Nil } from '@seedcompany/common';
-import { Connection } from 'cypher-query-builder';
 import { fileTypeFromStream } from 'file-type';
 import { intersection } from 'lodash';
 import { Duration } from 'luxon';
@@ -23,7 +22,13 @@ import {
   UnauthorizedException,
 } from '~/common';
 import { withAddedPath } from '~/common/url.util';
-import { ConfigService, IEventBus, ILogger, Logger } from '~/core';
+import {
+  ConfigService,
+  IEventBus,
+  ILogger,
+  Logger,
+  RollbackManager,
+} from '~/core';
 import { FileBucket } from './bucket';
 import {
   CreateDefinedFileVersionInput,
@@ -56,7 +61,7 @@ export class FileService {
   constructor(
     private readonly bucket: FileBucket,
     private readonly repo: FileRepository,
-    private readonly db: Connection,
+    private readonly rollbacks: RollbackManager,
     private readonly config: ConfigService,
     private readonly mediaService: MediaService,
     private readonly eventBus: IEventBus,
@@ -384,28 +389,17 @@ export class FileService {
     if (existingUpload.status === 'rejected') {
       await this.bucket.moveObject(`temp/${uploadId}`, uploadId);
 
-      // A bit of a hacky way to move files back to the temp/ folder on
-      // mutation error / transaction rollback. This prevents orphaned files in bucket.
-      const tx = this.db.currentTransaction;
-      // The mutation can be retried multiple times, when neo4j deems the error
-      // is retry-able, but we only want to attach this rollback logic once.
-      if (tx && !(tx as any).__S3_ROLLBACK) {
-        (tx as any).__S3_ROLLBACK = true;
-
-        const orig = tx.rollback.bind(tx);
-        tx.rollback = async () => {
-          // Undo above operation by moving it back to temp folder.
-          await this.bucket
-            .moveObject(uploadId, `temp/${uploadId}`)
-            .catch((e) => {
-              this.logger.error('Failed to move file back to temp holding', {
-                uploadId,
-                exception: e,
-              });
+      // Undo the above operation by moving it back to temp folder.
+      this.rollbacks.add(async () => {
+        await this.bucket
+          .moveObject(uploadId, `temp/${uploadId}`)
+          .catch((e) => {
+            this.logger.error('Failed to move file back to temp holding', {
+              uploadId,
+              exception: e,
             });
-          await orig();
-        };
-      }
+          });
+      });
     }
 
     await this.mediaService.detectAndSave(fv, media);

src/core/database/database.module.ts

@@ -9,6 +9,7 @@ import { DatabaseService } from './database.service';
 import { IndexerModule } from './indexer/indexer.module';
 import { MigrationModule } from './migration/migration.module';
 import { ParameterTransformer } from './parameter-transformer.service';
+import { RollbackManager } from './rollback-manager';
 import { TransactionalMutationsInterceptor } from './transactional-mutations.interceptor';
 
 @Module({
@@ -18,8 +19,9 @@ import { TransactionalMutationsInterceptor } from './transactional-mutations.int
     DatabaseService,
     ParameterTransformer,
     { provide: APP_INTERCEPTOR, useClass: TransactionalMutationsInterceptor },
+    RollbackManager,
   ],
-  exports: [CypherFactory, DatabaseService, IndexerModule],
+  exports: [CypherFactory, DatabaseService, IndexerModule, RollbackManager],
 })
 export class DatabaseModule implements OnApplicationShutdown {
   constructor(

src/core/database/index.ts

@@ -7,3 +7,4 @@ export { DtoRepository } from './dto.repository';
 export * from './indexer';
 export * from './migration';
 export * from './db-type';
+export * from './rollback-manager';

src/core/database/rollback-manager.ts

@@ -0,0 +1,47 @@
+import { Injectable } from '@nestjs/common';
+import { GqlContextHost } from '../graphql';
+
+type RollbackFn = () => Promise<void>;
+
+interface RollbackRef {
+  fn: RollbackFn;
+  remove: () => void;
+}
+
+@Injectable()
+export class RollbackManager {
+  private readonly functionsByContext = new WeakMap<
+    GqlContextHost['context'],
+    Set<RollbackFn>
+  >();
+
+  constructor(private readonly contextHost: GqlContextHost) {}
+
+  /**
+   * Add a function to call whenever a rollback is triggered.
+   */
+  add(fn: RollbackFn): RollbackRef {
+    const remove = () => this.functions.delete(fn);
+    this.functions.add(fn);
+    return { fn, remove };
+  }
+
+  /**
+   * Call all the stored rollback functions, and clear the list.
+   */
+  async runAndClear() {
+    const functions = this.functions;
+    for (const fn of functions) {
+      await fn();
+      functions.delete(fn);
+    }
+  }
+
+  private get functions() {
+    const contextId = this.contextHost.context;
+    if (!this.functionsByContext.has(contextId)) {
+      this.functionsByContext.set(contextId, new Set());
+    }
+    return this.functionsByContext.get(contextId)!;
+  }
+}

src/core/database/transactional-mutations.interceptor.ts

@@ -8,14 +8,18 @@ import { GqlContextType, GqlExecutionContext } from '@nestjs/graphql';
 import { Connection } from 'cypher-query-builder';
 import { GraphQLResolveInfo } from 'graphql';
 import { from, lastValueFrom } from 'rxjs';
+import { RollbackManager } from './rollback-manager';
 
 /**
  * Run all mutations in a neo4j transaction.
  * This allows automatic rollbacks on error.
  */
 @Injectable()
 export class TransactionalMutationsInterceptor implements NestInterceptor {
-  constructor(private readonly db: Connection) {}
+  constructor(
+    private readonly db: Connection,
+    private readonly rollbacks: RollbackManager,
+  ) {}
 
   async intercept(context: ExecutionContext, next: CallHandler) {
     if (context.getType<GqlContextType>() !== 'graphql') {
@@ -29,7 +33,14 @@ export class TransactionalMutationsInterceptor implements NestInterceptor {
     }
 
     return from(
-      this.db.runInTransaction(async () => await lastValueFrom(next.handle())),
+      this.db.runInTransaction(async () => {
+        try {
+          return await lastValueFrom(next.handle());
+        } catch (e) {
+          await this.rollbacks.runAndClear();
+          throw e;
+        }
+      }),
     );
   }
 }