Skip to content

Commit c9f288c

Browse files
steipetesteipete
committed
perf: extract memory atomic reindex helpers
1 parent 6b6c95b commit c9f288c

3 files changed

Lines changed: 156 additions & 171 deletions

File tree

‎extensions/memory-core/src/memory/manager-atomic-reindex.ts‎

Lines changed: 49 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,49 @@
1+
import { randomUUID } from "node:crypto";
2+
import fs from "node:fs/promises";
3+
4+
export async function moveMemoryIndexFiles(sourceBase: string, targetBase: string): Promise<void> {
5+
const suffixes = ["", "-wal", "-shm"];
6+
for (const suffix of suffixes) {
7+
const source = `${sourceBase}${suffix}`;
8+
const target = `${targetBase}${suffix}`;
9+
try {
10+
await fs.rename(source, target);
11+
} catch (err) {
12+
if ((err as NodeJS.ErrnoException).code !== "ENOENT") {
13+
throw err;
14+
}
15+
}
16+
}
17+
}
18+
19+
export async function removeMemoryIndexFiles(basePath: string): Promise<void> {
20+
const suffixes = ["", "-wal", "-shm"];
21+
await Promise.all(suffixes.map((suffix) => fs.rm(`${basePath}${suffix}`, { force: true })));
22+
}
23+
24+
export async function swapMemoryIndexFiles(targetPath: string, tempPath: string): Promise<void> {
25+
const backupPath = `${targetPath}.backup-${randomUUID()}`;
26+
await moveMemoryIndexFiles(targetPath, backupPath);
27+
try {
28+
await moveMemoryIndexFiles(tempPath, targetPath);
29+
} catch (err) {
30+
await moveMemoryIndexFiles(backupPath, targetPath);
31+
throw err;
32+
}
33+
await removeMemoryIndexFiles(backupPath);
34+
}
35+
36+
export async function runMemoryAtomicReindex<T>(params: {
37+
targetPath: string;
38+
tempPath: string;
39+
build: () => Promise<T>;
40+
}): Promise<T> {
41+
try {
42+
const result = await params.build();
43+
await swapMemoryIndexFiles(params.targetPath, params.tempPath);
44+
return result;
45+
} catch (err) {
46+
await removeMemoryIndexFiles(params.tempPath);
47+
throw err;
48+
}
49+
}

‎extensions/memory-core/src/memory/manager-sync-ops.ts‎

Lines changed: 54 additions & 84 deletions
Original file line numberDiff line numberDiff line change
@@ -42,6 +42,7 @@ import {
4242
type EmbeddingProviderId,
4343
type EmbeddingProviderRuntime,
4444
} from "./embeddings.js";
45+
import { runMemoryAtomicReindex } from "./manager-atomic-reindex.js";
4546
import { openMemoryDatabaseAtPath } from "./manager-db.js";
4647
import {
4748
applyMemoryFallbackProviderState,
@@ -313,38 +314,6 @@ export abstract class MemoryManagerSyncOps {
313314
}
314315
}
315316

