fix: speed up multi-file apply-patch
This commit is contained in:
Codex
@@ -115,6 +115,14 @@ export class ApplyPatchV2Error extends Error {
|
||||
type PlannedFileState = { exists: true; content: string } | { exists: false; content: "" };
|
||||
type PlannedOperation = { kind: "write"; path: string; content: string } | { kind: "delete"; path: string };
|
||||
type Replacement = [start: number, oldLength: number, newLines: string[]];
|
||||
interface BulkReplacementWritePlan {
|
||||
path: string;
|
||||
originalBytes: number;
|
||||
originalSha256: string;
|
||||
finalBytes: number;
|
||||
finalSha256: string;
|
||||
replacements: Replacement[];
|
||||
}
|
||||
|
||||
interface ApplyPatchV2Plan {
|
||||
changed: string[];
|
||||
@@ -516,7 +524,17 @@ function applyPatchV2ChangedCountFromOutcomes(outcomes: ApplyPatchV2Outcome[]):
|
||||
|
||||
async function applyPatchV2Hunks(executor: ApplyPatchV2Executor, hunks: PatchHunk[]): Promise<ApplyPatchV2Plan> {
|
||||
if (hunks.length === 0) throw new ApplyPatchV2Error("No files were modified.");
|
||||
if (shouldUseBulkUpdatePath(executor, hunks)) {
|
||||
try {
|
||||
return await applyPatchV2UpdateHunksBulk(executor, hunks as Array<Extract<PatchHunk, { kind: "update" }>>);
|
||||
} catch (error) {
|
||||
if (!isBulkReadFailure(error)) throw error;
|
||||
}
|
||||
}
|
||||
return applyPatchV2HunksSequential(executor, hunks);
|
||||
}
|
||||
|
||||
async function applyPatchV2HunksSequential(executor: ApplyPatchV2Executor, hunks: PatchHunk[]): Promise<ApplyPatchV2Plan> {
|
||||
const states = new Map<string, PlannedFileState>();
|
||||
const changed: string[] = [];
|
||||
const outcomes: ApplyPatchV2Outcome[] = [];
|
||||
@@ -596,6 +614,83 @@ async function applyPatchV2Hunks(executor: ApplyPatchV2Executor, hunks: PatchHun
|
||||
return { changed, outcomes };
|
||||
}
|
||||
|
||||
function isBulkReadFailure(error: unknown): boolean {
|
||||
if (!(error instanceof ApplyPatchV2Error)) return false;
|
||||
if (error.message.includes("bulk read")) return true;
|
||||
return JSON.stringify(error.details).includes("read-bulk-b64");
|
||||
}
|
||||
|
||||
function shouldUseBulkUpdatePath(executor: ApplyPatchV2Executor, hunks: PatchHunk[]): boolean {
|
||||
if (!executor.run || executor.fs !== undefined) return false;
|
||||
const paths = new Set<string>();
|
||||
for (const hunk of hunks) {
|
||||
if (hunk.kind !== "update" || hunk.movePath !== null) return false;
|
||||
if (paths.has(hunk.path)) return false;
|
||||
paths.add(hunk.path);
|
||||
}
|
||||
return paths.size >= 2;
|
||||
}
|
||||
|
||||
async function applyPatchV2UpdateHunksBulk(executor: ApplyPatchV2Executor, hunks: Array<Extract<PatchHunk, { kind: "update" }>>): Promise<ApplyPatchV2Plan> {
|
||||
const paths = uniqueStrings(hunks.map((hunk) => hunk.path));
|
||||
const originalFiles = await readRemoteTextsBulk(executor, paths);
|
||||
const states = new Map<string, PlannedFileState>();
|
||||
for (const [filePath, content] of originalFiles) states.set(filePath, { exists: true, content });
|
||||
|
||||
const changed: string[] = [];
|
||||
const outcomes: ApplyPatchV2Outcome[] = [];
|
||||
const pendingWrites = new Set<string>();
|
||||
const replacementPlans = new Map<string, BulkReplacementWritePlan>();
|
||||
|
||||
for (let index = 0; index < hunks.length; index += 1) {
|
||||
const hunk = hunks[index] as Extract<PatchHunk, { kind: "update" }>;
|
||||
const changedBefore = changed.length;
|
||||
try {
|
||||
const state = states.get(hunk.path);
|
||||
if (state === undefined || !state.exists) throw new ApplyPatchV2Error("cannot update a missing file in bulk patch", { path: hunk.path });
|
||||
const originalLines = splitContentLines(state.content);
|
||||
const replacements = computeReplacements(hunk.path, originalLines, hunk.chunks);
|
||||
const newContent = joinLinesWithFinalNewline(applyReplacements(originalLines, replacements));
|
||||
const originalBuffer = Buffer.from(state.content, "utf8");
|
||||
const finalBuffer = Buffer.from(newContent, "utf8");
|
||||
states.set(hunk.path, { exists: true, content: newContent });
|
||||
replacementPlans.set(hunk.path, {
|
||||
path: hunk.path,
|
||||
originalBytes: originalBuffer.length,
|
||||
originalSha256: sha256Hex(originalBuffer),
|
||||
finalBytes: finalBuffer.length,
|
||||
finalSha256: sha256Hex(finalBuffer),
|
||||
replacements,
|
||||
});
|
||||
pendingWrites.add(hunk.path);
|
||||
pushChanged(changed, `M ${hunk.path}`);
|
||||
outcomes.push({ ...outcomeBase(hunk, index), status: "applied", change: `M ${hunk.path}` });
|
||||
} catch (error) {
|
||||
const partialChanges = changed.slice(changedBefore);
|
||||
if (pendingWrites.size > 0) await writeRemoteReplacementsBulk(executor, pendingWrites, replacementPlans);
|
||||
outcomes.push({
|
||||
...outcomeBase(hunk, index),
|
||||
status: "failed",
|
||||
...(partialChanges.length > 0 ? { partialChanges } : {}),
|
||||
error: errorSummary(error),
|
||||
});
|
||||
throw new ApplyPatchV2Error(error instanceof Error ? error.message : String(error), {
|
||||
partialChanges: changed,
|
||||
outcomes,
|
||||
failed: outcomes.find((item) => item.status === "failed") ?? null,
|
||||
cause: error instanceof ApplyPatchV2Error ? error.details : undefined,
|
||||
});
|
||||
}
|
||||
}
|
||||
|
||||
if (pendingWrites.size > 0) await writeRemoteReplacementsBulk(executor, pendingWrites, replacementPlans);
|
||||
return { changed, outcomes };
|
||||
}
|
||||
|
||||
function uniqueStrings(values: string[]): string[] {
|
||||
return Array.from(new Set(values));
|
||||
}
|
||||
|
||||
function pushChanged(changed: string[], item: string): void {
|
||||
if (!changed.includes(item)) changed.push(item);
|
||||
}
|
||||
@@ -743,6 +838,7 @@ const readBlockBytes = 45_000;
|
||||
const writeB64ArgvLimit = 48_000;
|
||||
const writeB64ChunkChars = 12_000;
|
||||
const remoteReadBlockMarker = "UNIDESK_APPLY_PATCH_V2_BLOCK";
|
||||
const remoteBulkReadMarker = "UNIDESK_APPLY_PATCH_V2_BULK_READ";
|
||||
|
||||
async function readRemoteText(executor: ApplyPatchV2Executor, target: string): Promise<string> {
|
||||
if (executor.fs) {
|
||||
@@ -828,6 +924,17 @@ async function readRemoteText(executor: ApplyPatchV2Executor, target: string): P
|
||||
return contentBuffer.toString("utf8");
|
||||
}
|
||||
|
||||
async function readRemoteTextsBulk(executor: ApplyPatchV2Executor, targets: string[]): Promise<Map<string, string>> {
|
||||
if (targets.length === 0) return new Map();
|
||||
if (executor.fs || targets.length === 1) {
|
||||
const files = new Map<string, string>();
|
||||
for (const target of targets) files.set(target, await readRemoteText(executor, target));
|
||||
return files;
|
||||
}
|
||||
const result = await checkedRemoteV2(executor, "read-bulk-b64", targets);
|
||||
return decodeRemoteBulkRead(result.stdout, targets);
|
||||
}
|
||||
|
||||
function decodeRemoteReadBlock(stdout: string, target: string, blockIndex: number, expectedChunkBytes: number): Buffer {
|
||||
const lines = stdout.split(/\r?\n/u);
|
||||
const markerIndex = lines.findIndex((line) => line.startsWith(`${remoteReadBlockMarker} `));
|
||||
@@ -872,6 +979,47 @@ function decodeRemoteReadBlock(stdout: string, target: string, blockIndex: numbe
|
||||
return decoded.length > expectedChunkBytes ? decoded.subarray(0, expectedChunkBytes) : decoded;
|
||||
}
|
||||
|
||||
function decodeRemoteBulkRead(stdout: string, targets: string[]): Map<string, string> {
|
||||
const lines = stdout.split(/\r?\n/u);
|
||||
const markerIndex = lines.findIndex((line) => line.startsWith(`${remoteBulkReadMarker} `));
|
||||
if (markerIndex < 0) {
|
||||
throw new ApplyPatchV2Error("remote apply-patch v2 bulk read returned no marker", { expectedTargets: targets.length, stdout: stdout.slice(0, 500) });
|
||||
}
|
||||
const declaredCount = Number(lines[markerIndex].split(/\s+/u)[1] ?? "-1");
|
||||
if (!Number.isSafeInteger(declaredCount) || declaredCount !== targets.length) {
|
||||
throw new ApplyPatchV2Error("remote apply-patch v2 bulk read target count mismatch", { expectedTargets: targets.length, declaredCount });
|
||||
}
|
||||
const records = lines.slice(markerIndex + 1).filter((line) => line.trim().length > 0);
|
||||
if (records.length !== declaredCount) {
|
||||
throw new ApplyPatchV2Error("remote apply-patch v2 bulk read record count mismatch", { expectedTargets: targets.length, records: records.length });
|
||||
}
|
||||
const files = new Map<string, string>();
|
||||
for (const record of records) {
|
||||
const [pathB64, bytesText, sha256, contentB64] = record.split(/\s+/u);
|
||||
const target = Buffer.from(pathB64, "base64").toString("utf8");
|
||||
const expectedBytes = Number(bytesText);
|
||||
const contentBuffer = Buffer.from(contentB64 ?? "", "base64");
|
||||
if (!targets.includes(target)) {
|
||||
throw new ApplyPatchV2Error("remote apply-patch v2 bulk read returned unexpected target", { target, expectedTargets: targets });
|
||||
}
|
||||
if (!Number.isSafeInteger(expectedBytes) || expectedBytes < 0 || !/^[0-9a-f]{64}$/u.test(sha256)) {
|
||||
throw new ApplyPatchV2Error("remote apply-patch v2 bulk read returned invalid metadata", { target, expectedBytes: bytesText, sha256 });
|
||||
}
|
||||
if (contentBuffer.length !== expectedBytes) {
|
||||
throw new ApplyPatchV2Error(`remote apply-patch v2 bulk read byte count mismatch for ${target}: expected=${expectedBytes} actual=${contentBuffer.length}`, { target, expectedBytes, actualBytes: contentBuffer.length });
|
||||
}
|
||||
const actualSha256 = sha256Hex(contentBuffer);
|
||||
if (actualSha256 !== sha256) {
|
||||
throw new ApplyPatchV2Error(`remote apply-patch v2 bulk read sha256 mismatch for ${target}`, { target, expectedSha256: sha256, actualSha256 });
|
||||
}
|
||||
files.set(target, contentBuffer.toString("utf8"));
|
||||
}
|
||||
for (const target of targets) {
|
||||
if (!files.has(target)) throw new ApplyPatchV2Error("remote apply-patch v2 bulk read omitted target", { target });
|
||||
}
|
||||
return files;
|
||||
}
|
||||
|
||||
async function writeRemoteText(executor: ApplyPatchV2Executor, target: string, content: string): Promise<void> {
|
||||
const contentBuffer = Buffer.from(content, "utf8");
|
||||
if (executor.fs) {
|
||||
@@ -901,9 +1049,38 @@ async function writeRemoteText(executor: ApplyPatchV2Executor, target: string, c
|
||||
await checkedRemoteV2(executor, "write-b64-commit", [target, token, expectedBytes, expectedSha256]);
|
||||
}
|
||||
|
||||
async function writeRemoteReplacementsBulk(executor: ApplyPatchV2Executor, paths: Iterable<string>, plans: Map<string, BulkReplacementWritePlan>): Promise<void> {
|
||||
const targets = Array.from(paths);
|
||||
if (targets.length === 0) return;
|
||||
const records: string[] = [];
|
||||
for (const target of targets) {
|
||||
const plan = plans.get(target);
|
||||
if (plan === undefined) throw new ApplyPatchV2Error("missing bulk replacement plan", { target });
|
||||
const replacementFields: string[] = [];
|
||||
for (const [start, oldLength, newLines] of plan.replacements) {
|
||||
replacementFields.push([
|
||||
String(start),
|
||||
String(oldLength),
|
||||
Buffer.from(joinLinesWithFinalNewline(newLines), "utf8").toString("base64"),
|
||||
].join(","));
|
||||
}
|
||||
records.push([
|
||||
Buffer.from(plan.path, "utf8").toString("base64"),
|
||||
String(plan.originalBytes),
|
||||
plan.originalSha256,
|
||||
String(plan.finalBytes),
|
||||
plan.finalSha256,
|
||||
replacementFields.join(";"),
|
||||
].join(" "));
|
||||
}
|
||||
await checkedRemoteV2(executor, "apply-replacements-bulk-stdin", [String(targets.length)], `${records.join("\n")}\n`);
|
||||
}
|
||||
|
||||
type RemoteV2Operation =
|
||||
| "stat"
|
||||
| "read-b64-block"
|
||||
| "read-bulk-b64"
|
||||
| "apply-replacements-bulk-stdin"
|
||||
| "write-b64-argv"
|
||||
| "write-b64-stdin"
|
||||
| "write-b64-begin"
|
||||
@@ -989,6 +1166,80 @@ function remoteV2Script(operation: RemoteV2Operation, args: string[]): string[]
|
||||
" printf '\\n'",
|
||||
" rm -f -- \"$tmp\"",
|
||||
" ;;",
|
||||
" read-bulk-b64)",
|
||||
" count=$#",
|
||||
` printf '${remoteBulkReadMarker} %s\\n' "$count"`,
|
||||
" for target in \"$@\"; do",
|
||||
" if [ ! -e \"$target\" ]; then printf 'file not found: %s\\n' \"$target\" >&2; exit 1; fi",
|
||||
" if [ -d \"$target\" ]; then printf 'not a file: %s\\n' \"$target\" >&2; exit 1; fi",
|
||||
" path_b64=$(printf '%s' \"$target\" | base64 | tr -d '\\n')",
|
||||
" bytes=$(wc -c < \"$target\" | tr -d '[:space:]')",
|
||||
" digest=$(sha256_file \"$target\")",
|
||||
" printf '%s %s %s ' \"$path_b64\" \"$bytes\" \"$digest\"",
|
||||
" base64 < \"$target\" | tr -d '\\n'",
|
||||
" printf '\\n'",
|
||||
" done",
|
||||
" ;;",
|
||||
" apply-replacements-bulk-stdin)",
|
||||
" expected_count=$1",
|
||||
" tmp_payload=$(mktemp)",
|
||||
" tmp_dir=$(mktemp -d)",
|
||||
" trap 'rm -f -- \"$tmp_payload\"; rm -rf -- \"$tmp_dir\"' EXIT HUP INT TERM",
|
||||
" cat > \"$tmp_payload\"",
|
||||
" python3 - \"$tmp_payload\" \"$tmp_dir\" \"$expected_count\" <<'PY'",
|
||||
"import base64, hashlib, pathlib, shutil, sys",
|
||||
"payload_path, tmp_dir, expected_count_text = sys.argv[1:4]",
|
||||
"expected_count = int(expected_count_text)",
|
||||
"records = [line for line in pathlib.Path(payload_path).read_text(encoding='utf-8').splitlines() if line.strip()]",
|
||||
"if len(records) != expected_count:",
|
||||
" raise SystemExit(f'bulk replacement record count mismatch: expected={expected_count} actual={len(records)}')",
|
||||
"plans = []",
|
||||
"for index, record in enumerate(records, start=1):",
|
||||
" fields = record.split()",
|
||||
" if len(fields) != 6:",
|
||||
" raise SystemExit('bulk replacement malformed record')",
|
||||
" path_b64, original_bytes_text, original_sha256, final_bytes_text, final_sha256, replacements_text = fields",
|
||||
" target = base64.b64decode(path_b64).decode('utf-8')",
|
||||
" original_bytes = int(original_bytes_text)",
|
||||
" final_bytes = int(final_bytes_text)",
|
||||
" source = pathlib.Path(target).read_bytes()",
|
||||
" if len(source) != original_bytes or hashlib.sha256(source).hexdigest() != original_sha256:",
|
||||
" raise SystemExit(f'bulk replacement original integrity mismatch for {target}')",
|
||||
" lines = source.decode('utf-8').split('\\n')",
|
||||
" if lines and lines[-1] == '':",
|
||||
" lines.pop()",
|
||||
" replacements = []",
|
||||
" for item in replacements_text.split(';'):",
|
||||
" if not item:",
|
||||
" continue",
|
||||
" start_text, old_len_text, new_b64 = item.split(',', 2)",
|
||||
" new_text = base64.b64decode(new_b64).decode('utf-8')",
|
||||
" new_lines = new_text.split('\\n')",
|
||||
" if new_lines and new_lines[-1] == '':",
|
||||
" new_lines.pop()",
|
||||
" replacements.append((int(start_text), int(old_len_text), new_lines))",
|
||||
" for start, old_len, new_lines in sorted(replacements, reverse=True):",
|
||||
" if start < 0 or old_len < 0 or start + old_len > len(lines):",
|
||||
" raise SystemExit(f'bulk replacement out of bounds for {target}')",
|
||||
" lines[start:start + old_len] = new_lines",
|
||||
" output = (('\\n'.join(lines) + '\\n') if lines else '').encode('utf-8')",
|
||||
" if len(output) != final_bytes or hashlib.sha256(output).hexdigest() != final_sha256:",
|
||||
" raise SystemExit(f'bulk replacement final integrity mismatch for {target}')",
|
||||
" tmp_path = pathlib.Path(tmp_dir) / f'file-{index}.tmp'",
|
||||
" tmp_path.write_bytes(output)",
|
||||
" plans.append((target, final_sha256, tmp_path))",
|
||||
"for target, final_sha256, tmp_path in plans:",
|
||||
" target_path = pathlib.Path(target)",
|
||||
" if str(target_path.parent) not in ('', '.'):",
|
||||
" target_path.parent.mkdir(parents=True, exist_ok=True)",
|
||||
" shutil.move(str(tmp_path), str(target_path))",
|
||||
" if hashlib.sha256(target_path.read_bytes()).hexdigest() != final_sha256:",
|
||||
" raise SystemExit(f'bulk replacement final sha256 mismatch for {target}')",
|
||||
"PY",
|
||||
" rm -f -- \"$tmp_payload\"",
|
||||
" rm -rf -- \"$tmp_dir\"",
|
||||
" trap - EXIT HUP INT TERM",
|
||||
" ;;",
|
||||
" write-b64-argv)",
|
||||
" target=$1",
|
||||
" expected_bytes=$2",
|
||||
|
||||
Reference in New Issue
Block a user