quak/test/library/precache.test.ts

/**
 * The aggressive local precache (issue #48), driven from `Library.open`.
 *
 * Two background fills start with no caller input: every thumbnail in the
 * account newest first, and the originals of the pinned set (the favorites
 * album then the latest `precacheOriginalsDays` window ending at the newest
 * file). Both run through the shared pools at background priority, so on-demand
 * work always preempts them; both report through `status()`. The pinned set is
 * the eviction predicate (#47), so a pinned original is never evicted and a
 * file that leaves the set becomes an ordinary, evictable original.
 *
 * The unit tests drive `Precache` against a fake cache that records what it was
 * asked to fetch (order and kind) with no pool or network; the integration
 * tests drive the real wiring through `Library.open` with a stub content
 * source, and the eviction test drives the real `ContentCache`.
 */

import { describe, it, expect, beforeEach, afterEach } from "vitest";
import { mkdtempSync, rmSync, existsSync, utimesSync } from "node:fs";
import { writeFile } from "node:fs/promises";
import { tmpdir } from "node:os";
import { join } from "node:path";

import { Precache, type PrecacheCache } from "../../src/library/precache.js";
import {
    deriveRecords,
    type DerivedRecords,
} from "../../src/library/records.js";
import {
    ContentCache,
    type ContentSource,
    type EnsureResult,
    type StatFsFn,
} from "../../src/library/content.js";
import { RequestPools } from "../../src/library/pools.js";
import { Library } from "../../src/library/index.js";
import type { Collection, EnteFile } from "../../src/model/types.js";
import type { CollectionsPage, FilesPage } from "../../src/client.js";

const DAY_MICROS = 24 * 60 * 60 * 1000 * 1000;

const collection = (
    id: number,
    type: Collection["type"] = "album",
): Collection => ({
    id,
    ownerID: 1,
    key: new Uint8Array([id & 0xff]),
    name: `album-${id}`,
    type,
    updationTime: 1,
    isShared: false,
});

// A file whose creationTime (microseconds) places it `daysAgo` days before a
// fixed reference instant, so the latest-week window is deterministic.
const REFERENCE_MICROS = 1_000 * DAY_MICROS;
const file = (id: number, collectionID: number, daysAgo: number): EnteFile => ({
    id,
    collectionID,
    ownerID: 1,
    key: new Uint8Array([id & 0xff]),
    metadata: {
        title: `file-${id}.jpg`,
        fileType: "image",
        creationTime: REFERENCE_MICROS - daysAgo * DAY_MICROS,
        modificationTime: 0,
    },
    file: { decryptionHeader: "aGVhZGVy" },
    thumbnail: { decryptionHeader: "dGh1bWI=" },
    updationTime: 1,
});

const records = (
    collections: Collection[],
    files: EnteFile[],
): DerivedRecords => deriveRecords(collections, files);

// A fake cache: records every fetch (kind + order) and reports the files it has
// stored via `pathsFor`. `presentThumbs`/`presentOriginals` seed already-cached
// files so the precache skips them with a single lookup.
class FakeCache implements PrecacheCache {
    readonly thumbFetched: number[] = [];
    readonly originalFetched: number[] = [];
    readonly presentThumbs = new Set<number>();
    readonly presentOriginals = new Set<number>();

    pathsFor(fileID: number): {
        originalPath?: string;
        thumbnailPath?: string;
    } {
        const out: { originalPath?: string; thumbnailPath?: string } = {};
        if (this.presentThumbs.has(fileID))
            out.thumbnailPath = `/thumbs/${fileID}`;
        if (this.presentOriginals.has(fileID))
            out.originalPath = `/originals/${fileID}`;
        return out;
    }

    async ensureThumbnails(args: {
        fileIDs: number[];
        priority: "background";
    }): Promise<EnsureResult[]> {
        return args.fileIDs.map((fileID) => {
            this.thumbFetched.push(fileID);
            this.presentThumbs.add(fileID);
            return { fileID, path: `/thumbs/${fileID}` };
        });
    }

