diff --git a/src/backend/services/fs/FSService.test.ts b/src/backend/services/fs/FSService.test.ts index 11d60dff0..b8b251587 100644 --- a/src/backend/services/fs/FSService.test.ts +++ b/src/backend/services/fs/FSService.test.ts @@ -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): Promise => { */ const expirePendingSession = async (sessionId: string): Promise => { const key = toPendingUploadSessionKey(sessionId); - const { res: stored } = await server.stores.kv.get({ key }); - await server.stores.kv.batchPut({ - items: [ - { - key, - value: { - ...(stored as Record), - 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), + 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) => diff --git a/src/backend/stores/fs/FSEntryStore.test.ts b/src/backend/stores/fs/FSEntryStore.test.ts index f821a2fbd..6accc1c32 100644 --- a/src/backend/stores/fs/FSEntryStore.test.ts +++ b/src/backend/stores/fs/FSEntryStore.test.ts @@ -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, + ) => { + 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', () => { diff --git a/src/backend/stores/fs/FSEntryStore.ts b/src/backend/stores/fs/FSEntryStore.ts index c31c7ec3a..46870a2bf 100644 --- a/src/backend/stores/fs/FSEntryStore.ts +++ b/src/backend/stores/fs/FSEntryStore.ts @@ -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 => { + // 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 { - // 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( diff --git a/src/backend/stores/fs/pendingUploadSessionHelpers.test.ts b/src/backend/stores/fs/pendingUploadSessionHelpers.test.ts index 86cc26c73..172356127 100644 --- a/src/backend/stores/fs/pendingUploadSessionHelpers.test.ts +++ b/src/backend/stores/fs/pendingUploadSessionHelpers.test.ts @@ -17,18 +17,22 @@ * along with this program. If not, see . */ +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); diff --git a/src/backend/stores/fs/pendingUploadSessionHelpers.ts b/src/backend/stores/fs/pendingUploadSessionHelpers.ts index c3776c054..0d1aa8b8d 100644 --- a/src/backend/stores/fs/pendingUploadSessionHelpers.ts +++ b/src/backend/stores/fs/pendingUploadSessionHelpers.ts @@ -17,6 +17,7 @@ * along with this program. If not, see . */ +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( + values: T[], + sessionIdOf: (value: T) => string, +): Array<[partition: string, values: T[]]> { + const groups = new Map(); + 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 { diff --git a/src/backend/stores/systemKv/SystemKVStore.readCache.test.ts b/src/backend/stores/systemKv/SystemKVStore.readCache.test.ts index bf3bab103..b4f79c579 100644 --- a/src/backend/stores/systemKv/SystemKVStore.readCache.test.ts +++ b/src/backend/stores/systemKv/SystemKVStore.readCache.test.ts @@ -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); diff --git a/src/backend/stores/systemKv/SystemKVStore.test.ts b/src/backend/stores/systemKv/SystemKVStore.test.ts index 7693e0ea3..fbf9f9bac 100644 --- a/src/backend/stores/systemKv/SystemKVStore.test.ts +++ b/src/backend/stores/systemKv/SystemKVStore.test.ts @@ -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(); + }); + }); }); diff --git a/src/backend/stores/systemKv/SystemKVStore.ts b/src/backend/stores/systemKv/SystemKVStore.ts index 60855e871..b7a0db01d 100644 --- a/src/backend/stores/systemKv/SystemKVStore.ts +++ b/src/backend/stores/systemKv/SystemKVStore.ts @@ -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 { @@ -140,6 +147,8 @@ export interface RecursiveRecord { /** 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); } /** diff --git a/src/backend/stores/systemKv/kvMutationEvents.test.ts b/src/backend/stores/systemKv/kvMutationEvents.test.ts index 3faeee365..6c5751305 100644 --- a/src/backend/stores/systemKv/kvMutationEvents.test.ts +++ b/src/backend/stores/systemKv/kvMutationEvents.test.ts @@ -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 },