Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
5 changes: 5 additions & 0 deletions .changeset/memory-lock-exit-cleanup.md
Original file line number Diff line number Diff line change
@@ -0,0 +1,5 @@
---
"@modelcontextprotocol/server-memory": patch
---

Remove the memory file lock the server holds when it is stopped by SIGINT or SIGTERM, or exits, so the next writer does not wait for it to go stale.
6 changes: 6 additions & 0 deletions src/memory/Dockerfile
Original file line number Diff line number Diff line change
Expand Up @@ -72,4 +72,10 @@ WORKDIR /app
# not root, and refuses to start the container under `runAsNonRoot`.
USER 1000:1000

# The server handles SIGTERM itself, removing any lock it holds before it
# exits (#5485). Node running as PID 1 would otherwise ignore `docker stop`'s
# SIGTERM, leaving it to wait out its timeout and SIGKILL the server. SIGINT
# is handled the same way; it is the stop signal everything's image uses.
STOPSIGNAL SIGINT

ENTRYPOINT ["node", "dist/index.js"]
1 change: 1 addition & 0 deletions src/memory/README.md
Original file line number Diff line number Diff line change
Expand Up @@ -267,6 +267,7 @@ Saves are atomic. Every write tool writes the whole graph to a temporary file in

Each MCP client starts its own server process, so two clients using the same `MEMORY_FILE_PATH` (or both using the default) are two processes writing one file. Every write tool takes an exclusive lock file, `<memory file>.lock`, next to the memory file while it reads, changes and saves the graph, so writes from different processes take turns and none is lost. The lock file exists only while a write is in progress, so the directory holding the memory file must be writable (it already must be, for the atomic save).

- A server stopped mid-write by `SIGINT` or `SIGTERM` (Ctrl-C, `docker stop`, a client shutting it down), or exiting normally, removes the lock it holds before it exits, so the next writer does not wait on it.
- A lock left behind by a server that crashed mid-write is removed automatically: at once if its process is known to have exited on the same machine, otherwise once it has gone 30 seconds without being refreshed (a server holding the lock refreshes it every 15 seconds while it works). Only one waiting server breaks a stale lock at a time, guarded by a momentary `<memory file>.lock.break` file.
- A write that cannot get the lock within 60 seconds fails with a tool error rather than overwriting the file.
- The lock relies on exclusive file creation, which local filesystems provide; a network filesystem that does not honour it cannot be shared safely this way.
Expand Down
225 changes: 222 additions & 3 deletions src/memory/__tests__/file-lock.test.ts
Original file line number Diff line number Diff line change
@@ -1,12 +1,20 @@
// Unit tests for withFileLock, the lock file that orders mutations from
// separate server processes sharing one graph file (#4797): mutual exclusion,
// release, breaking an abandoned lock, and giving up on a held one.
// release, breaking an abandoned lock, and giving up on a held one; and the
// exit-time release of the locks this process still holds (#5485).
import { describe, it, expect, beforeEach, afterEach, vi } from "vitest";
import { promises as fs } from "fs";
import os from "os";
import path from "path";
import { spawnSync } from "child_process";
import { KnowledgeGraphManager, withFileLock } from "../index.js";
import { spawn, spawnSync } from "child_process";
import { fileURLToPath, pathToFileURL } from "url";
import {
KnowledgeGraphManager,
heldLockPaths,
installLockCleanup,
releaseHeldLocks,
withFileLock,
} from "../index.js";
import { makeTempGraph } from "./helpers.js";

const fast = { staleMs: 60_000, timeoutMs: 100, retryMs: 5 };
Expand Down Expand Up @@ -401,3 +409,214 @@ describe("withFileLock", () => {
expect(await fs.readdir(dir)).toEqual(["memory.jsonl"]);
});
});

