fix(kv): spread upload sessions across system KV partitions

Upload sessions shared the system namespace's single partition key. They
now go to one of 16 system partitions hashed from the session id, via a
new system-actor-only `systemPartition` KV option; partitions are never
cached and emit no mutation events. Reads fall back to the unpartitioned
key for sessions written before this (max session lifetime is one hour).
This commit is contained in:
Daniel Salazar committed 2026-10-09 18:18:02 -07:00
1 parent 43094a84a0
commit ae1bad89db
9 files changed
+466 -53

No files matched your search

+22 -13
View File
@@ -36,7 +36,10 @@ import { HttpError } from '../../core/http/HttpError.js';
import { appDataPermission } from '../permission/appDataScopes.js';
import { PuterServer } from '../../server.js';
import type { FSEntry } from '../../stores/fs/FSEntry.js';
import { toPendingUploadSessionKey } from '../../stores/fs/pendingUploadSessionHelpers.js';
import {
toPendingUploadSessionKey,
toPendingUploadSessionPartition,
} from '../../stores/fs/pendingUploadSessionHelpers.js';
import { setupTestServer } from '../../testUtil.js';
import type { IConfig } from '../../types.js';
import { generateDefaultFsentries } from '../../util/userProvisioning.js';
@@ -148,19 +151,25 @@ const caught = async (run: () => Promise<unknown>): Promise<HttpError> => {
*/
const expirePendingSession = async (sessionId: string): Promise<void> => {
const key = toPendingUploadSessionKey(sessionId);
const { res: stored } = await server.stores.kv.get({ key });
await server.stores.kv.batchPut({
items: [
{
key,
value: {
...(stored as Record<string, unknown>),
expiresAt: Date.now() - 1000,
const kvOpts = {
systemPartition: toPendingUploadSessionPartition(sessionId),
};
const { res: stored } = await server.stores.kv.get({ key }, kvOpts);
await server.stores.kv.batchPut(
{
items: [
{
key,
value: {
...(stored as Record<string, unknown>),
expiresAt: Date.now() - 1000,
},
expireAt: Math.ceil((Date.now() + 60 * 60 * 1000) / 1000),
},
expireAt: Math.ceil((Date.now() + 60 * 60 * 1000) / 1000),
},
],
});
],
},
kvOpts,
);
};
const entryAt = (user: TestUser, path: string) =>
+123
View File
@@ -28,6 +28,12 @@ import { decodeCursor, encodeCursor } from '../../util/pagination.js';
import { generateDefaultFsentries } from '../../util/userProvisioning.js';
import type { FSEntry, FSEntryCreateInput } from './FSEntry.js';
import { FSEntryStore } from './FSEntryStore.js';
import {
toPendingUploadSession,
toPendingUploadSessionExpiresAtSeconds,
toPendingUploadSessionKey,
toPendingUploadSessionPartition,
} from './pendingUploadSessionHelpers.js';
describe('FSEntryStore', () => {
it('quotes camelCase storage allowance aliases for Postgres', async () => {
@@ -1847,6 +1853,123 @@ describe('FSEntryStore pending upload sessions', () => {
).resolves.toMatchObject({ status: 'completed' });
}
});
describe('partitioning', () => {
const readStored = async (sessionId: string, partitioned: boolean) => {
const { res } = await server.stores.kv.get(
{ key: toPendingUploadSessionKey(sessionId) },
partitioned
? {
systemPartition:
toPendingUploadSessionPartition(sessionId),
}
: undefined,
);
return res;
};
/** Write a session where it lived before sessions were partitioned. */
const seedUnpartitioned = async (
input: ReturnType<typeof makeSession>,
) => {
await server.stores.kv.batchPut({
items: [
{
key: toPendingUploadSessionKey(input.sessionId),
value: toPendingUploadSession(input, Date.now()),
expireAt: toPendingUploadSessionExpiresAtSeconds(
input.expiresAt,
),
},
],
});
};
it('writes each session to its own partition, not the shared system namespace', async () => {
const user = await makeUser();
const sessions = Array.from({ length: 40 }, () =>
makeSession(user),
);
await store.batchCreatePendingEntries(sessions);
const partitions = new Set(
sessions.map((session) =>
toPendingUploadSessionPartition(session.sessionId),
),
);
expect(partitions.size).toBeGreaterThan(1);
for (const session of sessions) {
await expect(
readStored(session.sessionId, true),
).resolves.toMatchObject({ sessionId: session.sessionId });
await expect(
readStored(session.sessionId, false),
).resolves.toBeNull();
}
});
it('writes status changes to the partition the session was created in', async () => {
const user = await makeUser();
const session = makeSession(user);
await store.createPendingEntry(session);
await store.markPendingEntryFailed(session.sessionId, 'broke');
await expect(
readStored(session.sessionId, true),
).resolves.toMatchObject({
status: 'failed',
failureReason: 'broke',
});
await expect(
readStored(session.sessionId, false),
).resolves.toBeNull();
});
it('still reads a session written before partitioning', async () => {
const user = await makeUser();
const legacy = makeSession(user);
const current = makeSession(user);
await seedUnpartitioned(legacy);
await store.createPendingEntry(current);
await expect(
store.getPendingEntryBySessionId(legacy.sessionId),
).resolves.toMatchObject({
sessionId: legacy.sessionId,
status: 'pending',
});
const read = await store.getPendingEntriesBySessionIds([
legacy.sessionId,
'missing',
current.sessionId,
]);
expect(read.map((session) => session?.sessionId ?? null)).toEqual([
legacy.sessionId,
null,
current.sessionId,
]);
});
it('moves an unpartitioned session to its partition on its next update', async () => {
const user = await makeUser();
const legacy = makeSession(user);
await seedUnpartitioned(legacy);
await store.markPendingEntryCompleted(legacy.sessionId);
await expect(
readStored(legacy.sessionId, true),
).resolves.toMatchObject({ status: 'completed' });
// The partition wins over the stale copy left behind.
await expect(
readStored(legacy.sessionId, false),
).resolves.toMatchObject({ status: 'pending' });
await expect(
store.getPendingEntryBySessionId(legacy.sessionId),
).resolves.toMatchObject({ status: 'completed' });
});
});
});
describe('FSEntryStore storage allowance', () => {
+76 -36
View File
@@ -39,6 +39,7 @@ import {
PendingUploadSession,
} from './FSEntry.js';
import {
groupByPendingUploadSessionPartition,
normalizePendingUploadSession,
PendingUploadSessionStatus,
toPendingUploadSession,
@@ -459,15 +460,29 @@ export class FSEntryStore extends PuterStore {
}
try {
await this.stores.kv.batchPut({
items: sessions.map((session) => ({
key: toPendingUploadSessionKey(session.sessionId),
value: session,
expireAt: toPendingUploadSessionExpiresAtSeconds(
session.expiresAt,
await runWithConcurrencyLimit(
groupByPendingUploadSessionPartition(
sessions,
(session) => session.sessionId,
),
DEFAULT_DB_CHUNK_CONCURRENCY,
([systemPartition, group]) =>
this.stores.kv.batchPut(
{
items: group.map((session) => ({
key: toPendingUploadSessionKey(
session.sessionId,
),
value: session,
expireAt:
toPendingUploadSessionExpiresAtSeconds(
session.expiresAt,
),
})),
},
{ systemPartition },
),
})),
});
);
} catch (error) {
if (error instanceof Error) {
throw error;
@@ -485,29 +500,57 @@ export class FSEntryStore extends PuterStore {
return sessionsById;
}
const { res: rawValues } = await this.stores.kv.get({
key: uniqueSessionIds.map((sessionId) =>
toPendingUploadSessionKey(sessionId),
),
});
if (!Array.isArray(rawValues)) {
return sessionsById;
}
for (let index = 0; index < uniqueSessionIds.length; index++) {
const sessionId = uniqueSessionIds[index];
const rawValue = rawValues[index];
if (!sessionId) {
continue;
}
const normalizedSession = normalizePendingUploadSession(
rawValue,
sessionId,
const readSessions = async (
ids: string[],
systemPartition?: string,
): Promise<void> => {
// SystemKVStore returns `{ res, usage }`; only `res` is the value.
const { res: rawValues } = await this.stores.kv.get(
{
key: ids.map((sessionId) =>
toPendingUploadSessionKey(sessionId),
),
},
{ systemPartition },
);
if (normalizedSession) {
sessionsById.set(sessionId, normalizedSession);
if (!Array.isArray(rawValues)) {
return;
}
for (let index = 0; index < ids.length; index++) {
const sessionId = ids[index];
if (!sessionId) {
continue;
}
const normalizedSession = normalizePendingUploadSession(
rawValues[index],
sessionId,
);
if (normalizedSession) {
sessionsById.set(sessionId, normalizedSession);
}
}
};
await runWithConcurrencyLimit(
groupByPendingUploadSessionPartition(
uniqueSessionIds,
(sessionId) => sessionId,
),
DEFAULT_DB_CHUNK_CONCURRENCY,
([systemPartition, ids]) => readSessions(ids, systemPartition),
);
// Sessions written before partitioning sit in the unpartitioned system
// namespace. A session lives at most an hour (the presign clamp), so
// this fallback can go once partitioned writes have been deployed
// everywhere for longer than that.
const unpartitionedIds = uniqueSessionIds.filter(
(sessionId) => !sessionsById.has(sessionId),
);
if (unpartitionedIds.length > 0) {
await readSessions(unpartitionedIds);
}
return sessionsById;
@@ -2260,13 +2303,10 @@ export class FSEntryStore extends PuterStore {
async getPendingEntryBySessionId(
sessionId: string,
): Promise<PendingUploadSession | null> {
// SystemKVStore returns `{ res, usage }`. Hand the raw value (res)
// to the normalizer — passing the envelope would trip
// `isPendingUploadSession` and silently 404 the session.
const { res } = await this.stores.kv.get({
key: toPendingUploadSessionKey(sessionId),
});
return normalizePendingUploadSession(res, sessionId);
const sessionsById = await this.#getPendingUploadSessionsBySessionIds([
sessionId,
]);
return sessionsById.get(sessionId) ?? null;
}
async getPendingEntriesBySessionIds(
@@ -17,18 +17,22 @@
* along with this program. If not, see <https://www.gnu.org/licenses/>.
*/
import { randomUUID } from 'node:crypto';
import { describe, expect, it } from 'vitest';
import type {
PendingUploadCreateInput,
PendingUploadSession,
} from './FSEntry.js';
import {
groupByPendingUploadSessionPartition,
isPendingUploadSession,
normalizePendingUploadSession,
PENDING_UPLOAD_SESSION_KEY_PREFIX,
PENDING_UPLOAD_SESSION_PARTITIONS,
toPendingUploadSession,
toPendingUploadSessionExpiresAtSeconds,
toPendingUploadSessionKey,
toPendingUploadSessionPartition,
withPendingUploadSessionStatus,
} from './pendingUploadSessionHelpers.js';
@@ -69,6 +73,60 @@ describe('toPendingUploadSessionKey', () => {
});
});
describe('toPendingUploadSessionPartition', () => {
it('maps a session id to the same partition every time', () => {
// Golden values: a change here strands every live session.
expect(toPendingUploadSessionPartition('abc')).toBe(
'upload-session-10',
);
expect(
toPendingUploadSessionPartition(
'5d4adce0-a381-4982-9c02-6e2540026238',
),
).toBe('upload-session-6');
});
it('spreads session ids across every partition', () => {
const seen = new Set(
Array.from({ length: 2000 }, () =>
toPendingUploadSessionPartition(randomUUID()),
),
);
expect(seen).toEqual(
new Set(
Array.from(
{ length: PENDING_UPLOAD_SESSION_PARTITIONS },
(_, i) => `upload-session-${i}`,
),
),
);
});
});
describe('groupByPendingUploadSessionPartition', () => {
it('groups values by their session partition, keeping order within one', () => {
const ids = Array.from({ length: 50 }, () => randomUUID());
const groups = groupByPendingUploadSessionPartition(
ids.map((sessionId) => ({ sessionId })),
(value) => value.sessionId,
);
expect(new Set(groups.map(([partition]) => partition)).size).toBe(
groups.length,
);
for (const [partition, values] of groups) {
for (const { sessionId } of values) {
expect(toPendingUploadSessionPartition(sessionId)).toBe(
partition,
);
}
const order = values.map(({ sessionId }) => ids.indexOf(sessionId));
expect(order).toEqual([...order].sort((a, b) => a - b));
}
expect(groups.flatMap(([, values]) => values)).toHaveLength(ids.length);
});
});
describe('toPendingUploadSessionExpiresAtSeconds', () => {
it('rounds milliseconds up to whole seconds', () => {
expect(toPendingUploadSessionExpiresAtSeconds(1500)).toBe(2);
@@ -17,6 +17,7 @@
* along with this program. If not, see <https://www.gnu.org/licenses/>.
*/
import murmurhash from 'murmurhash';
import { PendingUploadCreateInput, PendingUploadSession } from './FSEntry.js';
export type PendingUploadSessionStatus =
@@ -27,10 +28,36 @@ export type PendingUploadSessionStatus =
export const PENDING_UPLOAD_SESSION_KEY_PREFIX = 'prodfsv2:upload-session:';
/**
* System KV partitions sessions are spread across. Changing it moves live
* sessions to partitions their reads won't look in.
*/
export const PENDING_UPLOAD_SESSION_PARTITIONS = 16;
export function toPendingUploadSessionKey(sessionId: string): string {
return `${PENDING_UPLOAD_SESSION_KEY_PREFIX}${sessionId}`;
}
/** The system KV partition a session lives in, from its id alone. */
export function toPendingUploadSessionPartition(sessionId: string): string {
return `upload-session-${murmurhash.v3(sessionId) % PENDING_UPLOAD_SESSION_PARTITIONS}`;
}
/** `values` grouped by the partition of the session each belongs to. */
export function groupByPendingUploadSessionPartition<T>(
values: T[],
sessionIdOf: (value: T) => string,
): Array<[partition: string, values: T[]]> {
const groups = new Map<string, T[]>();
for (const value of values) {
const partition = toPendingUploadSessionPartition(sessionIdOf(value));
const group = groups.get(partition);
if (group) group.push(value);
else groups.set(partition, [value]);
}
return [...groups];
}
export function toPendingUploadSessionExpiresAtSeconds(
expiresAtMs: number,
): number {
@@ -8,7 +8,7 @@ import {
it,
vi,
} from 'vitest';
import type { Actor } from '../../core/actor.ts';
import { SYSTEM_ACTOR_UUID, type Actor } from '../../core/actor.ts';
import { PuterServer } from '../../server.ts';
import { setupTestServer } from '../../testUtil.ts';
import type { SystemKVStore } from './SystemKVStore.ts';
@@ -148,6 +148,41 @@ describe('SystemKVStore read cache', () => {
expect(get).toHaveBeenCalledTimes(2);
});
it('never caches a system partition', async () => {
const key = `sys-${Math.random().toString(36).slice(2)}`;
const partition = { systemPartition: 'cache-test' };
const partitionNamespace = `v1:${SYSTEM_ACTOR_UUID}:os-global:cache-test`;
await target.batchPut({ items: [{ key, value: 'v' }] }, partition);
const get = vi.spyOn(server.clients.dynamo, 'get');
const batchGet = vi.spyOn(server.clients.dynamo, 'batchGet');
const batchGetsHere = () =>
batchGet.mock.calls.filter((call) =>
call[0].some(
(request) =>
request.items.namespace === partitionNamespace,
),
);
await target.get({ key }, partition);
await target.get({ key: [key] }, partition);
await settle();
await target.get({ key }, partition);
const second = await target.get({ key: [key] }, partition);
expect(second.res).toEqual(['v']);
expect(
get.mock.calls.filter(
(call) =>
(call[1] as { namespace?: string })?.namespace ===
partitionNamespace,
),
).toHaveLength(2);
expect(batchGetsHere()).toHaveLength(2);
await expect(
server.clients.redis.get(kvCacheKey(partitionNamespace, key)),
).resolves.toBeNull();
});
it('leaves an oversized value uncached', async () => {
await seed('big', 'x'.repeat(40 * 1024));
await target.get({ key: 'big' }, opts);
@@ -428,6 +463,20 @@ describe('SystemKVStore read cache', () => {
);
});
it('says nothing for a write to a system partition', async () => {
const emit = vi.spyOn(server.clients.event, 'emit');
await target.batchPut(
{ items: [{ key: 'sys-key', value: 'v' }] },
{ systemPartition: 'cache-test' },
);
expect(emit).not.toHaveBeenCalledWith(
'outer.kv.cacheInvalidated',
expect.anything(),
expect.anything(),
);
});
it('applies an invalidation that arrived from another region', async () => {
await seed('k', 'v');
await target.get({ key: 'k' }, opts);
@@ -15,12 +15,13 @@ import {
incrExpressionBytes,
KV_GLOBAL_APP_KEY,
kvNamespace,
parseKvNamespace,
PathPrefixIds,
type SystemKVStore,
} from './SystemKVStore.ts';
import { PUTER_KV_STORE_TABLE_NAME } from './tableDefinition.ts';
import { PuterServer } from '../../server.ts';
import type { Actor } from '../../core/actor.ts';
import { SYSTEM_ACTOR_UUID, type Actor } from '../../core/actor.ts';
describe('incr expression sizing', () => {
const longPath = (i: number): string =>
@@ -4350,4 +4351,74 @@ describe('SystemKVStore', () => {
expect(listed.res).not.toContain('pastMarker');
});
});
describe('system partitions', () => {
const systemNamespace = `v1:${SYSTEM_ACTOR_UUID}:${KV_GLOBAL_APP_KEY}`;
// Mirrors the store's own naming for a partition.
const partitionNamespace = (partition: string) =>
`${systemNamespace}:${partition}`;
const rawItem = async (namespace: string, key: string) =>
(
await server.clients.dynamo.get(PUTER_KV_STORE_TABLE_NAME, {
namespace,
key,
})
).Item;
let key: string;
beforeEach(() => {
key = `part-${Math.random().toString(36).slice(2)}`;
});
it('files a system write under its own partition key', async () => {
await target.batchPut(
{ items: [{ key, value: 'v' }] },
{ systemPartition: 'p-1' },
);
expect((await rawItem(partitionNamespace('p-1'), key))?.value).toBe(
'v',
);
expect(await rawItem(systemNamespace, key)).toBeUndefined();
});
it('reads only from the partition it names', async () => {
await target.set({ key, value: 'v' }, { systemPartition: 'p-1' });
expect(
(await target.get({ key }, { systemPartition: 'p-1' })).res,
).toBe('v');
expect(
(await target.get({ key: [key] }, { systemPartition: 'p-1' }))
.res,
).toEqual(['v']);
expect((await target.get({ key })).res).toBeNull();
expect(
(await target.get({ key }, { systemPartition: 'p-2' })).res,
).toBeNull();
});
it('does not parse as a user namespace', () => {
expect(parseKvNamespace(partitionNamespace('p-1'))).toBeNull();
});
it('ignores a partition asked for by any other actor', async () => {
await target.set(
{ key, value: 'v' },
{ ...opts, systemPartition: 'p-1' },
);
expect(
(
await rawItem(
kvNamespace(actor.user.uuid!, KV_GLOBAL_APP_KEY),
key,
)
)?.value,
).toBe('v');
expect(
await rawItem(partitionNamespace('p-1'), key),
).toBeUndefined();
});
});
});
+20 -2
View File
@@ -129,6 +129,13 @@ export interface KVOpts {
* only after its cross-app permission check has passed.
*/
namespaceAppUuid?: string;
/**
* Files system data under its own partition of the system namespace, so a
* high-volume key family doesn't share one partition key with the rest of
* the internal state. Honored only for the system actor; reads must name
* the same partition.
*/
systemPartition?: string;
}
export interface RecursiveRecord<T> {
@@ -140,6 +147,8 @@ export interface RecursiveRecord<T> {
/** Namespace app component for an actor acting without an app. */
export const KV_GLOBAL_APP_KEY = 'os-global';
const SYSTEM_NAMESPACE = `v1:${SYSTEM_ACTOR_UUID}:${KV_GLOBAL_APP_KEY}`;
// The extra segment keeps a partition from parsing as a user's namespace.
const SYSTEM_PARTITION_PREFIX = `${SYSTEM_NAMESPACE}:`;
const MAX_KEY_BYTES = 1024;
/**
@@ -345,7 +354,11 @@ const isCrossApp = (opts?: KVOpts): boolean => Boolean(opts?.namespaceAppUuid);
// than the driver treating it as app-scoped while the store files its data
// under the shared global key.
const getNamespace = (actor: Actor, opts?: KVOpts): string => {
if (isSystemActor(actor)) return SYSTEM_NAMESPACE;
if (isSystemActor(actor)) {
return opts?.systemPartition
? `${SYSTEM_PARTITION_PREFIX}${opts.systemPartition}`
: SYSTEM_NAMESPACE;
}
const appUuid =
opts?.namespaceAppUuid ??
actor.effectiveApp?.uid ??
@@ -358,6 +371,11 @@ const getNamespace = (actor: Actor, opts?: KVOpts): string => {
export const kvNamespace = (userUuid: string, appUid: string): string =>
`v1:${userUuid}:${appUid}`;
/** The system namespace or one of its partitions. */
const isSystemNamespace = (namespace: string): boolean =>
namespace === SYSTEM_NAMESPACE ||
namespace.startsWith(SYSTEM_PARTITION_PREFIX);
/** Split a namespace back into its parts, or `null` if it is not one. */
export const parseKvNamespace = (
namespace: string,
@@ -1214,7 +1232,7 @@ export class SystemKVStore extends PuterStore {
* to opt in or out: the namespace already says which kind of data it is.
*/
#cacheable(namespace: string): boolean {
return this.#cache.enabled && namespace !== SYSTEM_NAMESPACE;
return this.#cache.enabled && !isSystemNamespace(namespace);
}
/**
@@ -290,6 +290,24 @@ describe('who is left out', () => {
expect(mutations()).toEqual([]);
});
it('says nothing for a system partition', async () => {
const partition = { systemPartition: 'events-test' };
await store.batchPut(
{
items: [
{ key: 'a', value: 1 },
{ key: 'b', value: 2 },
],
},
partition,
);
await store.set({ key: KEY, value: 1 }, partition);
await store.del({ key: 'a' }, partition);
await store.flush(partition);
expect(mutations()).toEqual([]);
});
it('says nothing for an actor with no user row to key on', async () => {
await store.set(
{ key: KEY, value: 1 },