    async ensureOriginals(args: {
        fileIDs: number[];
    }): Promise<EnsureResult[]> {
        return args.fileIDs.map((fileID) => {
            this.originalFetched.push(fileID);
            this.presentOriginals.add(fileID);
            return { fileID, path: `/originals/${fileID}` };
        });
    }
}

// Resolve once a predicate holds, polling the microtask queue; fails fast
// rather than hanging the suite.
const until = async (predicate: () => boolean): Promise<void> => {
    for (let i = 0; i < 1000; i++) {
        if (predicate()) return;
        await new Promise((r) => setTimeout(r, 1));
    }
    throw new Error("condition not met in time");
};

describe("Precache unit", () => {
    it("precaches every thumbnail newest first, skipping present ones", async () => {
        const cols = [collection(1)];
        const files = [
            file(1, 1, 0),
            file(2, 1, 1),
            file(3, 1, 2),
            file(4, 1, 3),
        ];
        const cache = new FakeCache();
        cache.presentThumbs.add(3); // already on disk: skipped
        const pre = new Precache({ originals: false });
        pre.bind(cache);
        pre.update(records(cols, files));
        pre.start();

        await until(() => cache.thumbFetched.length === 3);
        // Newest first (file 1 is newest), file 3 skipped by a lookup.
        expect(cache.thumbFetched).toEqual([1, 2, 4]);
        expect(cache.originalFetched).toEqual([]);
    });

    it("pins favorites then the latest-week window and precaches their originals in that order", async () => {
        const cols = [collection(1), collection(2, "favorites")];
        // File 10 is an old favorite (30 days old); files 1..3 are within the
        // 7-day window; file 4 is outside it.
        const files = [
            file(1, 1, 0),
            file(2, 1, 2),
            file(3, 1, 6),
            file(4, 1, 20),
            file(10, 2, 30), // favorite, old
        ];
        const cache = new FakeCache();
        const pre = new Precache({ originalsDays: 7 });
        pre.bind(cache);
        pre.update(records(cols, files));
        pre.start();

        await until(() => cache.originalFetched.length === 4);
        // Favorite (10) first, then the window newest-first (1, 2, 3). File 4
        // is outside the window and never pinned.
        expect(cache.originalFetched).toEqual([10, 1, 2, 3]);
        expect(pre.isPinned(10)).toBe(true);
        expect(pre.isPinned(1)).toBe(true);
        expect(pre.isPinned(4)).toBe(false);
    });

    it("drops a file from the pinned set when the window moves past it", () => {
        const cols = [collection(1)];
        const cache = new FakeCache();
        const pre = new Precache({ originalsDays: 7 });
        pre.bind(cache);

        pre.update(records(cols, [file(1, 1, 0), file(2, 1, 3)]));
        expect(pre.isPinned(2)).toBe(true);

        // A newer file arrives; the window's newest end moves forward so the
        // 3-day-old file 2 (now 13 days behind the newest) falls out.
        pre.update(
            records(cols, [file(3, 1, -10), file(1, 1, 0), file(2, 1, 3)]),
        );
        expect(pre.isPinned(3)).toBe(true);
        expect(pre.isPinned(2)).toBe(false);
    });

    it("reports progress through status()", async () => {
        const cols = [collection(1), collection(2, "favorites")];
        const files = [file(1, 1, 0), file(2, 1, 1), file(10, 2, 0)];
        const cache = new FakeCache();
        const pre = new Precache({ originalsDays: 7 });
        pre.bind(cache);
        pre.update(records(cols, files));

        const before = pre.status();
        expect(before.thumbnailsTotal).toBe(3);
        expect(before.thumbnailsCached).toBe(0);
        expect(before.originalsPinned).toBe(3); // files 1, 2, 10 all in window
        expect(before.originalsCached).toBe(0);

        pre.start();
        await until(
            () =>
                cache.thumbFetched.length === 3 &&
                cache.originalFetched.length === 3,
        );
        const after = pre.status();
        expect(after.thumbnailsCached).toBe(3);
        expect(after.originalsCached).toBe(3);
    });

    it("honours the disable flags", async () => {
        const cols = [collection(1)];
        const files = [file(1, 1, 0)];
        const cache = new FakeCache();
        const pre = new Precache({ thumbnails: false, originals: false });
        pre.bind(cache);
        pre.update(records(cols, files));
        pre.start();

        await new Promise((r) => setTimeout(r, 20));
        expect(cache.thumbFetched).toEqual([]);
        expect(cache.originalFetched).toEqual([]);
        expect(pre.isPinned(1)).toBe(false);
        expect(pre.status().thumbnailsTotal).toBe(0);
    });
});

