Fallback to in-memory sqllite db in all cases (#591)

This lets up remove a bunch of logic for the old  disk cache because it only needs to be read and then eventually deleted

Author: mjbvz

src/platform/workspaceChunkSearch/node/workspaceChunkAndEmbeddingCache.ts

@@ -6,16 +6,14 @@ import fs from 'fs';
 import { IDisposable } from 'monaco-editor';
 import sql from 'node:sqlite';
 import path from 'path';
-import { CancelablePromise, ThrottledDelayer, createCancelablePromise, raceTimeout } from '../../../util/vs/base/common/async';
+import { CancelablePromise, createCancelablePromise } from '../../../util/vs/base/common/async';
 import { CancellationToken } from '../../../util/vs/base/common/cancellation';
-import { Disposable } from '../../../util/vs/base/common/lifecycle';
 import { ResourceMap } from '../../../util/vs/base/common/map';
 import { Schemas } from '../../../util/vs/base/common/network';
 import { URI } from '../../../util/vs/base/common/uri';
 import { IRange, Range } from '../../../util/vs/editor/common/core/range';
 import { IInstantiationService, ServicesAccessor } from '../../../util/vs/platform/instantiation/common/instantiation';
-import { FileChunk, FileChunkWithEmbedding } from '../../chunking/common/chunk';
-import { stripChunkTextMetadata } from '../../chunking/common/chunkingStringUtils';
+import { FileChunkWithEmbedding } from '../../chunking/common/chunk';
 import { Embedding, EmbeddingType, EmbeddingVector } from '../../embeddings/common/embeddingsComputer';
 import { IFileSystemService } from '../../filesystem/common/fileSystemService';
 import { ILogService } from '../../log/common/logService';
@@ -78,24 +76,13 @@ export async function createWorkspaceChunkAndEmbeddingCache(
 	workspaceIndex: IWorkspaceFileIndex
 ): Promise<IWorkspaceChunkAndEmbeddingCache> {
 	const instantiationService = accessor.get(IInstantiationService);
-	if (cacheRoot) {
-		const db = await instantiationService.invokeFunction(accessor => DbCache.create(accessor, embeddingType, cacheRoot, workspaceIndex));
-		if (db) {
-			return db;
-		}
-	}
-	return instantiationService.invokeFunction(accessor => DiskCache.load(accessor, embeddingType, cacheRoot, workspaceIndex));
+	return instantiationService.invokeFunction(accessor => DbCache.create(accessor, embeddingType, cacheRoot ?? ':memory:', workspaceIndex));
 }
 
-class DiskCache extends Disposable implements IWorkspaceChunkAndEmbeddingCache {
+class OldDiskCache {
 	private static readonly version = '1.0.0';
 	private static cacheFileName = 'workspace-chunks.json';
 
-	private static encodeEmbedding(embedding: EmbeddingVector): string {
-		const floatArray = Float32Array.from(embedding);
-		return Buffer.from(floatArray.buffer).toString('base64');
-	}
-
 	public static decodeEmbedding(base64Str: string): EmbeddingVector {
 		const decoded = Buffer.from(base64Str, 'base64');
 		const float32Array = new Float32Array(decoded.buffer, decoded.byteOffset, decoded.byteLength / Float32Array.BYTES_PER_ELEMENT);
@@ -105,7 +92,7 @@ class DiskCache extends Disposable implements IWorkspaceChunkAndEmbeddingCache {
 	public static async readDiskCache(accessor: ServicesAccessor, embeddingType: EmbeddingType, cacheRoot: URI, logService: ILogService): Promise<Iterable<[string, PersistedCacheEntry]> | undefined> {
 		const fileSystem = accessor.get(IFileSystemService);
 
-		const cachePath = URI.joinPath(cacheRoot, DiskCache.cacheFileName);
+		const cachePath = URI.joinPath(cacheRoot, OldDiskCache.cacheFileName);
 		try {
 			let file: Uint8Array | undefined;
 			try {
@@ -116,8 +103,8 @@ class DiskCache extends Disposable implements IWorkspaceChunkAndEmbeddingCache {
 			}
 
 			const data: PersistedCache = JSON.parse(new TextDecoder().decode(file));
-			if (data.version !== DiskCache.version) {
-				logService.debug(`WorkspaceChunkAndEmbeddingCache: invalidating cache due to version mismatch. Expected ${DiskCache.version} but found ${data.version}`);
+			if (data.version !== OldDiskCache.version) {
+				logService.debug(`WorkspaceChunkAndEmbeddingCache: invalidating cache due to version mismatch. Expected ${OldDiskCache.version} but found ${data.version}`);
 				return undefined;
 			}
 
@@ -139,215 +126,15 @@ class DiskCache extends Disposable implements IWorkspaceChunkAndEmbeddingCache {
 
 	static async deleteDiskCache(accessor: ServicesAccessor, cacheRoot: URI) {
 		const fileSystem = accessor.get(IFileSystemService);
-		const cachePath = URI.joinPath(cacheRoot, DiskCache.cacheFileName);
+		const cachePath = URI.joinPath(cacheRoot, OldDiskCache.cacheFileName);
 		try {
 			await fileSystem.delete(cachePath);
 		} catch {
 			// noop
 		}
 	}
 
-	static async load(
-		accessor: ServicesAccessor,
-		embeddingType: EmbeddingType,
-		cacheRoot: URI | undefined,
-		workspaceIndex: IWorkspaceFileIndex
-	): Promise<DiskCache> {
-		const fileSystem = accessor.get(IFileSystemService);
-		const instantiationService = accessor.get(IInstantiationService);
-		const logService = accessor.get(ILogService);
-
-		const cachePath = cacheRoot ? URI.joinPath(cacheRoot, DiskCache.cacheFileName) : undefined;
-		const cache = new DiskCache(embeddingType, cachePath, workspaceIndex, fileSystem, logService);
-
-		if (cacheRoot && cachePath) {
-			await workspaceIndex.initialize();
-
-			const cacheValues = await instantiationService.invokeFunction(accessor => DiskCache.readDiskCache(accessor, embeddingType, cacheRoot, logService));
-			if (cacheValues) {
-				logService.debug(`Restoring workspace chunk + embeddings cache from ${cachePath.fsPath}`);
-
-				for (const [uriStr, entry] of cacheValues) {
-					const docUri = URI.parse(uriStr);
-					if (!workspaceIndex.get(docUri)) {
-						continue;
-					}
-
-					cache._cache.set(docUri, {
-						contentVersionId: entry.contentVersionId,
-						fileHash: entry.hash,
-						state: 'resolved',
-						value: entry.entries.map((x): FileChunkWithEmbedding => ({
-							embedding: {
-								value: typeof x.embedding === 'string' ? DiskCache.decodeEmbedding(x.embedding) : x.embedding,
-								type: embeddingType,
-							},
-							chunkHash: x.chunkHash,
-							chunk: {
-								file: docUri,
-								text: stripChunkTextMetadata(x.text),
-								rawText: undefined,
-								range: Range.lift(x.range),
-							} satisfies FileChunk
-						}))
-					});
-				}
-			}
-		}
-
-		return cache;
-	}
-
-	private readonly _cache = new ResourceMap<CacheEntry>();
-
-	private _isDisposed = false;
-
-	private readonly _writeDelayer = this._register(new ThrottledDelayer<void>(5000));
-
-	private constructor(
-		private readonly embeddingType: EmbeddingType,
-		private readonly cachePath: URI | undefined,
-		@IWorkspaceFileIndex private readonly _workspaceIndex: IWorkspaceFileIndex,
-		@IFileSystemService private readonly fileSystem: IFileSystemService,
-		@ILogService private readonly logService: ILogService
-	) {
-		super();
-
-		this._register(this._workspaceIndex.onDidDeleteFiles(uris => {
-			for (const uri of uris) {
-				this._cache.delete(uri);
-			}
-		}));
-	}
-
-	public override dispose(): void {
-		this._isDisposed = true;
-		super.dispose();
-	}
-
-	/**
-	 * Checks if {@linkcode file} is currently indexed. Does not wait for any current indexing operation to complete.
-	 */
-	async isIndexed(file: FileRepresentation): Promise<boolean> {
-		const entry = await this.getEntry(file);
-		return entry?.state === 'resolved';
-	}
-
-	async get(file: FileRepresentation): Promise<readonly FileChunkWithEmbedding[] | undefined> {
-		return (await this.getEntry(file))?.value;
-	}
-
-	getCurrentChunksForUri(uri: URI): ReadonlyMap<string, FileChunkWithEmbedding> | undefined {
-		const entry = this._cache.get(uri);
-		if (entry?.state === 'resolved' || entry?.state === 'rejected') {
-			if (entry.value) {
-				const out = new Map<string, FileChunkWithEmbedding>();
-				for (const x of entry.value) {
-					if (x.chunkHash) {
-						out.set(x.chunkHash, x);
-					}
-				}
-				return out;
-			}
-		}
-		return undefined;
-	}
-
-	private async getEntry(file: FileRepresentation): Promise<CacheEntry | undefined> {
-		const entry = this._cache.get(file.uri);
-		if (!entry) {
-			return undefined;
-		}
-
-		if (entry.contentVersionId === await file.getFastContentVersionId()) {
-			return entry;
-		}
-
-		return undefined;
-	}
-
-	async update(file: FileRepresentation, compute: (token: CancellationToken) => Promise<readonly FileChunkWithEmbedding[] | undefined>): Promise<readonly FileChunkWithEmbedding[] | undefined> {
-		const existing = this._cache.get(file.uri);
-		const inContentVersionId = await file.getFastContentVersionId();
-		if (existing?.contentVersionId === inContentVersionId) {
-			// Already up to date
-			return existing.value;
-		}
-
-		// Overwrite
-		if (existing?.state === 'pending') {
-			existing.value.cancel();
-		}
-		const chunks = createCancelablePromise(compute);
-		const entry: CacheEntry = {
-			contentVersionId: inContentVersionId,
-			fileHash: undefined,
-			state: 'pending',
-			value: chunks
-		};
-		this._cache.set(file.uri, entry);
-
-		chunks
-			.then((result): CacheEntry => {
-				return { contentVersionId: inContentVersionId, fileHash: undefined, state: Array.isArray(result) ? 'resolved' : 'rejected', value: result };
-			}, (): CacheEntry => {
-				return { contentVersionId: inContentVersionId, fileHash: undefined, state: 'rejected', value: undefined };
-			})
-			.then(newEntry => {
-				const current = this._cache.get(file.uri);
-				if (entry === current) {
-					this._cache.set(file.uri, newEntry);
-					return this._writeDelayer.trigger(() => this.save());
-				}
-			});
-
-		return chunks;
-	}
-
-	private async save() {
-		if (!this.cachePath || this._isDisposed) {
-			return;
-		}
-
-		const entries: Record<string, PersistedCacheEntry> = {};
-		await Promise.all(Array.from(this._cache.entries(), async ([uri, entry]) => {
-			let chunkAndEmbeddings: readonly FileChunkWithEmbedding[] | undefined;
-			try {
-				// Don't block saving on entries that are still resolving
-				chunkAndEmbeddings = entry.state === 'pending' ? await raceTimeout(entry.value, 1000) : entry.value;
-			} catch {
-				// noop
-			}
-
-			if (!chunkAndEmbeddings) {
-				return;
-			}
-
-			entries[uri.toString()] = {
-				contentVersionId: entry.contentVersionId,
-				hash: undefined,
-				entries: chunkAndEmbeddings.map(x => ({
-					text: x.chunk.text,
-					range: x.chunk.range.toJSON(),
-					embedding: DiskCache.encodeEmbedding(x.embedding.value),
-					chunkHash: x.chunkHash,
-				})),
-			};
-		}));
-
-		if (this._isDisposed) {
-			return;
-		}
-
-		const data: PersistedCache = {
-			version: DiskCache.version,
-			embeddingModel: this.embeddingType.id,
-			entries: entries,
-		};
-		await this.fileSystem.writeFile(this.cachePath, new TextEncoder().encode(JSON.stringify(data)));
-
-		this.logService.debug(`Wrote workspace chunk + embeddings cache to ${this.cachePath.fsPath}`);
-	}
+	private constructor() { }
 }
 
 
@@ -358,29 +145,30 @@ class DbCache implements IWorkspaceChunkAndEmbeddingCache {
 	public static async create(
 		accessor: ServicesAccessor,
 		embeddingType: EmbeddingType,
-		cacheRoot: URI,
+		cacheRoot: URI | ':memory:',
 		workspaceIndex: IWorkspaceFileIndex,
-	): Promise<DbCache | undefined> {
+	): Promise<DbCache> {
 		const instantiationService = accessor.get(IInstantiationService);
 
 		const syncOptions: sql.DatabaseSyncOptions = {
 			open: true,
 			enableForeignKeyConstraints: true
 		};
 
-		const dbPath = URI.joinPath(cacheRoot, `workspace-chunks.db`);
 
 		let db: sql.DatabaseSync | undefined;
-		if (dbPath.scheme === Schemas.file) {
+		if (cacheRoot !== ':memory:' && cacheRoot.scheme === Schemas.file) {
+			const dbPath = URI.joinPath(cacheRoot, `workspace-chunks.db`);
 			try {
 				await fs.promises.mkdir(path.dirname(dbPath.fsPath), { recursive: true });
 				db = new sql.DatabaseSync(dbPath.fsPath, syncOptions);
 			} catch (e) {
 				console.error('Failed to open SQLite database on disk', e);
 			}
 		}
+
 		if (!db) {
-			return;
+			db = new sql.DatabaseSync(':memory:', syncOptions);
 		}
 
 		db.exec(`
@@ -431,12 +219,14 @@ class DbCache implements IWorkspaceChunkAndEmbeddingCache {
 		db.prepare('INSERT INTO CacheMeta (version, embeddingModel) VALUES (?, ?)').run(this.version, embeddingType.id);
 
 		// Load existing disk db if it exists
-		const diskCache = await instantiationService.invokeFunction(accessor => DiskCache.readDiskCache(
-			accessor,
-			embeddingType,
-			cacheRoot,
-			accessor.get(ILogService)
-		));
+		const diskCache = cacheRoot !== ':memory:' ?
+			await instantiationService.invokeFunction(accessor => OldDiskCache.readDiskCache(
+				accessor,
+				embeddingType,
+				cacheRoot,
+				accessor.get(ILogService)
+			))
+			: undefined;
 		if (diskCache) {
 			try {
 				const insertFileStatement = db.prepare('INSERT OR REPLACE INTO Files (uri, contentVersionId) VALUES (?, ?)');
@@ -457,7 +247,7 @@ class DbCache implements IWorkspaceChunkAndEmbeddingCache {
 							chunk.range.endColumn,
 							packEmbedding({
 								type: embeddingType,
-								value: typeof chunk.embedding === 'string' ? DiskCache.decodeEmbedding(chunk.embedding) : chunk.embedding,
+								value: typeof chunk.embedding === 'string' ? OldDiskCache.decodeEmbedding(chunk.embedding) : chunk.embedding,
 							}),
 							chunk.chunkHash ?? ''
 						);
@@ -467,7 +257,9 @@ class DbCache implements IWorkspaceChunkAndEmbeddingCache {
 				db.exec('COMMIT');
 			}
 
-			void instantiationService.invokeFunction(accessor => DiskCache.deleteDiskCache(accessor, cacheRoot));
+			if (cacheRoot !== ':memory:') {
+				void instantiationService.invokeFunction(accessor => OldDiskCache.deleteDiskCache(accessor, cacheRoot));
+			}
 		}
 
 		// Validate all files in the database against the workspace index and remove any that are no longer present