describe("releasing held locks on exit (#5485)", () => {
let dir: string;
let cleanup: () => Promise<void>;

beforeEach(async () => {
({ dir, cleanup } = await makeTempGraph());
});

afterEach(async () => {
vi.restoreAllMocks();
await cleanup();
});

async function exists(file: string): Promise<boolean> {
return fs.access(file).then(
() => true,
() => false,
);
}

it("tracks a lock only while it is held", async () => {
const lockPath = path.join(dir, "memory.jsonl.lock");
await withFileLock(lockPath, async () => {
expect(heldLockPaths()).toEqual([lockPath]);
});
expect(heldLockPaths()).toEqual([]);
});

it("removes each lock still held, and only those still ours", async () => {
const ours = path.join(dir, "ours.lock");
const taken = path.join(dir, "taken.lock");
const gone = path.join(dir, "gone.lock");
await withFileLock(ours, () =>
withFileLock(taken, () =>
withFileLock(gone, async () => {
// Another process broke and re-took one lock, and another vanished.
await fs.writeFile(taken, "someone else");
await fs.unlink(gone);

releaseHeldLocks();

expect(heldLockPaths()).toEqual([]);
expect(await exists(ours)).toBe(false);
expect(await fs.readFile(taken, "utf-8")).toBe("someone else");
}),
),
);
expect(await fs.readFile(taken, "utf-8")).toBe("someone else");
});

it("removes a lock whose token was still being written at exit", async () => {
const lockPath = path.join(dir, "memory.jsonl.lock");
const realOpen = fs.open;
let finishWrite!: () => void;
let writing!: () => void;
const writeStarted = new Promise<void>((resolve) => (writing = resolve));
vi.spyOn(fs, "open").mockImplementationOnce(async (...args) => {
const handle = await realOpen(...args);
const realWrite = handle.writeFile.bind(handle);
// Hold the token write open, with nothing written yet.
handle.writeFile = async (data) => {
writing();
await new Promise<void>((resolve) => (finishWrite = resolve));
return realWrite(data);
};
return handle;
});

const held = withFileLock(lockPath, async () => "ran");
await writeStarted;
expect(await fs.readFile(lockPath, "utf-8")).toBe("");

releaseHeldLocks();

expect(await exists(lockPath)).toBe(false);
finishWrite();
expect(await held).toBe("ran");
});

it("leaves another writer's empty replacement of a written lock in place", async () => {
const lockPath = path.join(dir, "memory.jsonl.lock");
await withFileLock(lockPath, async () => {
// Our lock was broken as stale, and a new writer has created its own
// but not yet written its token.
await fs.writeFile(lockPath, "");

releaseHeldLocks();

expect(await exists(lockPath)).toBe(true);
});
});

it("still tracks a lock while its release is in progress", async () => {
const lockPath = path.join(dir, "memory.jsonl.lock");
const realUnlink = fs.unlink;
let trackedDuringRelease: string[] = [];
await withFileLock(lockPath, async () => {
// The release's removal of the lock is the next unlink.
vi.spyOn(fs, "unlink").mockImplementationOnce(async (file) => {
trackedDuringRelease = heldLockPaths();
return realUnlink(file);
});
});
expect(trackedDuringRelease).toEqual([lockPath]);
expect(heldLockPaths()).toEqual([]);
expect(await exists(lockPath)).toBe(false);
});

it("keeps tracking a lock the same process re-took after breaking it", async () => {
const lockPath = path.join(dir, "memory.jsonl.lock");
let releaseFirst!: () => void;
const firstHolds = new Promise<void>((resolve) => (releaseFirst = resolve));
let firstStarted!: () => void;
const started = new Promise<void>((resolve) => (firstStarted = resolve));
vi.spyOn(console, "error").mockImplementation(() => {});

const first = withFileLock(
lockPath,
async () => {
firstStarted();
await firstHolds;
},
{ staleMs: 60_000 },
);
await started;
// Age the first holder's lock past a short staleMs, so a second holder in
// this process breaks it and takes it over.
const old = new Date(Date.now() - 10_000);
await fs.utimes(lockPath, old, old);
await withFileLock(
lockPath,
async () => {
releaseFirst();
await first;
// The first holder's release left the second's lock, and its entry.
expect(heldLockPaths()).toEqual([lockPath]);
expect(await exists(lockPath)).toBe(true);
},
{ staleMs: 1_000, timeoutMs: 1_000, retryMs: 5 },
);
expect(heldLockPaths()).toEqual([]);
});

it("installs an exit hook and SIGINT and SIGTERM handlers, once", () => {
const handlers = new Map<string | symbol, () => void>();
const on = vi.spyOn(process, "on").mockImplementation((event, listener) => {
handlers.set(event, listener as () => void);
return process;
});
const exit = vi
.spyOn(process, "exit")
.mockImplementation((() => undefined) as typeof process.exit);

installLockCleanup();
installLockCleanup();

expect(on).toHaveBeenCalledTimes(3);
expect(handlers.get("exit")).toBe(releaseHeldLocks);
handlers.get("SIGINT")?.();
expect(exit).toHaveBeenLastCalledWith(130);
handlers.get("SIGTERM")?.();
expect(exit).toHaveBeenLastCalledWith(143);
});

// A real signal, sent to a real process holding a lock mid-operation.
// Windows has no SIGTERM to deliver: process.kill() there ends the process
// outright, so no handler can run.
it.skipIf(process.platform === "win32")(
"removes the lock when SIGTERM stops a process mid-mutation",
async () => {
const lockPath = path.join(dir, "memory.jsonl.lock");
const distIndex = pathToFileURL(
path.join(
path.dirname(fileURLToPath(import.meta.url)),
"..",
"dist",
"index.js",
),
).href;
const script = `
const { installLockCleanup, withFileLock } = await import(${JSON.stringify(distIndex)});
installLockCleanup();
await withFileLock(${JSON.stringify(lockPath)}, () => {
console.log("held");
return new Promise(() => setInterval(() => {}, 1000));
});
`;
const child = spawn(
process.execPath,
["--input-type=module", "-e", script],
{ stdio: ["ignore", "pipe", "inherit"] },
);
const exited = new Promise<number | null>((resolve) =>
child.on("exit", (code) => resolve(code)),
);
await new Promise<void>((resolve, reject) => {
child.stdout.on("data", (data: Buffer) => {
if (data.toString().includes("held")) resolve();
});
child.on("exit", () => reject(new Error("exited before holding")));
});
expect(await exists(lockPath)).toBe(true);

child.kill("SIGTERM");

expect(await exited).toBe(143);
expect(await exists(lockPath)).toBe(false);
},
);
});
83 changes: 82 additions & 1 deletion src/memory/index.ts
Original file line number Diff line number Diff line change
Expand Up @@ -15,7 +15,7 @@ import {
import { McpServer } from "@modelcontextprotocol/server";
import type { Transport } from "@modelcontextprotocol/server";
import { z } from "zod";
import { promises as fs, realpathSync } from "fs";
import { promises as fs, readFileSync, realpathSync, unlinkSync } from "fs";
import type { FileHandle } from "fs/promises";
import path from "path";
import os from "os";
Expand Down Expand Up @@ -276,6 +276,72 @@ async function breakStaleLock(
}
}

