refactor: use a heap-based second run approach

gagik · gagik · commit 5bc4cb762a43 · 2025-02-10T14:46:11.000+01:00
Includes test for random OS-based file order as well as mixing different settings together.
diff --git a/packages/mongodb-log-writer/src/mongo-log-manager.spec.ts b/packages/mongodb-log-writer/src/mongo-log-manager.spec.ts
@@ -1,7 +1,7 @@
 import { MongoLogManager, mongoLogId } from '.';
 import { ObjectId } from 'bson';
 import { once } from 'events';
-import type { Stats } from 'fs';
+import type { Stats, Dir } from 'fs';
 import { promises as fs } from 'fs';
 import path from 'path';
 import os from 'os';
@@ -178,7 +178,7 @@ describe('MongoLogManager', function () {
     expect(leftoverFiles).deep.equals([faultyFile, ...validFiles.slice(3)]);
   });
 
-  it('cleans up least recent log files when requested with a storage limit', async function () {
+  it('cleans up least recent log files when over a storage limit', async function () {
     const manager = new MongoLogManager({
       directory,
       retentionDays,
@@ -207,6 +207,172 @@ describe('MongoLogManager', function () {
     expect(await getFilesState(paths)).to.equal('0000111111');
   });
 
+  describe('with a random file order', function () {
+    let paths: string[] = [];
+    const times = [92, 90, 1, 2, 3, 91];
+
+    beforeEach(async function () {
+      const fileNames: string[] = [];
+      paths = [];
+      const offset = Math.floor(Date.now() / 1000);
+
+      for (const time of times) {
+        const fileName =
+          ObjectId.createFromTime(offset - time).toHexString() + '_log';
+        const fullPath = path.join(directory, fileName);
+        await fs.writeFile(fullPath, '0'.repeat(1024));
+        fileNames.push(fileName);
+        paths.push(fullPath);
+      }
+
+      sinon.replace(fs, 'opendir', async () =>
+        Promise.resolve({
+          [Symbol.asyncIterator]: function* () {
+            for (const fileName of fileNames) {
+              yield {
+                name: fileName,
+                isFile: () => true,
+              };
+            }
+          },
+        } as unknown as Dir)
+      );
+    });
+
+    it('cleans up in the expected order with maxLogFileCount', async function () {
+      const manager = new MongoLogManager({
+        directory,
+        retentionDays,
+        maxLogFileCount: 3,
+        onwarn,
+        onerror,
+      });
+
+      expect(await getFilesState(paths)).to.equal('111111');
+
+      await manager.cleanupOldLogFiles();
+
+      expect(await getFilesState(paths)).to.equal('001110');
+    });
+
+    it('cleans up in the expected order with retentionGB', async function () {
+      const manager = new MongoLogManager({
+        directory,
+        retentionDays,
+        retentionGB: 3 / 1024 / 1024,
+        onwarn,
+        onerror,
+      });
+
+      expect(await getFilesState(paths)).to.equal('111111');
+
+      await manager.cleanupOldLogFiles();
+
+      expect(await getFilesState(paths)).to.equal('001110');
+    });
+  });
+
+  describe('with multiple log retention settings', function () {
+    it('with retention days, file count, and max size maintains all conditions', async function () {
+      const manager = new MongoLogManager({
+        directory,
+        retentionDays: 1,
+        maxLogFileCount: 3,
+        retentionGB: 2 / 1024 / 1024,
+        onwarn,
+        onerror,
+      });
+
+      const paths: string[] = [];
+
+      // Create 4 files which are all older than 1 day and 4 which are from today.
+      for (let i = 0; i < 4; i++) {
+        const today = Math.floor(Date.now() / 1000);
+        const yesterday = today - 25 * 60 * 60;
+        const todayFile = path.join(
+          directory,
+          ObjectId.createFromTime(today - i).toHexString() + '_log'
+        );
+        await fs.writeFile(todayFile, '0'.repeat(1024));
+
+        const yesterdayFile = path.join(
+          directory,
+          ObjectId.createFromTime(yesterday - i).toHexString() + '_log'
+        );
+        await fs.writeFile(yesterdayFile, '0'.repeat(1024));
+
+        paths.unshift(todayFile);
+        paths.unshift(yesterdayFile);
+      }
+
+      expect(await getFilesState(paths)).to.equal('11111111');
+
+      await manager.cleanupOldLogFiles();
+
+      // All yesterdays files, 2 of today's files should be deleted.
+      // (because of file count and file size)
+      expect(await getFilesState(paths)).to.equal('00000101');
+    });
+
+    it('with low GB but high file count maintains both conditions', async function () {
+      const manager = new MongoLogManager({
+        directory,
+        retentionDays,
+        maxLogFileCount: 3,
+        // 2 KB, so 2 files
+        retentionGB: 2 / 1024 / 1024,
+        onwarn,
+        onerror,
+      });
+
+      const paths: string[] = [];
+      const offset = Math.floor(Date.now() / 1000);
+
+      // Create 10 files of 1 KB each.
+      for (let i = 0; i < 10; i++) {
+        const filename = path.join(
+          directory,
+          ObjectId.createFromTime(offset - i).toHexString() + '_log'
+        );
+        await fs.writeFile(filename, '0'.repeat(1024));
+        paths.unshift(filename);
+      }
+
+      expect(await getFilesState(paths)).to.equal('1111111111');
+      await manager.cleanupOldLogFiles();
+      expect(await getFilesState(paths)).to.equal('0000000011');
+    });
+
+    it('with high GB but low file count maintains both conditions', async function () {
+      const manager = new MongoLogManager({
+        directory,
+        retentionDays,
+        maxLogFileCount: 2,
+        // 3 KB, so 3 files
+        retentionGB: 3 / 1024 / 1024,
+        onwarn,
+        onerror,
+      });
+
+      const paths: string[] = [];
+      const offset = Math.floor(Date.now() / 1000);
+
+      // Create 10 files of 1 KB each.
+      for (let i = 0; i < 10; i++) {
+        const filename = path.join(
+          directory,
+          ObjectId.createFromTime(offset - i).toHexString() + '_log'
+        );
+        await fs.writeFile(filename, '0'.repeat(1024));
+        paths.unshift(filename);
+      }
+
+      expect(await getFilesState(paths)).to.equal('1111111111');
+      await manager.cleanupOldLogFiles();
+      expect(await getFilesState(paths)).to.equal('0000000011');
+    });
+  });
+
   it('cleaning up old log files is a no-op by default', async function () {
     const manager = new MongoLogManager({
       directory: path.join('directory', 'nonexistent'),
diff --git a/packages/mongodb-log-writer/src/mongo-log-manager.ts b/packages/mongodb-log-writer/src/mongo-log-manager.ts
@@ -1,10 +1,12 @@
 import path from 'path';
 import { ObjectId } from 'bson';
 import { once } from 'events';
+import type { Dir } from 'fs';
 import { createWriteStream, promises as fs } from 'fs';
 import { createGzip, constants as zlibConstants } from 'zlib';
 import { MongoLogWriter } from './mongo-log-writer';
 import { Writable } from 'stream';
+import Heap from 'heap-js';
 
 /** Options used by MongoLogManager instances. */
 interface MongoLogOptions {
@@ -36,40 +38,23 @@ export class MongoLogManager {
     this._options = options;
   }
 
+  private async deleteFile(path: string) {
+    try {
+      await fs.unlink(path);
+      // eslint-disable-next-line @typescript-eslint/no-explicit-any
+    } catch (err: any) {
+      if (err?.code !== 'ENOENT') {
+        this._options.onerror(err as Error, path);
+      }
+    }
+  }
+
   /** Clean up log files older than `retentionDays`. */
   async cleanupOldLogFiles(maxDurationMs = 5_000): Promise<void> {
     const dir = this._options.directory;
-    const sortedLogFiles: {
-      fullPath: string;
-      id: string;
-      size?: number;
-    }[] = [];
-    let usedStorageSize = this._options.retentionGB ? 0 : -Infinity;
-
+    let dirHandle;
     try {
-      const files = await fs.readdir(dir, { withFileTypes: true });
-      for (const file of files) {
-        const { id } =
-          /^(?<id>[a-f0-9]{24})_log(\.gz)?$/i.exec(file.name)?.groups ?? {};
-
-        if (!file.isFile() || !id) {
-          continue;
-        }
-
-        const fullPath = path.join(dir, file.name);
-        let size: number | undefined;
-        if (this._options.retentionGB) {
-          try {
-            size = (await fs.stat(fullPath)).size;
-            usedStorageSize += size;
-          } catch (err) {
-            this._options.onerror(err as Error, fullPath);
-            continue;
-          }
-        }
-
-        sortedLogFiles.push({ fullPath, id, size });
-      }
+      dirHandle = await fs.opendir(dir);
     } catch {
       return;
     }
@@ -78,58 +63,80 @@ export class MongoLogManager {
     // Delete files older than N days
     const deletionCutoffTimestamp =
       deletionStartTimestamp - this._options.retentionDays * 86400 * 1000;
+    // Store the known set of least recent files in a heap in order to be able to
+    // delete all but the most recent N files.
+    const leastRecentFileHeap = new Heap<{
+      fileTimestamp: number;
+      fullPath: string;
+      fileSize: number | undefined;
+    }>((a, b) => a.fileTimestamp - b.fileTimestamp);
 
-    const storageSizeLimit = this._options.retentionGB
-      ? this._options.retentionGB * 1024 * 1024 * 1024
-      : Infinity;
+    let usedStorageSize = this._options.retentionGB ? 0 : -Infinity;
 
-    for await (const { id, fullPath } of [...sortedLogFiles]) {
+    for await (const dirent of dirHandle) {
       // Cap the overall time spent inside this function. Consider situations like
       // a large number of machines using a shared network-mounted $HOME directory
       // where lots and lots of log files end up and filesystem operations happen
       // with network latency.
       if (Date.now() - deletionStartTimestamp > maxDurationMs) break;
 
+      const { id } =
+        /^(?<id>[a-f0-9]{24})_log(\.gz)?$/i.exec(dirent.name)?.groups ?? {};
+
+      if (!dirent.isFile() || !id) {
+        continue;
+      }
+
       const fileTimestamp = +new ObjectId(id).getTimestamp();
-      let toDelete:
-        | {
-            fullPath: string;
-            /** If the file wasn't deleted right away and there is a
-             *  retention size limit, its size should be accounted */
-            fileSize?: number;
-          }
-        | undefined;
+      const fullPath = path.join(dir, dirent.name);
 
       // If the file is older than expected, delete it. If the file is recent,
       // add it to the list of seen files, and if that list is too large, remove
       // the least recent file we've seen so far.
       if (fileTimestamp < deletionCutoffTimestamp) {
-        toDelete = {
-          fullPath,
-        };
-      } else if (this._options.retentionGB || this._options.maxLogFileCount) {
-        const reachedMaxStorageSize = usedStorageSize > storageSizeLimit;
-        const reachedMaxFileCount =
-          this._options.maxLogFileCount &&
-          sortedLogFiles.length > this._options.maxLogFileCount;
-
-        if (reachedMaxStorageSize || reachedMaxFileCount) {
-          toDelete = sortedLogFiles.shift();
-        }
+        await this.deleteFile(fullPath);
+        continue;
       }
 
-      if (!toDelete) continue;
-      try {
-        await fs.unlink(toDelete.fullPath);
-        if (toDelete.fileSize) {
-          usedStorageSize -= toDelete.fileSize;
-        }
-        // eslint-disable-next-line @typescript-eslint/no-explicit-any
-      } catch (err: any) {
-        if (err?.code !== 'ENOENT') {
+      let fileSize: number | undefined;
+      if (this._options.retentionGB) {
+        try {
+          fileSize = (await fs.stat(fullPath)).size;
+          usedStorageSize += fileSize;
+        } catch (err) {
           this._options.onerror(err as Error, fullPath);
+          continue;
         }
       }
+
+      if (this._options.maxLogFileCount || this._options.retentionGB) {
+        leastRecentFileHeap.push({ fullPath, fileTimestamp, fileSize });
+      }
+
+      if (
+        this._options.maxLogFileCount &&
+        leastRecentFileHeap.size() > this._options.maxLogFileCount
+      ) {
+        const toDelete = leastRecentFileHeap.pop();
+        if (!toDelete) continue;
+        await this.deleteFile(toDelete.fullPath);
+        usedStorageSize -= toDelete.fileSize ?? 0;
+      }
+    }
+
+    if (this._options.retentionGB) {
+      const storageSizeLimit = this._options.retentionGB * 1024 * 1024 * 1024;
+
+      for (const file of leastRecentFileHeap) {
+        if (Date.now() - deletionStartTimestamp > maxDurationMs) break;
+
+        if (usedStorageSize <= storageSizeLimit) break;
+
+        if (!file.fileSize) continue;
+
+        await this.deleteFile(file.fullPath);
+        usedStorageSize -= file.fileSize;
+      }
     }
   }
 
