@davideasden/pi-undo 0.1.2 → 0.2.1

package/src/restore-engine.ts

@@ -19,10 +19,15 @@ import {
 	fingerprintAbsent,
 	fingerprintBytes,
 	fingerprintSymlink,
+	type DeleteLeafRequest,
+	type ReplaceFileRequest,
 } from "./quarantine.ts";
 import { SnapshotStoreError, type SnapshotStore } from "./snapshot-store.ts";
 
 const PREPARED_PLAN_CACHE_LIMIT = 16;
+const RESTORE_FILE_BATCH_MAX_ENTRIES = 128;
+const RESTORE_FILE_BATCH_MAX_BYTES = 32 * 1024 * 1024;
+const RESTORE_FILE_PREPARE_CONCURRENCY = 32;
 
 export interface RestorePlan {
 	currentManifestId: ManifestId;
@@ -128,17 +133,18 @@ export class RestoreEngine {
 	): Promise<RestorePlan> {
 		assertManifest(current);
 		assertManifest(target);
-		assertCompatibleManifests(current, target);
 		const scope = scopePaths === undefined ? undefined : this.canonicalScope(scopePaths);
+		const canonicalScopePaths = scope === undefined ? undefined : [...scope];
+		assertCompatibleManifests(current, target, scope);
 		const isScopedPath = (path: string): boolean => scope === undefined || scope.has(path);
 		await Promise.all([
-			this.store.assertComplete(current.manifestId),
-			this.store.assertComplete(target.manifestId),
+			this.store.assertComplete(current.manifestId, canonicalScopePaths),
+			this.store.assertComplete(target.manifestId, canonicalScopePaths),
 		]);
 
 		const [currentPaths, targetPaths] = await Promise.all([
-			this.readOwnedPaths(current),
-			this.readOwnedPaths(target),
+			this.readOwnedPaths(current, canonicalScopePaths),
+			this.readOwnedPaths(target, canonicalScopePaths),
 		]);
 		const targetIgnoredPaths = ignoredWorkspacePaths(target);
 		const deleteByRoot = new Map<string, string[]>();
@@ -315,7 +321,11 @@ export class RestoreEngine {
 		if (canonicalJson(storedTarget) !== canonicalJson(target)) {
 			throw new Error("target manifest 与 store 内容不一致");
 		}
-		assertCompatibleManifests(current, target);
+		assertCompatibleManifests(
+			current,
+			target,
+			plan.scopePaths === undefined ? undefined : this.canonicalScope(plan.scopePaths),
+		);
 		let expectedPlan: RestorePlan;
 		let prepared: PreparedRestorePlan | undefined;
 		try {
@@ -342,7 +352,10 @@ export class RestoreEngine {
 			return { code: "restore_failed_safe", verifiedPaths: 0, totalPaths: 0 };
 		}
 		const [currentPaths, targetPaths] = prepared === undefined
-			? await Promise.all([this.readOwnedPaths(current), this.readOwnedPaths(target)])
+			? await Promise.all([
+				this.readOwnedPaths(current, plan.scopePaths),
+				this.readOwnedPaths(target, plan.scopePaths),
+			])
 			: [prepared.currentPaths, prepared.targetPaths];
 		const quarantine = new QuarantineManager({
 			workspaceRoot: this.requestedWorkspaceRoot,
@@ -382,12 +395,12 @@ export class RestoreEngine {
 			quarantine,
 		};
 		try {
-			for (const path of plan.deletePaths) {
-				if (await this.pathIsShadowedByTarget(target.manifestId, path, targetPaths)) {
-					continue;
-				}
-				await this.deletePath(path, mutationContext);
-			}
+			await this.deletePlannedPaths(
+				target.manifestId,
+				targetPaths,
+				plan.deletePaths,
+				mutationContext,
+			);
 			await this.writePlannedPaths(target.manifestId, targetPaths, plan.writePaths, mutationContext);
 
 			const topologyAfter = await this.discovery.discover(this.workspaceRoot);
@@ -425,11 +438,16 @@ export class RestoreEngine {
 		}
 	}
 
-	private async readOwnedPaths(manifest: SnapshotManifest): Promise<Map<string, OwnedPath>> {
+	private async readOwnedPaths(
+		manifest: SnapshotManifest,
+		scopePaths?: readonly string[],
+	): Promise<Map<string, OwnedPath>> {
 		const result = new Map<string, OwnedPath>();
 		for (const root of manifest.roots) {
 			assertNotGitMetadata(root.relativeRoot);
-			const entries = await this.store.listTree(manifest.manifestId, root.relativeRoot);
+			const rootScope = rootRelativeScopePaths(root.relativeRoot, scopePaths);
+			if (rootScope !== undefined && rootScope.length === 0) continue;
+			const entries = await this.store.listTree(manifest.manifestId, root.relativeRoot, rootScope);
 			if (entries.length > 0) {
 				for (const boundaryPath of rootBoundaryDirectories(root.relativeRoot)) {
 					if (!result.has(boundaryPath)) {
@@ -523,6 +541,63 @@ export class RestoreEngine {
 		return { ok: true, verifiedPaths, totalPaths: paths.length };
 	}
 
+	private async deletePlannedPaths(
+		targetManifestId: ManifestId,
+		targetPaths: ReadonlyMap<string, OwnedPath>,
+		deletePaths: readonly string[],
+		context: MutationContext,
+	): Promise<void> {
+		let fileBatch: OwnedPath[] = [];
+		let fileBatchRoot: string | undefined;
+		const flushFiles = async (): Promise<void> => {
+			if (fileBatch.length === 0) return;
+			const requests: DeleteLeafRequest[] = [];
+			for (const source of fileBatch) {
+				context.ordinal += 1;
+				const ordinal = context.ordinal;
+				await this.beforeMutation?.({
+					phase: context.phase,
+					ordinal,
+					kind: "delete",
+					path: source.absolutePath,
+				});
+				await this.assertMutationPath(source.absolutePath);
+				await this.assertMutationState(context, "delete", source.absolutePath);
+				requests.push({
+					path: source.absolutePath,
+					sourceFingerprint: await this.expectedMutationFingerprint(context, source.absolutePath),
+					targetFingerprint: fingerprintAbsent(source.absolutePath),
+				});
+			}
+			await context.quarantine.deleteFiles(requests);
+			fileBatch = [];
+			fileBatchRoot = undefined;
+		};
+		for (const path of deletePaths) {
+			if (await this.pathIsShadowedByTarget(targetManifestId, path, targetPaths)) continue;
+			const source = context.sourcePaths.get(path);
+			const live = await lstat(this.absolutePath(path)).catch((error) => {
+				if (hasErrorCode(error, "ENOENT") || hasErrorCode(error, "ENOTDIR")) return null;
+				throw error;
+			});
+			if (live === null) continue;
+			if (source?.entry.kind !== "file" || live.isSymbolicLink() || !live.isFile()) {
+				await flushFiles();
+				await this.deletePath(path, context);
+				continue;
+			}
+			if (
+				fileBatch.length > 0 &&
+				(fileBatchRoot !== source.root.relativeRoot || fileBatch.length >= RESTORE_FILE_BATCH_MAX_ENTRIES)
+			) {
+				await flushFiles();
+			}
+			fileBatch.push(source);
+			fileBatchRoot = source.root.relativeRoot;
+		}
+		await flushFiles();
+	}
+
 	private async deletePath(path: string, context: MutationContext): Promise<void> {
 		const absolutePath = this.absolutePath(path);
 		try {
@@ -557,12 +632,12 @@ export class RestoreEngine {
 			return;
 		}
 		await this.mutate(context, "delete", path, async () => {
-			await context.quarantine.deleteLeaf({
+			const record = await context.quarantine.deleteLeaf({
 				path,
 				sourceFingerprint: await this.expectedMutationFingerprint(context, path),
 				targetFingerprint: fingerprintAbsent(path),
 			});
-			await this.cleanupLatestMutation(context);
+			await context.quarantine.cleanupMutation(record);
 		});
 	}
 
@@ -604,14 +679,14 @@ export class RestoreEngine {
 				throw new Error(`symlink 存在类型或内容冲突：${path}`);
 			}
 			await this.mutate(context, "symlink", path, async (beforeInstall) => {
-				await context.quarantine.replaceSymlink({
+				const record = await context.quarantine.replaceSymlink({
 					path,
 					targetLinkText: linkText,
 					sourceFingerprint: await this.expectedMutationFingerprint(context, path),
 					targetFingerprint: fingerprintSymlink(path, linkText),
 					beforeInstall,
 				});
-				await this.cleanupLatestMutation(context);
+				await context.quarantine.cleanupMutation(record);
 			}, true);
 			return;
 		}
@@ -619,7 +694,12 @@ export class RestoreEngine {
 		if (target.entry.blobId === null) {
 			throw new Error(`普通文件缺少 blob：${path}`);
 		}
-		const bytes = await this.store.readBlob(manifestId, target.root.relativeRoot, target.entry.blobId);
+		const bytes = await this.store.readBlob(
+			manifestId,
+			target.root.relativeRoot,
+			target.entry.blobId,
+			target.entry.relativePath,
+		);
 		if (bytes.byteLength !== target.entry.size) {
 			throw new Error(`普通文件 blob 大小不匹配：${path}`);
 		}
@@ -628,7 +708,7 @@ export class RestoreEngine {
 			"write",
 			path,
 			async (beforeInstall) => {
-				await context.quarantine.replaceFile({
+				const record = await context.quarantine.replaceFile({
 					path,
 					targetBytes: bytes,
 					targetMode: target.entry.mode & 0o777,
@@ -636,7 +716,7 @@ export class RestoreEngine {
 					targetFingerprint: fingerprintBytes(path, bytes, target.entry.mode),
 					beforeInstall,
 				});
-				await this.cleanupLatestMutation(context);
+				await context.quarantine.cleanupMutation(record);
 			},
 			true,
 		);
@@ -669,12 +749,12 @@ export class RestoreEngine {
 					journal: options.mutationJournal,
 				}),
 			};
-			for (const path of rollbackPlan.deletePaths) {
-				if (await this.pathIsShadowedByTarget(current.manifestId, path, currentPaths)) {
-					continue;
-				}
-				await this.deletePath(path, context);
-			}
+			await this.deletePlannedPaths(
+				current.manifestId,
+				currentPaths,
+				rollbackPlan.deletePaths,
+				context,
+			);
 			await this.writePlannedPaths(current.manifestId, currentPaths, rollbackPlan.writePaths, context);
 			const topologyAfter = await this.discovery.discover(this.workspaceRoot);
 			assertUnchangedTopology(topologyBefore, topologyAfter);
@@ -766,8 +846,9 @@ export class RestoreEngine {
 		deferHook = false,
 	): Promise<void> {
 		context.ordinal += 1;
+		const ordinal = context.ordinal;
 		const beforeInstall = async (): Promise<void> => {
-			await this.beforeMutation?.({ phase: context.phase, ordinal: context.ordinal, kind, path });
+			await this.beforeMutation?.({ phase: context.phase, ordinal, kind, path });
 		};
 		if (!deferHook) await beforeInstall();
 		await this.assertMutationPath(path);
@@ -781,26 +862,95 @@ export class RestoreEngine {
 		writePaths: readonly string[],
 		context: MutationContext,
 	): Promise<void> {
+		let fileBatch: OwnedPath[] = [];
+		let fileBatchBytes = 0;
+		let fileBatchRoot: string | undefined;
+		const flushFiles = async (): Promise<void> => {
+			if (fileBatch.length === 0) return;
+			const pending = fileBatch.map((target) => {
+				context.ordinal += 1;
+				return { target, ordinal: context.ordinal };
+			});
+			const requests = await mapConcurrentOrdered(
+				pending,
+				RESTORE_FILE_PREPARE_CONCURRENCY,
+				({ target, ordinal }) => this.prepareFileReplacement(manifestId, target, context, ordinal),
+			);
+			await context.quarantine.replaceFiles(requests);
+			fileBatch = [];
+			fileBatchBytes = 0;
+			fileBatchRoot = undefined;
+		};
 		for (const kind of ["directory", "leaf"] as const) {
 			for (const path of writePaths) {
 				const target = targetPaths.get(path);
 				if (target === undefined) {
 					throw new Error(`${context.phase} plan 引用了 manifest 外路径：${path}`);
 				}
-				if ((target.entry.kind === "directory") !== (kind === "directory")) {
-					continue;
-				}
+				if ((target.entry.kind === "directory") !== (kind === "directory")) continue;
 				if (context.sourceIgnoredPaths.has(path)) {
-					if (await this.entryMatches(manifestId, target)) {
-						continue;
-					}
+					await flushFiles();
+					if (await this.entryMatches(manifestId, target)) continue;
 					throw new Error(`${context.phase} 的 ignored-present 路径与目标内容冲突：${path}`);
 				}
-				await this.writePath(manifestId, target, context);
+				if (target.entry.kind !== "file") {
+					await flushFiles();
+					await this.writePath(manifestId, target, context);
+					continue;
+				}
+				if (
+					fileBatch.length > 0 &&
+					(fileBatchRoot !== target.root.relativeRoot ||
+						fileBatch.length >= RESTORE_FILE_BATCH_MAX_ENTRIES ||
+						fileBatchBytes + target.entry.size > RESTORE_FILE_BATCH_MAX_BYTES)
+				) {
+					await flushFiles();
+				}
+				fileBatch.push(target);
+				fileBatchBytes += target.entry.size;
+				fileBatchRoot = target.root.relativeRoot;
 			}
+			await flushFiles();
 		}
 	}
 
+	private async prepareFileReplacement(
+		manifestId: ManifestId,
+		target: OwnedPath,
+		context: MutationContext,
+		ordinal: number,
+	): Promise<ReplaceFileRequest> {
+		if (target.entry.kind !== "file" || target.entry.blobId === null) {
+			throw new Error(`批量普通文件缺少 blob：${target.absolutePath}`);
+		}
+		const bytes = await this.store.readBlob(
+			manifestId,
+			target.root.relativeRoot,
+			target.entry.blobId,
+			target.entry.relativePath,
+		);
+		if (bytes.byteLength !== target.entry.size) {
+			throw new Error(`普通文件 blob 大小不匹配：${target.absolutePath}`);
+		}
+		await this.assertMutationPath(target.absolutePath);
+		await this.assertMutationState(context, "write", target.absolutePath);
+		return {
+			path: target.absolutePath,
+			targetBytes: bytes,
+			targetMode: target.entry.mode & 0o777,
+			sourceFingerprint: await this.expectedMutationFingerprint(context, target.absolutePath),
+			targetFingerprint: fingerprintBytes(target.absolutePath, bytes, target.entry.mode),
+			beforeInstall: async () => {
+				await this.beforeMutation?.({
+					phase: context.phase,
+					ordinal,
+					kind: "write",
+					path: target.absolutePath,
+				});
+			},
+		};
+	}
+
 	private async assertMutationState(
 		context: MutationContext,
 		kind: RestoreMutation["kind"],
@@ -933,7 +1083,12 @@ export class RestoreEngine {
 				return fingerprintSymlink(path, owned.entry.linkText!);
 			}
 			if (owned.entry.blobId === null) throw new Error(`普通文件缺少 blob：${path}`);
-			const bytes = await this.store.readBlob(manifestId, owned.root.relativeRoot, owned.entry.blobId);
+			const bytes = await this.store.readBlob(
+				manifestId,
+				owned.root.relativeRoot,
+				owned.entry.blobId,
+				owned.entry.relativePath,
+			);
 			return fingerprintBytes(path, bytes, owned.entry.mode);
 		}
 		if (await this.pathIsAbsent(path)) return fingerprintAbsent(path);
@@ -949,13 +1104,6 @@ export class RestoreEngine {
 		}
 	}
 
-	private async cleanupLatestMutation(context: MutationContext): Promise<void> {
-		const records = await context.mutationJournal.load();
-		const record = records.at(-1);
-		if (record === undefined) throw new Error("quarantine mutation 记录缺失");
-		await context.quarantine.cleanupMutation(record);
-	}
-
 	private async restorePendingMutations(
 		quarantine: QuarantineManager,
 		journal: MutationJournal,
@@ -1012,7 +1160,12 @@ export class RestoreEngine {
 		}
 		const [actual, expected] = await Promise.all([
 			readFile(this.absolutePath(path)),
-			this.store.readBlob(manifestId, owned.root.relativeRoot, owned.entry.blobId),
+			this.store.readBlob(
+				manifestId,
+				owned.root.relativeRoot,
+				owned.entry.blobId,
+				owned.entry.relativePath,
+			),
 		]);
 		if (!actual.equals(Buffer.from(expected))) {
 			throw new Error(`普通文件内容校验失败：${path}`);
@@ -1081,16 +1234,37 @@ function sameEntry(left: RestorePath, right: RestorePath): boolean {
 		left.linkText === right.linkText;
 }
 
-function assertCompatibleManifests(current: SnapshotManifest, target: SnapshotManifest): void {
+function assertCompatibleManifests(
+	current: SnapshotManifest,
+	target: SnapshotManifest,
+	scope?: ReadonlySet<string>,
+): void {
 	if (current.workspaceIdentity !== target.workspaceIdentity) {
 		throw new Error("restore manifest 不属于同一 workspace");
 	}
-	if (current.coverage !== target.coverage) {
-		throw new Error("restore manifest coverage 不一致，不能推断缺失路径");
-	}
 	if (current.roots.some((root) => root.state === "broken") || target.roots.some((root) => root.state === "broken")) {
 		throw new Error("broken root 不能用于 restore");
 	}
+	const scopedCoverage = scope === undefined
+		? undefined
+		: `paths:${checksum(canonicalJson([...scope]))}`;
+	if (current.coverage === target.coverage) {
+		if (
+			scopedCoverage !== undefined && current.coverage !== "complete" &&
+			current.coverage !== scopedCoverage
+		) {
+			throw new Error("restore manifest coverage 与 scope 不匹配");
+		}
+		return;
+	}
+	if (
+		scopedCoverage === undefined ||
+		![current.coverage, target.coverage].every(
+			(coverage) => coverage === "complete" || coverage === scopedCoverage,
+		)
+	) {
+		throw new Error("restore manifest coverage 不一致，不能推断缺失路径");
+	}
 }
 
 function orderedWritePaths(
@@ -1120,6 +1294,26 @@ function appendPath(paths: Map<string, string[]>, root: string, path: string): v
 	owned.push(path);
 }
 
+function rootRelativeScopePaths(
+	root: string,
+	scopePaths: readonly string[] | undefined,
+): string[] | undefined {
+	if (scopePaths === undefined) return undefined;
+	if (scopePaths.length === 0) return [];
+	const result = new Set<string>();
+	for (const path of scopePaths) {
+		if (path === "." || path === root || isStrictWorkspaceAncestor(path, root)) return undefined;
+		if (isStrictWorkspaceAncestor(root, path)) {
+			result.add(root === "." ? path : path.slice(root.length + 1));
+		}
+	}
+	return [...result].sort(comparePaths);
+}
+
+function isStrictWorkspaceAncestor(parent: string, child: string): boolean {
+	return parent === "." ? child !== "." : child.startsWith(`${parent}/`);
+}
+
 function workspacePath(root: string, path: string): string {
 	return root === "." ? path : path === "." ? root : `${root}/${path}`;
 }
@@ -1227,6 +1421,32 @@ function hasValidPlanDigest(plan: RestorePlan): boolean {
 	}
 }
 
+async function mapConcurrentOrdered<T, R>(
+	values: readonly T[],
+	concurrency: number,
+	operation: (value: T) => Promise<R>,
+): Promise<R[]> {
+	const results: R[] = new Array(values.length);
+	let nextIndex = 0;
+	let failed = false;
+	let failure: unknown;
+	async function worker(): Promise<void> {
+		while (!failed && nextIndex < values.length) {
+			const index = nextIndex;
+			nextIndex += 1;
+			try {
+				results[index] = await operation(values[index]!);
+			} catch (error) {
+				if (!failed) failure = error;
+				failed = true;
+			}
+		}
+	}
+	await Promise.all(Array.from({ length: Math.min(concurrency, values.length) }, () => worker()));
+	if (failed) throw failure;
+	return results;
+}
+
 function hasErrorCode(error: unknown, code: string): error is NodeJS.ErrnoException {
 	return typeof error === "object" && error !== null && "code" in error && error.code === code;
 }