// ---- Integration through the real ContentCache and Library ----

const enteFile = (id: number, collectionID: number): EnteFile =>
    file(id, collectionID, 0);

describe("Precache eviction integration", () => {
    let root: string;

    beforeEach(() => {
        root = mkdtempSync(join(tmpdir(), "quak-precache-evict-"));
    });
    afterEach(() => {
        if (root && existsSync(root))
            rmSync(root, { recursive: true, force: true });
    });

    it("never evicts a pinned original the precache put in place", async () => {
        const cacheDir = join(root, "cache");
        // File 1 is the favorites album's only file (pinned regardless of age);
        // files 2 and 3 sit outside the latest-week window, so only file 1 is
        // pinned. The by-id map serves the bytes for each fetch.
        const cols = [collection(1), collection(2, "favorites")];
        const files = [file(1, 2, 30), file(2, 1, 40), file(3, 1, 50)];
        const byID = new Map<number, EnteFile>(files.map((f) => [f.id, f]));
        const source: ContentSource = {
            original: async ({ destination }) => {
                await writeFile(destination, Buffer.alloc(10, 1));
                return { bytesWritten: 10 };
            },
            thumbnail: async ({ destination }) => {
                await writeFile(destination, Buffer.alloc(10, 1));
                return { bytesWritten: 10 };
            },
        };
        const statfs: StatFsFn = async () => ({
            bsize: 1,
            bavail: 1_000_000_000,
        });
        const pre = new Precache({ originalsDays: 7 });
        const cache = new ContentCache({
            pools: new RequestPools(),
            source,
            cacheDirectory: cacheDir,
            getFile: (id) => byID.get(id),
            statfs,
            cacheOriginalsMaxBytes: 25, // holds two 10-byte originals
            freeBelowBytes: 0,
            isPinned: (id) => pre.isPinned(id),
        });
        pre.bind(cache);
        pre.update(records(cols, files));
        expect(pre.isPinned(1)).toBe(true);
        expect(pre.isPinned(2)).toBe(false);
        await cache.open();

        // Fill three originals; the 25-byte cap forces an eviction on the
        // third, and the pinned file 1 must survive it even though it is the
        // least-recently-used.
        await cache.original(1);
        utimesSync(join(cacheDir, "originals", "1.jpg"), 1000, 1000); // oldest
        await cache.original(2);
        utimesSync(join(cacheDir, "originals", "2.jpg"), 2000, 2000);
        await cache.original(3);

        expect(existsSync(join(cacheDir, "originals", "1.jpg"))).toBe(true);
        expect(existsSync(join(cacheDir, "originals", "2.jpg"))).toBe(false);
        expect(existsSync(join(cacheDir, "originals", "3.jpg"))).toBe(true);
    });
});