// The lock files this process holds right now, each mapped to the token it
// writes into it and whether that write has finished. withFileLock releases its lock in a `finally`, which never
// runs when a signal ends the process, so the exit hook releaseHeldLocks()
// removes them instead (#5485). Without it, a lock left by a stopped server
// stalls the next writer until it goes stale, which takes the full staleMs
// when that writer is on another host (another container) and cannot see
// that this process has gone.
interface HeldLock {
token: string;
written: boolean;
}
const heldLocks = new Map<string, HeldLock>();

// The lock paths this process holds right now, for tests.
export function heldLockPaths(): string[] {
return [...heldLocks.keys()];
}

// Synchronously remove every lock this process still holds, as the process
// exits: an `exit` listener cannot wait on a promise. As in withFileLock's own
// release, a lock is removed only while it still holds this process's token,
// so a lock broken as stale and re-taken by another process stays in place.
// A lock is tracked from the moment its exclusive create returns, before its
// token is written. Until that write finishes, a file holding only the start
// of the token (or none of it) is ours too: the exit interrupted the write.
// Nothing else can have replaced it by then, since a lock that new, with a
// live owner, is never judged stale. Once the token is written only an exact
// match is ours, so another writer's replacement, empty while its own token
// is being written, is left alone.
export function releaseHeldLocks(): void {
for (const [lockPath, { token, written }] of heldLocks) {
try {
const contents = readFileSync(lockPath, "utf-8");
if (written ? contents === token : token.startsWith(contents)) {
unlinkSync(lockPath);
}
} catch {
// Already gone, or unreadable: there is nothing of ours to remove.
}
}
heldLocks.clear();
}

// The exit codes a shell reports for a process its signal ended: 128 + the
// signal's number.
const SIGNAL_EXIT_CODES = { SIGINT: 130, SIGTERM: 143 } as const;

let lockCleanupInstalled = false;

// Release held locks however the server process ends (#5485): on a normal
// exit or process.exit() through the `exit` event, and on SIGINT or SIGTERM,
// whose default action would end the process without emitting `exit`. A
// signal handler replaces that default, so it exits itself, with the status
// the default would have given. Installed once, by main(): only the server
// process should take over its signals, not a test that imports this module.
export function installLockCleanup(): void {
if (lockCleanupInstalled) {
return;
}
lockCleanupInstalled = true;
process.on("exit", releaseHeldLocks);
for (const [signal, code] of Object.entries(SIGNAL_EXIT_CODES)) {
process.on(signal, () => process.exit(code));
}
}

// Run operation while holding an exclusive lock file at lockPath, so that
// mutations from separate server processes sharing one graph file take turns
// and each one loads the graph the previous one saved (#4797). The lock is a
Expand Down Expand Up @@ -305,14 +371,21 @@ export async function withFileLock<T>(
}
}
if (handle) {
// Track the lock before anything else is awaited: a signal is handled
// only between tasks, never between the open's completion and this
// line, so from here on the exit hook can remove it (#5485).
const held: HeldLock = { token, written: false };
heldLocks.set(lockPath, held);
try {
await handle.writeFile(token);
held.written = true;
await handle.close();
} catch (error) {
// The lock was created but not written: remove it rather than leave
// every other writer waiting for it to go stale.
await handle.close().catch(ignoreError);
await fs.unlink(lockPath).catch(ignoreError);
heldLocks.delete(lockPath);
throw error;
}
break;
Expand Down Expand Up @@ -381,6 +454,13 @@ export async function withFileLock<T>(
if (current === token) {
await fs.unlink(lockPath).catch(ignoreError);
}
// Stop tracking the lock only once the release is done, so an exit during
// it still removes the lock. A lock this same process broke as stale and
// re-took in another withFileLock call carries that call's token, and
// stays tracked for it.
if (heldLocks.get(lockPath)?.token === token) {
heldLocks.delete(lockPath);
}
}
}

Expand Down Expand Up @@ -1435,6 +1515,7 @@ export async function main(transport?: Transport): Promise<StdioServerHandle> {
if (locationWarning) {
console.error(locationWarning);
}
installLockCleanup();
const handle = serveStdio(({ era }) => createServer(memoryFilePath, era), {
transport,
// serveStdio reports a transport that fails, including one that cannot
Expand Down
Loading