316-
private async swapIndexFiles(targetPath: string, tempPath: string): Promise<void> {
317-
const backupPath = `${targetPath}.backup-${randomUUID()}`;
318-
await this.moveIndexFiles(targetPath, backupPath);
319-
try {
320-
await this.moveIndexFiles(tempPath, targetPath);
321-
} catch (err) {
322-
await this.moveIndexFiles(backupPath, targetPath);
323-
throw err;
324-
}
325-
await this.removeIndexFiles(backupPath);
326-
}
327-
328-
private async moveIndexFiles(sourceBase: string, targetBase: string): Promise<void> {
329-
const suffixes = ["", "-wal", "-shm"];
330-
for (const suffix of suffixes) {
331-
const source = `${sourceBase}${suffix}`;
332-
const target = `${targetBase}${suffix}`;
333-
try {
334-
await fs.rename(source, target);
335-
} catch (err) {
336-
if ((err as NodeJS.ErrnoException).code !== "ENOENT") {
337-
throw err;
338-
}
339-
}
340-
}
341-
}
342-
343-
private async removeIndexFiles(basePath: string): Promise<void> {
344-
const suffixes = ["", "-wal", "-shm"];
345-
await Promise.all(suffixes.map((suffix) => fs.rm(`${basePath}${suffix}`, { force: true })));
346-
}
347-
348317
protected ensureSchema() {
349318
const result = ensureMemoryIndexSchema({
350319
db: this.db,
@@ -1159,62 +1128,64 @@ export abstract class MemoryManagerSyncOps {
11591128
let nextMeta: MemoryIndexMeta | null = null;
11601129

11611130
try {
1162-
this.seedEmbeddingCache(originalDb);
1163-
const shouldSyncMemory = this.sources.has("memory");
1164-
const shouldSyncSessions = this.shouldSyncSessions(
1165-
{ reason: params.reason, force: params.force },
1166-
true,
1167-
);
1168-
1169-
if (shouldSyncMemory) {
1170-
await this.syncMemoryFiles({ needsFullReindex: true, progress: params.progress });
1171-
this.dirty = false;
1172-
}
1173-
1174-
if (shouldSyncSessions) {
1175-
await this.syncSessionFiles({ needsFullReindex: true, progress: params.progress });
1176-
this.sessionsDirty = false;
1177-
this.sessionsDirtyFiles.clear();
1178-
} else if (this.sessionsDirtyFiles.size > 0) {
1179-
this.sessionsDirty = true;
1180-
} else {
1181-
this.sessionsDirty = false;
1182-
}
1131+
nextMeta = await runMemoryAtomicReindex({
1132+
targetPath: dbPath,
1133+
tempPath: tempDbPath,
1134+
build: async () => {
1135+
this.seedEmbeddingCache(originalDb);
1136+
const shouldSyncMemory = this.sources.has("memory");
1137+
const shouldSyncSessions = this.shouldSyncSessions(
1138+
{ reason: params.reason, force: params.force },
1139+
true,
1140+
);
11831141

1184-
nextMeta = {
1185-
model: this.provider?.model ?? "fts-only",
1186-
provider: this.provider?.id ?? "none",
1187-
providerKey: this.providerKey!,
1188-
sources: resolveConfiguredSourcesForMeta(this.sources),
1189-
scopeHash: resolveConfiguredScopeHash({
1190-
workspaceDir: this.workspaceDir,
1191-
extraPaths: this.settings.extraPaths,
1192-
multimodal: {
1193-
enabled: this.settings.multimodal.enabled,
1194-
modalities: this.settings.multimodal.modalities,
1195-
maxFileBytes: this.settings.multimodal.maxFileBytes,
1196-
},
1197-
}),
1198-
chunkTokens: this.settings.chunking.tokens,
1199-
chunkOverlap: this.settings.chunking.overlap,
1200-
ftsTokenizer: this.settings.store.fts.tokenizer,
1201-
};
1202-
if (!nextMeta) {
1203-
throw new Error("Failed to compute memory index metadata for reindexing.");
1204-
}
1142+
if (shouldSyncMemory) {
1143+
await this.syncMemoryFiles({ needsFullReindex: true, progress: params.progress });
1144+
this.dirty = false;
1145+
}
12051146

1206-
if (this.vector.available && this.vector.dims) {
1207-
nextMeta.vectorDims = this.vector.dims;
1208-
}
1147+
if (shouldSyncSessions) {
1148+
await this.syncSessionFiles({ needsFullReindex: true, progress: params.progress });
1149+
this.sessionsDirty = false;
1150+
this.sessionsDirtyFiles.clear();
1151+
} else if (this.sessionsDirtyFiles.size > 0) {
1152+
this.sessionsDirty = true;
1153+
} else {
1154+
this.sessionsDirty = false;
1155+
}
12091156

1210-
this.writeMeta(nextMeta);
1211-
this.pruneEmbeddingCacheIfNeeded?.();
1157+
const meta: MemoryIndexMeta = {
1158+
model: this.provider?.model ?? "fts-only",
1159+
provider: this.provider?.id ?? "none",
1160+
providerKey: this.providerKey!,
1161+
sources: resolveConfiguredSourcesForMeta(this.sources),
1162+
scopeHash: resolveConfiguredScopeHash({
1163+
workspaceDir: this.workspaceDir,
1164+
extraPaths: this.settings.extraPaths,
1165+
multimodal: {
1166+
enabled: this.settings.multimodal.enabled,
1167+
modalities: this.settings.multimodal.modalities,
1168+
maxFileBytes: this.settings.multimodal.maxFileBytes,
1169+
},
1170+
}),
1171+
chunkTokens: this.settings.chunking.tokens,
1172+
chunkOverlap: this.settings.chunking.overlap,
1173+
ftsTokenizer: this.settings.store.fts.tokenizer,
1174+
};
1175+
1176+
if (this.vector.available && this.vector.dims) {
1177+
meta.vectorDims = this.vector.dims;
1178+
}
12121179

1213-
this.db.close();
1214-
originalDb.close();
1215-
originalDbClosed = true;
1180+
this.writeMeta(meta);
1181+
this.pruneEmbeddingCacheIfNeeded?.();
12161182

1217-
await this.swapIndexFiles(dbPath, tempDbPath);
1183+
this.db.close();
1184+
originalDb.close();
1185+
originalDbClosed = true;
1186+
return meta;
1187+
},
1188+
});
12181189

12191190
this.db = openMemoryDatabaseAtPath(dbPath, this.settings.store.vector.enabled);
12201191
this.vectorReady = null;
@@ -1226,7 +1197,6 @@ export abstract class MemoryManagerSyncOps {
12261197
try {
12271198
this.db.close();
12281199
} catch {}
1229-
await this.removeIndexFiles(tempDbPath);
12301200
restoreOriginalState();
12311201
throw err;
12321202
}

‎extensions/memory-core/src/memory/manager.atomic-reindex.test.ts‎

Lines changed: 53 additions & 87 deletions
Original file line numberDiff line numberDiff line change
@@ -1,117 +1,83 @@
11
import fs from "node:fs/promises";
22
import os from "node:os";
33
import path from "node:path";
4-
import type { OpenClawConfig } from "openclaw/plugin-sdk/memory-core-host-engine-foundation";
5-
import { afterAll, afterEach, beforeAll, beforeEach, describe, expect, it, vi } from "vitest";
6-
import type { MemoryIndexManager } from "./index.js";
7-
8-
let shouldFail = false;
9-
10-
type EmbeddingTestMocksModule = typeof import("./embedding.test-mocks.js");
11-
type TestManagerHelpersModule = typeof import("./test-manager-helpers.js");
12-
type MemoryIndexModule = typeof import("./index.js");
13-
type MemoryEmbeddingProvidersModule =
14-
typeof import("../../../../src/plugins/memory-embedding-providers.js");
4+
import { DatabaseSync } from "node:sqlite";
5+
import { afterAll, beforeAll, beforeEach, describe, expect, it } from "vitest";
6+
import { runMemoryAtomicReindex } from "./manager-atomic-reindex.js";
157

168
describe("memory manager atomic reindex", () => {
179
let fixtureRoot = "";
1810
let caseId = 0;
19-
let workspaceDir: string;
2011
let indexPath: string;
21-
let manager: MemoryIndexManager | null = null;
22-
let embedBatch: ReturnType<EmbeddingTestMocksModule["getEmbedBatchMock"]>;
23-
let resetEmbeddingMocks: EmbeddingTestMocksModule["resetEmbeddingMocks"];
24-
let getRequiredMemoryIndexManager: TestManagerHelpersModule["getRequiredMemoryIndexManager"];
25-
let closeAllMemorySearchManagers: MemoryIndexModule["closeAllMemorySearchManagers"];
26-
let clearRegistry: MemoryEmbeddingProvidersModule["clearMemoryEmbeddingProviders"];
27-
let registerAdapter: MemoryEmbeddingProvidersModule["registerMemoryEmbeddingProvider"];
12+
let tempIndexPath: string;
2813

2914
beforeAll(async () => {
30-
vi.resetModules();
31-
const embeddingMocks = await import("./embedding.test-mocks.js");
32-
embedBatch = embeddingMocks.getEmbedBatchMock();
33-
resetEmbeddingMocks = embeddingMocks.resetEmbeddingMocks;
34-
({ getRequiredMemoryIndexManager } = await import("./test-manager-helpers.js"));
35-
({ closeAllMemorySearchManagers } = await import("./index.js"));
36-
({
37-
clearMemoryEmbeddingProviders: clearRegistry,
38-
registerMemoryEmbeddingProvider: registerAdapter,
39-
} = await import("../../../../src/plugins/memory-embedding-providers.js"));
4015
fixtureRoot = await fs.mkdtemp(path.join(os.tmpdir(), "openclaw-mem-atomic-"));
4116
});
4217

4318
beforeEach(async () => {
44-
vi.stubEnv("OPENCLAW_TEST_MEMORY_UNSAFE_REINDEX", "0");
45-
clearRegistry();
46-
registerAdapter({
47-
id: "openai",
48-
defaultModel: "mock-embed",
49-
transport: "remote",
50-
create: async () => ({ provider: null }),
51-
});
52-
resetEmbeddingMocks();
53-
shouldFail = false;
54-
embedBatch.mockImplementation(async (texts: string[]) => {
55-
if (shouldFail) {
56-
throw new Error("embedding failure");
57-
}
58-
return texts.map((_, index) => [index + 1, 0, 0]);
59-
});
60-
workspaceDir = path.join(fixtureRoot, `case-${caseId++}`);
19+
const workspaceDir = path.join(fixtureRoot, `case-${caseId++}`);
6120
await fs.mkdir(workspaceDir, { recursive: true });
6221
indexPath = path.join(workspaceDir, "index.sqlite");
63-
await fs.mkdir(path.join(workspaceDir, "memory"));
64-
await fs.writeFile(path.join(workspaceDir, "MEMORY.md"), "Hello memory.");
65-
});
66-
67-
afterEach(async () => {
68-
if (manager) {
69-
await manager.close();
70-
manager = null;
71-
}
72-
await closeAllMemorySearchManagers();
73-
clearRegistry();
74-
vi.unstubAllEnvs();
22+
tempIndexPath = `${indexPath}.tmp`;
7523
});
7624

7725
afterAll(async () => {
78-
if (!fixtureRoot) {
79-
vi.resetModules();
80-
return;
81-
}
8226
await fs.rm(fixtureRoot, { recursive: true, force: true });
83-
vi.resetModules();
8427
});
8528

8629
it("keeps the prior index when a full reindex fails", async () => {
87-
const cfg = {
88-
agents: {
89-
defaults: {
90-
workspace: workspaceDir,
91-
memorySearch: {
92-
provider: "openai",
93-
model: "mock-embed",
94-
store: { path: indexPath, vector: { enabled: false } },
95-
cache: { enabled: false },
96-
// Perf: keep test indexes to a single chunk to reduce sqlite work.
97-
chunking: { tokens: 4000, overlap: 0 },
98-
sync: { watch: false, onSessionStart: false, onSearch: false },
99-
},
30+
writeChunkMarker(indexPath, "before");
31+
writeChunkMarker(tempIndexPath, "after");
32+
33+
await expect(
34+
runMemoryAtomicReindex({
35+
targetPath: indexPath,
36+
tempPath: tempIndexPath,
37+
build: async () => {
38+
throw new Error("embedding failure");
10039
},
101-
list: [{ id: "main", default: true }],
102-
},
103-
} as OpenClawConfig;
40+
}),
41+
).rejects.toThrow("embedding failure");
10442

105-
manager = await getRequiredMemoryIndexManager({ cfg, agentId: "main" });
43+
expect(readChunkMarker(indexPath)).toBe("before");
44+
await expect(fs.access(tempIndexPath)).rejects.toThrow();
45+
});
10646

107-
await manager.sync({ force: true });
108-
const beforeStatus = manager.status();
109-
expect(beforeStatus.chunks).toBeGreaterThan(0);
47+
it("replaces the old index after a successful temp reindex", async () => {
48+
writeChunkMarker(indexPath, "before");
49+
writeChunkMarker(tempIndexPath, "after");
11050

111-
shouldFail = true;
112-
await expect(manager.sync({ force: true })).rejects.toThrow("embedding failure");
51+
await runMemoryAtomicReindex({
52+
targetPath: indexPath,
53+
tempPath: tempIndexPath,
54+
build: async () => undefined,
55+
});
11356

114-
const afterStatus = manager.status();
115-
expect(afterStatus.chunks).toBeGreaterThan(0);
57+
expect(readChunkMarker(indexPath)).toBe("after");
58+
await expect(fs.access(tempIndexPath)).rejects.toThrow();
11659
});
11760
});
61+
62+
function writeChunkMarker(dbPath: string, marker: string): void {
63+
const db = new DatabaseSync(dbPath);
64+
try {
65+
db.exec("CREATE TABLE chunks (id TEXT PRIMARY KEY, text TEXT NOT NULL)");
66+
db.prepare("INSERT INTO chunks (id, text) VALUES (?, ?)").run("chunk-1", marker);
67+
} finally {
68+
db.close();
69+
}
70+
}
71+
72+
function readChunkMarker(dbPath: string): string | undefined {
73+
const db = new DatabaseSync(dbPath);
74+
try {
75+
return (
76+
db.prepare("SELECT text FROM chunks WHERE id = ?").get("chunk-1") as
77+
| { text: string }
78+
| undefined
79+
)?.text;
80+
} finally {
81+
db.close();
82+
}
83+
}

0 commit comments

Comments
 (0)