describe("Precache preemption", () => {
    let root: string;

    beforeEach(() => {
        root = mkdtempSync(join(tmpdir(), "quak-precache-preempt-"));
    });
    afterEach(() => {
        if (root && existsSync(root))
            rmSync(root, { recursive: true, force: true });
    });

    it("lets an on-demand original preempt the background originals fill", async () => {
        const byID = new Map<number, EnteFile>([
            [1, enteFile(1, 1)],
            [2, enteFile(2, 1)],
            [3, enteFile(3, 1)],
        ]);
        const finished: number[] = [];
        let openGate!: () => void;
        const gate = new Promise<void>((r) => (openGate = r));
        let sawFirst!: () => void;
        const firstStarted = new Promise<void>((r) => (sawFirst = r));
        let started = 0;
        const source: ContentSource = {
            original: async ({ file: f, destination }) => {
                if (++started === 1) sawFirst();
                await gate;
                await writeFile(destination, Buffer.alloc(10, 1));
                finished.push(f.id);
                return { bytesWritten: 10 };
            },
            thumbnail: async ({ destination }) => {
                await writeFile(destination, Buffer.alloc(10, 1));
                return { bytesWritten: 10 };
            },
        };
        // One content slot, so file 1 holds it while 2 and 3 wait.
        const cache = new ContentCache({
            pools: new RequestPools({ contentConcurrency: 1 }),
            source,
            cacheDirectory: join(root, "cache"),
            getFile: (id) => byID.get(id),
            statfs: async () => ({ bsize: 1, bavail: 1_000_000_000 }),
            freeBelowBytes: 0,
        });
        await cache.open();

        const pA = cache.ensureOriginals({ fileIDs: [1] }); // background
        await firstStarted; // file 1 now holds the only slot
        const pB = cache.original(2); // on-demand, queued behind file 1
        const pC = cache.ensureOriginals({ fileIDs: [3] }); // background, queued
        await new Promise((r) => setTimeout(r, 5)); // let both enqueue
        openGate();
        await Promise.all([pA, pB, pC]);

        // On-demand file 2 was served before the background file 3.
        expect(finished).toEqual([1, 2, 3]);
    });
});

describe("Precache through Library.open", () => {
    let root: string;

    beforeEach(() => {
        root = mkdtempSync(join(tmpdir(), "quak-precache-lib-"));
    });
    afterEach(() => {
        if (root && existsSync(root))
            rmSync(root, { recursive: true, force: true });
    });

    class MockClient {
        served = false;
        whoami(): { email: string; userID: number } {
            return { email: "u@example.com", userID: 7 };
        }
        async collectionsSince(): Promise<CollectionsPage> {
            if (this.served) return { collections: [], deleted: [], cursor: 1 };
            this.served = true;
            return {
                collections: [collection(1), collection(2, "favorites")],
                deleted: [],
                cursor: 1,
            };
        }
        async filesSince(args: { collectionID: number }): Promise<FilesPage> {
            const files =
                args.collectionID === 1
                    ? [enteFile(1, 1), enteFile(2, 1)]
                    : [enteFile(3, 2)];
            return { files, deleted: [], cursor: 1 };
        }
    }

    it("starts both precaches from open() and reports them in status()", async () => {
        const thumbFetched = new Set<number>();
        const origFetched = new Set<number>();
        const source: ContentSource = {
            original: async ({ file: f, destination }) => {
                origFetched.add(f.id);
                await writeFile(destination, Buffer.alloc(10, 1));
                return { bytesWritten: 10 };
            },
            thumbnail: async ({ file: f, destination }) => {
                thumbFetched.add(f.id);
                await writeFile(destination, Buffer.alloc(10, 1));
                return { bytesWritten: 10 };
            },
        };
        const lib = await Library.open({
            client: new MockClient(),
            cacheDirectory: join(root, "cache"),
            contentSource: source,
            refreshIntervalSeconds: 3600,
        });

        // Every file's thumbnail is precached; the favorite (file 3) and the
        // week's files (1, 2) all have their originals precached.
        await until(() => thumbFetched.size === 3 && origFetched.size === 3);
        const status = lib.status();
        expect(status.thumbnailsTotal).toBe(3);
        expect(status.thumbnailsCached).toBe(3);
        expect(status.originalsPinned).toBe(3);
        expect(status.originalsCached).toBe(3);
        lib.close();
    });
});