From fa7697cabae1310e991c17e9a2df9ba2da710e22 Mon Sep 17 00:00:00 2001
From: automated-signal <37887102+automated-signal@users.noreply.github.com>
Date: Fri, 2 Aug 2024 18:05:13 -0500
Subject: [PATCH] Run orphaned attachment cleanup after timeout

Co-authored-by: Fedor Indutny <79877362+indutny-signal@users.noreply.github.com>
---
 ts/util/encryptLegacyAttachment.ts | 31 +++++++++++++++++++++++++++++-
 1 file changed, 30 insertions(+), 1 deletion(-)

diff --git a/ts/util/encryptLegacyAttachment.ts b/ts/util/encryptLegacyAttachment.ts
index d4ee0d0a2c..72ee878012 100644
--- a/ts/util/encryptLegacyAttachment.ts
+++ b/ts/util/encryptLegacyAttachment.ts
@@ -8,9 +8,17 @@ import type {
   LocalAttachmentV2Type,
 } from '../types/Attachment';
 import * as log from '../logging/log';
+import { DataWriter } from '../sql/Client';
 import { AttachmentDisposition } from './getLocalAttachmentUrl';
+import { drop } from './drop';
+import { MINUTE } from './durations';
 
 let setCheck = false;
+let orphanedCount = 0;
+let cleanupTimeout: NodeJS.Timeout | undefined;
+
+// Max number of orphaned attachments before we schedule a cleanup.
+const MAX_ORPHANED_COUNT = 10000;
 
 const lru = new LRU<string, Promise<LocalAttachmentV2Type>>({
   max: 1000,
@@ -70,11 +78,32 @@ async function doEncrypt<T extends Partial<AddressableAttachmentType>>(
   const data = await readAttachmentData(attachment);
   const result = await writeNewAttachmentData(data);
 
+  orphanedCount += 1;
+
   // Remove fully migrated attachments without references on next startup.
-  if (!setCheck) {
+  if (orphanedCount > MAX_ORPHANED_COUNT) {
+    log.error('encryptLegacyAttachment: too many orphaned, cleanup now');
+    if (cleanupTimeout !== undefined) {
+      clearTimeout(cleanupTimeout);
+      cleanupTimeout = undefined;
+    }
+    cleanup();
+  } else if (!setCheck) {
     setCheck = true;
     await window.storage.put('needOrphanedAttachmentCheck', true);
+    log.error('encryptLegacyAttachment: scheduling orphaned cleanup');
+    cleanupTimeout = setTimeout(cleanup, 15 * MINUTE);
   }
 
   return result;
 }
+
+function cleanup(): void {
+  log.error('encryptLegacyAttachment: running orphaned cleanup');
+
+  cleanupTimeout = undefined;
+  setCheck = false;
+  orphanedCount = 0;
+  drop(window.storage.remove('needOrphanedAttachmentCheck'));
+  drop(DataWriter.cleanupOrphanedAttachments());
+}