diff --git a/packages/@wispplace/tiered-storage/scripts/bench-stream-promotion.ts b/packages/@wispplace/tiered-storage/scripts/bench-stream-promotion.ts new file mode 100644 index 0000000..facf170 --- /dev/null +++ b/packages/@wispplace/tiered-storage/scripts/bench-stream-promotion.ts @@ -0,0 +1,188 @@ +/** + * Reproducible local bench for streamed-read promotion. + * + * Reads the same cold object N times through `TieredStorage.getStream` (the + * hosting read path) with hosting's tier layout: memory hot, disk warm, and a + * counting cold tier that stands in for S3. It reports cold-tier calls + * (GET = getStream/get/getWithMetadata, HEAD = getMetadata, metadata writes), + * the tier that served each read, per-read latency, and the RSS growth of one + * large streamed read. + * + * Usage: bun scripts/bench-stream-promotion.ts [--reads 20] [--cold-latency-ms 25] [--large-mib 256] + */ +import { mkdtemp, rm } from 'node:fs/promises' +import { tmpdir } from 'node:os' +import { join } from 'node:path' +import { Readable } from 'node:stream' +import { parseArgs } from 'node:util' +import { DiskStorageTier, MemoryStorageTier, type StorageMetadata, type StorageTier, TieredStorage } from '../src/index' + +const { values } = parseArgs({ + options: { + reads: { type: 'string', default: '20' }, + 'cold-latency-ms': { type: 'string', default: '25' }, + 'large-mib': { type: 'string', default: '256' }, + }, +}) +const reads = Number(values.reads) +const coldLatencyMs = Number(values['cold-latency-ms']) +const largeBytes = Number(values['large-mib']) * 1024 * 1024 +const CHUNK = 64 * 1024 + +const counts = { get: 0, head: 0, metadataWrite: 0 } + +const metadataFor = (key: string, size: number): StorageMetadata => ({ + key, + size, + createdAt: new Date(0), + lastAccessed: new Date(0), + accessCount: 0, + compressed: false, + checksum: `bench-${size}`, +}) + +/** Deterministic bytes generated on demand, so the source itself never holds an object in memory. */ +function* chunksOf(size: number): Generator { + for (let offset = 0; offset < size; offset += CHUNK) { + yield Buffer.alloc(Math.min(CHUNK, size - offset), offset / CHUNK) + } +} + +const delay = () => new Promise((resolve) => setTimeout(resolve, coldLatencyMs)) + +/** Read-only S3 stand-in: every call pays one simulated round trip and is counted. */ +function countingColdTier(objects: Map): StorageTier { + const body = async (key: string) => { + counts.get++ + await delay() + const size = objects.get(key) + return size === undefined ? null : { size, metadata: metadataFor(key, size) } + } + return { + get: async (key) => { + const found = await body(key) + return found && new Uint8Array(Buffer.concat([...chunksOf(found.size)])) + }, + getWithMetadata: async (key) => { + const found = await body(key) + return found && { data: new Uint8Array(Buffer.concat([...chunksOf(found.size)])), metadata: found.metadata } + }, + getStream: async (key) => { + const found = await body(key) + return found && { stream: Readable.from(chunksOf(found.size)), metadata: found.metadata } + }, + getMetadata: async (key) => { + counts.head++ + await delay() + const size = objects.get(key) + return size === undefined ? null : metadataFor(key, size) + }, + setMetadata: async () => { + counts.metadataWrite++ + }, + set: async () => {}, + delete: async () => {}, + exists: async (key) => objects.has(key), + listKeys: async function* () {}, + deleteMany: async () => {}, + getStats: async () => ({ bytes: 0, items: objects.size }), + clear: async () => {}, + } +} + +async function drain(stream: NodeJS.ReadableStream): Promise { + let bytes = 0 + for await (const chunk of stream) bytes += (chunk as Buffer).byteLength + return bytes +} + +const directory = await mkdtemp(join(tmpdir(), 'tiered-storage-bench-')) +const objects = new Map([ + ['did:plc:bench/site/app.css', 40 * 1024], + ['did:plc:bench/site/photo.png', 2 * 1024 * 1024], + ['did:plc:bench/site/video.mp4', largeBytes], +]) +const storage = new TieredStorage({ + tiers: { + hot: new MemoryStorageTier({ maxSizeBytes: 100 * 1024 * 1024, maxItems: 500 }), + warm: new DiskStorageTier({ directory, encodeColons: false }), + cold: countingColdTier(objects), + }, + // Hosting's placement shape: css/js may live in memory, media stays on disk. + placementRules: [ + { pattern: '**/*.{css,js}', tiers: ['hot', 'warm', 'cold'] }, + { pattern: '**', tiers: ['warm', 'cold'] }, + ], + compression: false, + promotionStrategy: 'eager', + serialization: { serialize: async (data) => data as Uint8Array, deserialize: async (data) => data }, +}) + +const settle = () => new Promise((resolve) => setTimeout(resolve, 2 * coldLatencyMs + 10)) + +async function repeatedReads(key: string) { + const before = { ...counts } + const sources: Record = {} + const latencies: number[] = [] + for (let read = 0; read < reads; read++) { + const started = performance.now() + const result = await storage.getStream(key) + if (!result) throw new Error(`missing ${key}`) + const bytes = await drain(result.stream) + latencies.push(performance.now() - started) + if (bytes !== objects.get(key)) throw new Error(`short read of ${key}: ${bytes}`) + sources[result.source] = (sources[result.source] ?? 0) + 1 + // Let fire-and-forget access-stat updates land in this read's counts. + await settle() + } + const sorted = [...latencies].sort((a, b) => a - b) + return { + key, + reads, + coldGet: counts.get - before.get, + coldHead: counts.head - before.head, + coldMetadataWrite: counts.metadataWrite - before.metadataWrite, + sources, + firstReadMs: Number(latencies[0]!.toFixed(2)), + medianLaterReadMs: Number(sorted[Math.floor(sorted.length / 2)]!.toFixed(2)), + totalMs: Number(latencies.reduce((sum, value) => sum + value, 0).toFixed(1)), + } +} + +async function largeRead(key: string) { + Bun.gc(true) + const baseline = process.memoryUsage().rss + let peak = baseline + const sampler = setInterval(() => { + peak = Math.max(peak, process.memoryUsage().rss) + }, 2) + const result = await storage.getStream(key) + if (!result) throw new Error(`missing ${key}`) + const bytes = await drain(result.stream) + clearInterval(sampler) + peak = Math.max(peak, process.memoryUsage().rss) + const mib = (value: number) => Number((value / 1024 / 1024).toFixed(1)) + return { + key, + bytes: mib(bytes), + source: result.source, + rssBaselineMiB: mib(baseline), + rssPeakGrowthMiB: mib(peak - baseline), + } +} + +try { + // The large read runs first so its RSS growth is not hidden by earlier allocations. + const large = await largeRead('did:plc:bench/site/video.mp4') + const largeAgain = await storage.getStream('did:plc:bench/site/video.mp4') + if (largeAgain) await drain(largeAgain.stream) + await settle() + const report = { + coldLatencyMs, + large: { ...large, secondReadSource: largeAgain?.source ?? null }, + repeated: [await repeatedReads('did:plc:bench/site/app.css'), await repeatedReads('did:plc:bench/site/photo.png')], + } + console.log(JSON.stringify(report, null, 2)) +} finally { + await rm(directory, { recursive: true, force: true }) +} diff --git a/packages/@wispplace/tiered-storage/src/TieredStorage.ts b/packages/@wispplace/tiered-storage/src/TieredStorage.ts index d635ada..88ef3d2 100644 --- a/packages/@wispplace/tiered-storage/src/TieredStorage.ts +++ b/packages/@wispplace/tiered-storage/src/TieredStorage.ts @@ -6,6 +6,7 @@ import type { AllTierStats, SetOptions, SetResult, + StagedTierWrite, StorageMetadata, StorageResult, StorageSnapshot, @@ -33,6 +34,9 @@ type InFlightBufferedRead = { promise: Promise } +/** Default for {@link TieredStorageConfig.streamHotPromotionMaxBytes}. */ +const DEFAULT_STREAM_HOT_PROMOTION_MAX_BYTES = 256 * 1024 + export interface UpperTierInvalidationFailure { tier: 'hot' | 'warm' reason: unknown @@ -109,6 +113,7 @@ export class TieredStorage { private deserialize: (data: Uint8Array) => Promise private keyFences = new Map() private inFlightBufferedReads = new Map() + private inFlightStreamPromotions = new Set() private upperTierPromotionEpoch = 0 private upperTierMutationQueue: Promise = Promise.resolve() @@ -224,6 +229,7 @@ export class TieredStorage { /** * Promote only if no invalidation began after the source read, and order that * promotion with deletion so a write already in progress is deleted afterward. + * Staged writes are committed under the same check and discarded when it fails. */ private async promoteIfCurrent( key: string, @@ -231,13 +237,18 @@ export class TieredStorage { generation: number, upperTierPromotionEpoch: number, promotions: Array<{ tier: StorageTier; data: Uint8Array; metadata: StorageMetadata }>, + staged: StagedTierWrite[] = [], ): Promise { await this.enqueueKeyMutation(key, fence, async () => { await this.enqueueUpperTierMutation(async () => { if (fence.generation !== generation || this.upperTierPromotionEpoch !== upperTierPromotionEpoch) { + await Promise.allSettled(staged.map((write) => write.abort())) return } - await Promise.allSettled(promotions.map(({ tier, data, metadata }) => tier.set(key, data, metadata))) + await Promise.allSettled([ + ...promotions.map(({ tier, data, metadata }) => tier.set(key, data, metadata)), + ...staged.map((write) => write.commit()), + ]) }) }) } @@ -379,8 +390,6 @@ export class TieredStorage { await this.promoteIfCurrent(key, fence, generation, upperTierPromotionEpoch, promotions) } - // Fire-and-forget access stats update (non-critical) - void this.updateAccessStats(key, 'cold') const consumerResult = cloneTierGetResult(result) return { data: (await this.deserializeData(consumerResult.data)) as T, @@ -460,8 +469,9 @@ export class TieredStorage { * The stream must be consumed or destroyed by the caller. * * Checks tiers in order: hot → warm → cold. - * On cache miss, does NOT promote data to upper tiers (streaming would - * require buffering, defeating the purpose). + * With eager promotion, a warm or cold hit is copied into upper tiers as the + * caller consumes it (see {@link promoteStream}); nothing is promoted from a + * stream that is destroyed or fails before its end. * * Decompression is automatically handled if the data was stored with * compression enabled (metadata.compressed = true). @@ -483,6 +493,7 @@ export class TieredStorage { throwIfAborted(options.signal) const fence = this.acquireReadFence(key) const generation = fence.generation + const upperTierPromotionEpoch = this.upperTierPromotionEpoch let ownsFence = false try { const tiers = [ @@ -508,9 +519,13 @@ export class TieredStorage { await this.delete(key) return null } - void this.updateAccessStats(key, name) + if (name !== 'cold') void this.updateAccessStats(key, name) + const source = + name === 'hot' + ? result.stream + : await this.promoteStream(key, name, result, fence, generation, upperTierPromotionEpoch) const owned = this.wrapStreamWithDecompression( - { stream: result.stream, metadata: cloneStorageMetadata(result.metadata) }, + { stream: source, metadata: cloneStorageMetadata(result.metadata) }, name, ) let released = false @@ -540,6 +555,95 @@ export class TieredStorage { } } + /** + * Copy a streamed warm or cold hit into upper tiers while the caller reads it. + * + * @remarks + * Stored bytes pass through unchanged. A cold hit is written to the warm + * tier's staged write, and each chunk is handed on only once that write has + * taken it, so memory stays bounded by the stream's buffers. The hot tier + * gets a copy only when placement allows hot and the object stays within + * `streamHotPromotionMaxBytes`. When the source ends, both are published + * through {@link promoteIfCurrent}, before the returned stream ends, while + * the caller's read fence is still held; a write or invalidation that began + * after this read was opened discards them. A stream that is destroyed, + * fails, or ends short of `metadata.size` promotes nothing, and its staged + * bytes are removed. Concurrent reads of a key promote once. + */ + private async promoteStream( + key: string, + source: 'warm' | 'cold', + result: { stream: NodeJS.ReadableStream; metadata: StorageMetadata }, + fence: KeyFence, + generation: number, + upperTierPromotionEpoch: number, + ): Promise { + if (this.config.promotionStrategy !== 'eager' || this.inFlightStreamPromotions.has(key)) return result.stream + const warm = source === 'cold' ? this.config.tiers.warm : undefined + const hot = this.config.tiers.hot && this.getTiersForKey(key).includes('hot') ? this.config.tiers.hot : undefined + if (!warm?.stageWrite && !hot) return result.stream + + this.inFlightStreamPromotions.add(key) + const metadata = cloneStorageMetadata(result.metadata) + const hotMaxBytes = this.config.streamHotPromotionMaxBytes ?? DEFAULT_STREAM_HOT_PROMOTION_MAX_BYTES + let staged: StagedTierWrite | undefined + try { + staged = await warm?.stageWrite?.(key, cloneStorageMetadata(metadata)) + } catch { + // Promotion is best-effort; the read itself still succeeds. + } + let hotChunks: Buffer[] | undefined = hot ? [] : undefined + let bytesRead = 0 + const discardStaged = async () => { + const discarded = staged + staged = undefined + await discarded?.abort() + } + + const tee = new Transform({ + transform: (chunk: Buffer, _encoding, callback) => { + bytesRead += chunk.byteLength + if (hotChunks && bytesRead > hotMaxBytes) hotChunks = undefined + // Copy: the caller may mutate the chunk it receives. + hotChunks?.push(Buffer.from(chunk)) + if (!staged) return callback(null, chunk) + staged.write(chunk).then( + () => callback(null, chunk), + () => { + void discardStaged() + callback(null, chunk) + }, + ) + }, + flush: (callback) => { + // A compressed object's recorded size may be its uncompressed length, + // so only uncompressed objects can be checked for a short source. + if (!metadata.compressed && bytesRead !== metadata.size) { + void discardStaged().then(() => callback()) + return + } + const promotions = + hot && hotChunks ? [{ tier: hot, data: new Uint8Array(Buffer.concat(hotChunks)), metadata }] : [] + const commits = staged ? [staged] : [] + hotChunks = undefined + this.promoteIfCurrent(key, fence, generation, upperTierPromotionEpoch, promotions, commits).then( + () => callback(), + () => callback(), + ) + }, + }) + tee.once('close', () => { + this.inFlightStreamPromotions.delete(key) + hotChunks = undefined + // A no-op after commit; otherwise removes the partial staged file. + void discardStaged() + }) + pipeline(result.stream as Readable, tee).catch(() => { + // pipeline destroys the tee with the error, which the caller observes. + }) + return tee + } + /** * Wrap a stream result with decompression if needed. */ @@ -1461,12 +1565,15 @@ export class TieredStorage { } /** - * Update access statistics for a key. + * Update access statistics for a key in a cache tier. + * + * @remarks + * The cold tier is never touched on read: it is the source of truth, and an + * S3 metadata update is a HEAD plus a full-object copy that can race a + * concurrent write and replace its metadata with the stale copy read here. */ - private async updateAccessStats(key: string, tier: 'hot' | 'warm' | 'cold'): Promise { - const tierObj = - tier === 'hot' ? this.config.tiers.hot : tier === 'warm' ? this.config.tiers.warm : this.config.tiers.cold - + private async updateAccessStats(key: string, tier: 'hot' | 'warm'): Promise { + const tierObj = tier === 'hot' ? this.config.tiers.hot : this.config.tiers.warm if (!tierObj) return const metadata = await tierObj.getMetadata(key) diff --git a/packages/@wispplace/tiered-storage/src/index.ts b/packages/@wispplace/tiered-storage/src/index.ts index 3457b84..4b3e5e7 100644 --- a/packages/@wispplace/tiered-storage/src/index.ts +++ b/packages/@wispplace/tiered-storage/src/index.ts @@ -24,6 +24,7 @@ export type { PlacementRule, SetOptions, SetResult, + StagedTierWrite, StorageMetadata, StorageResult, StorageSnapshot, diff --git a/packages/@wispplace/tiered-storage/src/tiers/DiskStorageTier.ts b/packages/@wispplace/tiered-storage/src/tiers/DiskStorageTier.ts index 94577e4..4fe638f 100644 --- a/packages/@wispplace/tiered-storage/src/tiers/DiskStorageTier.ts +++ b/packages/@wispplace/tiered-storage/src/tiers/DiskStorageTier.ts @@ -2,9 +2,21 @@ import { createWriteStream, existsSync } from 'node:fs' import { mkdir, open, readdir, readFile, rename, rm, stat, unlink, writeFile } from 'node:fs/promises' import { dirname, join, relative, resolve, sep } from 'node:path' import { pipeline } from 'node:stream/promises' -import type { StorageMetadata, StorageTier, TierGetResult, TierStats, TierStreamResult } from '../types/index.js' +import type { + StagedTierWrite, + StorageMetadata, + StorageTier, + TierGetResult, + TierStats, + TierStreamResult, +} from '../types/index.js' import { encodeKey } from '../utils/path-encoding.js' +/** Cache-root directories that hold in-progress work, never keys. */ +const INVALIDATED_DIRECTORY = '.invalidated' +const STAGING_DIRECTORY = '.staging' +const INTERNAL_DIRECTORIES = new Set([INVALIDATED_DIRECTORY, STAGING_DIRECTORY]) + function getErrnoCode(error: unknown): string | undefined { if (typeof error !== 'object' || error === null) return undefined const maybeCode = (error as { code?: unknown }).code @@ -151,7 +163,8 @@ export class DiskStorageTier implements StorageTier { */ private async initialize(): Promise { await this.ensureDirectory() - await rm(join(resolve(this.config.directory), '.invalidated'), { recursive: true, force: true }) + await rm(join(resolve(this.config.directory), INVALIDATED_DIRECTORY), { recursive: true, force: true }) + await rm(join(resolve(this.config.directory), STAGING_DIRECTORY), { recursive: true, force: true }) this.metadataIndex.clear() this.currentSize = 0 await this.rebuildIndex() @@ -200,7 +213,7 @@ export class DiskStorageTier implements StorageTier { const fullPath = join(dir, entry.name) if (entry.isDirectory()) { - if (entry.name === '.invalidated') continue + if (INTERNAL_DIRECTORIES.has(entry.name)) continue await this.rebuildIndexRecursive(fullPath) } else if (!entry.name.endsWith('.meta')) { try { @@ -425,10 +438,23 @@ export class DiskStorageTier implements StorageTier { } private async setUnlocked(key: string, data: Uint8Array, metadata: StorageMetadata): Promise { + await this.storeUnlocked(key, metadata, data.byteLength, (filePath) => writeFile(filePath, data)) + } + + /** + * Replace a key's entry with `size` bytes that `writeData` puts at its data path. + * The caller holds the mutation lock. + */ + private async storeUnlocked( + key: string, + metadata: StorageMetadata, + size: number, + writeData: (filePath: string) => Promise, + ): Promise { const filePath = this.getFilePath(key) const metaPath = this.getMetaPath(key) - if (this.exceedsSizeLimit(data.byteLength)) { + if (this.exceedsSizeLimit(size)) { // Preserve unrelated cached entries, but never serve a stale value for a // key whose replacement is too large for this tier. await this.deleteUnlocked(key) @@ -442,7 +468,7 @@ export class DiskStorageTier implements StorageTier { } if (this.config.maxSizeBytes) { - await this.evictIfNeeded(data.byteLength) + await this.evictIfNeeded(size) } // Eviction may have removed this newly empty nested directory. @@ -453,7 +479,7 @@ export class DiskStorageTier implements StorageTier { try { if (!(await this.writeMetadataAtomically(metaPath, metadata))) return - await writeFile(filePath, data) + await writeData(filePath) } catch (error) { // Do not recreate a directory that an external invalidation removed. The // entry stays absent and will be re-fetched from the source of truth. @@ -463,12 +489,61 @@ export class DiskStorageTier implements StorageTier { } this.metadataIndex.set(key, { - size: data.byteLength, + size, createdAt: metadata.createdAt, lastAccessed: metadata.lastAccessed, ...(metadata.ttl && { ttl: metadata.ttl }), }) - this.currentSize += data.byteLength + this.currentSize += size + } + + /** + * Stage a write in a private file under the cache root and publish it with + * one rename on commit, so readers never observe a partial object. + * + * @param key - The key the staged bytes will be stored under + * @param metadata - Metadata stored when the write commits + * @returns A staged write; abandoned staging files are removed on startup + */ + async stageWrite(key: string, metadata: StorageMetadata): Promise { + await this.initialization + // Reject an invalid key before any bytes are staged for it. + this.getFilePath(key) + const stagingDirectory = join(resolve(this.config.directory), STAGING_DIRECTORY) + await mkdir(stagingDirectory, { recursive: true }) + const stagedPath = join(stagingDirectory, `${process.pid}-${Date.now()}-${Math.random().toString(16).slice(2)}`) + const handle = await open(stagedPath, 'wx', 0o600) + let settled = false + const settle = async (): Promise => { + if (settled) return false + settled = true + await handle.close().catch(() => {}) + return true + } + + return { + write: async (chunk) => { + for (let offset = 0; offset < chunk.byteLength; ) { + const { bytesWritten } = await handle.write(chunk, offset) + offset += bytesWritten + } + }, + commit: async () => { + if (!(await settle())) return + try { + const { size } = await stat(stagedPath) + await this.withMutation(() => + this.storeUnlocked(key, metadata, size, (filePath) => rename(stagedPath, filePath)), + ) + } finally { + // Gone after a successful rename; removes the staged bytes otherwise. + await unlink(stagedPath).catch(() => {}) + } + }, + abort: async () => { + if (await settle()) await unlink(stagedPath).catch(() => {}) + }, + } } async delete(key: string): Promise { @@ -532,7 +607,7 @@ export class DiskStorageTier implements StorageTier { const fullPath = join(dir, entry.name) if (entry.isDirectory()) { - if (entry.name === '.invalidated') continue + if (INTERNAL_DIRECTORIES.has(entry.name)) continue // Recurse into subdirectory for await (const key of this.listKeysRecursive(fullPath, prefix)) { yield key @@ -580,7 +655,7 @@ export class DiskStorageTier implements StorageTier { const prefixDirectory = this.getFilePath(prefix) if (!existsSync(prefixDirectory)) return 0 - const invalidationDirectory = join(resolve(this.config.directory), '.invalidated') + const invalidationDirectory = join(resolve(this.config.directory), INVALIDATED_DIRECTORY) await mkdir(invalidationDirectory, { recursive: true }) const detachedDirectory = join( invalidationDirectory, @@ -695,7 +770,7 @@ export class DiskStorageTier implements StorageTier { const fullPath = join(dir, entry.name) if (entry.isDirectory()) { - if (entry.name === '.invalidated') continue + if (INTERNAL_DIRECTORIES.has(entry.name)) continue const subStats = await this.getStatsRecursive(fullPath) bytes += subStats.bytes items += subStats.items diff --git a/packages/@wispplace/tiered-storage/src/types/index.ts b/packages/@wispplace/tiered-storage/src/types/index.ts index 89fd4ac..b8f9c9c 100644 --- a/packages/@wispplace/tiered-storage/src/types/index.ts +++ b/packages/@wispplace/tiered-storage/src/types/index.ts @@ -227,6 +227,18 @@ export interface StorageTier { */ setStream?(key: string, stream: NodeJS.ReadableStream, metadata: StorageMetadata): Promise + /** + * Stage a write that stays invisible to readers until it is committed. + * + * @param key - The key the staged bytes will be stored under + * @param metadata - Metadata stored alongside the data on commit + * + * @remarks + * Lets a streamed read copy an object into this tier chunk by chunk without + * buffering it and without ever exposing a partial object. + */ + stageWrite?(key: string, metadata: StorageMetadata): Promise + /** * Store data with associated metadata. * @@ -338,6 +350,20 @@ export interface StorageTier { clear(): Promise } +/** + * An in-progress write returned by {@link StorageTier.stageWrite}. + */ +export interface StagedTierWrite { + /** Append bytes; resolves once the tier no longer needs the chunk. */ + write(chunk: Uint8Array): Promise + + /** Publish the staged bytes under the key, replacing any current value. */ + commit(): Promise + + /** Discard the staged bytes. A no-op after commit() or an earlier abort(). */ + abort(): Promise +} + /** * Rule for automatic tier placement based on key patterns. * @@ -443,6 +469,18 @@ export interface TieredStorageConfig { */ promotionStrategy?: 'eager' | 'lazy' + /** + * Largest streamed object copied into the hot tier during eager promotion. + * + * @defaultValue 262144 (256 KiB) + * + * @remarks + * A streamed read buffers at most this many bytes for the hot tier, and only + * for keys whose placement includes hot. Warm promotion streams to a staged + * file and has no such limit beyond the warm tier's own capacity. + */ + streamHotPromotionMaxBytes?: number + /** * Custom serialization/deserialization functions. * diff --git a/packages/@wispplace/tiered-storage/test/TieredStorage.stream-promotion.test.ts b/packages/@wispplace/tiered-storage/test/TieredStorage.stream-promotion.test.ts new file mode 100644 index 0000000..7d04617 --- /dev/null +++ b/packages/@wispplace/tiered-storage/test/TieredStorage.stream-promotion.test.ts @@ -0,0 +1,307 @@ +import { afterEach, beforeEach, describe, expect, test } from 'bun:test' +import { mkdtemp, readdir, rm } from 'node:fs/promises' +import { tmpdir } from 'node:os' +import { join } from 'node:path' +import { PassThrough, type Readable } from 'node:stream' +import { TieredStorage } from '../src/TieredStorage.js' +import { DiskStorageTier } from '../src/tiers/DiskStorageTier.js' +import { MemoryStorageTier } from '../src/tiers/MemoryStorageTier.js' +import type { StorageMetadata, TieredStorageConfig } from '../src/types/index.js' + +const KEY = 'did:plc:a/site/photo.png' + +function metadata(key: string, size: number): StorageMetadata { + return { + key, + size, + createdAt: new Date(), + lastAccessed: new Date(), + accessCount: 0, + compressed: false, + checksum: 'test', + } +} + +const bytes = (size: number, fill = 7) => new Uint8Array(size).fill(fill) + +async function readAll(stream: NodeJS.ReadableStream): Promise { + const chunks: Buffer[] = [] + for await (const chunk of stream) chunks.push(chunk as Buffer) + return Buffer.concat(chunks) +} + +let directory: string +let hot: MemoryStorageTier +let warm: DiskStorageTier +let cold: MemoryStorageTier +let coldCalls: { getStream: number; getMetadata: number; setMetadata: number } + +function createStorage(config: Partial = {}): TieredStorage { + return new TieredStorage({ + tiers: { hot, warm, cold }, + promotionStrategy: 'eager', + serialization: { serialize: async (data) => data as Uint8Array, deserialize: async (data) => data }, + ...config, + }) +} + +/** Replace the cold stream for the next read with one the test writes to. */ +function controlNextColdStream(): PassThrough { + const stream = new PassThrough() + const original = cold.getStream.bind(cold) + cold.getStream = async (key, options) => { + cold.getStream = original + const result = await original(key, options) + return result && { stream, metadata: result.metadata } + } + return stream +} + +const stagingEntries = async () => readdir(join(directory, '.staging')).catch(() => []) + +describe('TieredStorage streamed-read promotion', () => { + beforeEach(async () => { + directory = await mkdtemp(join(tmpdir(), 'tiered-storage-stream-promotion-')) + hot = new MemoryStorageTier({ maxSizeBytes: 1024 * 1024 }) + warm = new DiskStorageTier({ directory }) + cold = new MemoryStorageTier({ maxSizeBytes: 64 * 1024 * 1024 }) + coldCalls = { getStream: 0, getMetadata: 0, setMetadata: 0 } + for (const method of Object.keys(coldCalls) as Array) { + const original = (cold[method] as (...args: never[]) => unknown).bind(cold) + ;(cold as unknown as Record)[method] = (...args: never[]) => { + coldCalls[method]++ + return original(...args) + } + } + }) + afterEach(async () => rm(directory, { recursive: true, force: true })) + + test('copies a cold read into warm so later reads stop reaching cold', async () => { + const data = bytes(300 * 1024) + await cold.set(KEY, data, metadata(KEY, data.byteLength)) + const storage = createStorage() + + const first = await storage.getStream(KEY) + expect(first?.source).toBe('cold') + expect(await readAll(first!.stream)).toEqual(Buffer.from(data)) + + const second = await storage.getStream(KEY) + expect(second?.source).toBe('warm') + expect(await readAll(second!.stream)).toEqual(Buffer.from(data)) + expect(coldCalls.getStream).toBe(1) + expect(await stagingEntries()).toEqual([]) + }) + + test('never reads or rewrites cold metadata for access stats', async () => { + const data = bytes(1024) + await cold.set(KEY, data, metadata(KEY, data.byteLength)) + const storage = createStorage({ promotionStrategy: 'lazy' }) + + for (let read = 0; read < 3; read++) await readAll((await storage.getStream(KEY))!.stream) + await storage.get(KEY) + await new Promise((resolve) => setTimeout(resolve, 10)) + expect(coldCalls.getMetadata).toBe(0) + expect(coldCalls.setMetadata).toBe(0) + }) + + test('promotes to hot only when placement allows it and the object is small', async () => { + const small = 'did:plc:a/site/app.css' + const large = 'did:plc:a/site/vendor.js' + const media = 'did:plc:a/site/icon.png' + for (const [key, size] of [ + [small, 1024], + [large, 4096], + [media, 1024], + ] as const) { + await cold.set(key, bytes(size), metadata(key, size)) + } + const storage = createStorage({ + streamHotPromotionMaxBytes: 2048, + placementRules: [ + { pattern: '**/*.{css,js}', tiers: ['hot', 'warm', 'cold'] }, + { pattern: '**', tiers: ['warm', 'cold'] }, + ], + }) + + for (const key of [small, large, media]) await readAll((await storage.getStream(key))!.stream) + expect(await hot.exists(small)).toBe(true) + expect(await hot.exists(large)).toBe(false) + expect(await hot.exists(media)).toBe(false) + for (const key of [small, large, media]) expect(await warm.exists(key)).toBe(true) + expect((await storage.getStream(small))?.source).toBe('hot') + }) + + test('promotes a warm read into hot', async () => { + const data = bytes(1024) + await warm.set(KEY, data, metadata(KEY, data.byteLength)) + const storage = createStorage() + + const result = await storage.getStream(KEY) + expect(result?.source).toBe('warm') + await readAll(result!.stream) + expect(await hot.get(KEY)).toEqual(data) + }) + + test('does not promote with the lazy strategy', async () => { + await cold.set(KEY, bytes(1024), metadata(KEY, 1024)) + const storage = createStorage({ promotionStrategy: 'lazy' }) + + await readAll((await storage.getStream(KEY))!.stream) + expect(await warm.exists(KEY)).toBe(false) + expect(await hot.exists(KEY)).toBe(false) + }) + + test('a stream destroyed before its end leaves no partial object', async () => { + await cold.set(KEY, bytes(1024), metadata(KEY, 1024)) + const source = controlNextColdStream() + const storage = createStorage() + + const result = await storage.getStream(KEY) + const stream = result!.stream as Readable + source.write(bytes(512)) + await new Promise((resolve) => stream.once('data', resolve)) + stream.destroy() + await new Promise((resolve) => setTimeout(resolve, 10)) + + expect(source.destroyed).toBe(true) + expect(await warm.exists(KEY)).toBe(false) + expect(await hot.exists(KEY)).toBe(false) + expect(await stagingEntries()).toEqual([]) + }) + + test('a failed source stream leaves no partial object', async () => { + await cold.set(KEY, bytes(1024), metadata(KEY, 1024)) + const source = controlNextColdStream() + const storage = createStorage() + + const result = await storage.getStream(KEY) + source.write(bytes(512)) + source.destroy(new Error('connection reset')) + await expect(readAll(result!.stream)).rejects.toThrow('connection reset') + await new Promise((resolve) => setTimeout(resolve, 10)) + + expect(await warm.exists(KEY)).toBe(false) + expect(await stagingEntries()).toEqual([]) + }) + + test('a source that ends short of its recorded size is not promoted', async () => { + await cold.set(KEY, bytes(1024), metadata(KEY, 1024)) + const source = controlNextColdStream() + const storage = createStorage() + + const result = await storage.getStream(KEY) + source.end(bytes(512)) + expect((await readAll(result!.stream)).byteLength).toBe(512) + + expect(await warm.exists(KEY)).toBe(false) + expect(await hot.exists(KEY)).toBe(false) + expect(await stagingEntries()).toEqual([]) + }) + + test('an invalidation during the stream discards its promotion', async () => { + await cold.set(KEY, bytes(1024), metadata(KEY, 1024)) + const source = controlNextColdStream() + const storage = createStorage() + + const result = await storage.getStream(KEY) + const body = readAll(result!.stream) + source.write(bytes(512)) + await storage.invalidateUpperCacheKey(KEY) + source.end(bytes(512)) + expect((await body).byteLength).toBe(1024) + + expect(await warm.exists(KEY)).toBe(false) + expect(await hot.exists(KEY)).toBe(false) + expect(await stagingEntries()).toEqual([]) + }) + + test('a write during the stream wins over the stale streamed bytes', async () => { + await cold.set(KEY, bytes(1024, 1), metadata(KEY, 1024)) + const source = controlNextColdStream() + const storage = createStorage() + + const result = await storage.getStream(KEY) + const body = readAll(result!.stream) + source.write(bytes(512, 1)) + const fresh = bytes(1024, 2) + const write = storage.set(KEY, fresh) + source.end(bytes(512, 1)) + await body + await write + + expect(await warm.get(KEY)).toEqual(fresh) + expect(await hot.get(KEY)).toEqual(fresh) + expect(await stagingEntries()).toEqual([]) + }) + + test('concurrent cold reads of one key stage a single copy', async () => { + await cold.set(KEY, bytes(1024), metadata(KEY, 1024)) + const storage = createStorage() + let staged = 0 + const stageWrite = warm.stageWrite.bind(warm) + warm.stageWrite = async (...args) => { + staged++ + return stageWrite(...args) + } + + const reads = await Promise.all([storage.getStream(KEY), storage.getStream(KEY)]) + await Promise.all(reads.map((result) => readAll(result!.stream))) + expect(reads.map((result) => result?.source)).toEqual(['cold', 'cold']) + expect(staged).toBe(1) + expect(await warm.exists(KEY)).toBe(true) + }) +}) + +describe('DiskStorageTier staged writes', () => { + beforeEach(async () => { + directory = await mkdtemp(join(tmpdir(), 'disk-tier-staged-')) + warm = new DiskStorageTier({ directory }) + }) + afterEach(async () => rm(directory, { recursive: true, force: true })) + + test('stay invisible until commit and are published whole', async () => { + const staged = await warm.stageWrite(KEY, metadata(KEY, 6)) + await staged.write(new TextEncoder().encode('abc')) + await staged.write(new TextEncoder().encode('def')) + expect(await warm.exists(KEY)).toBe(false) + const keys: string[] = [] + for await (const key of warm.listKeys()) keys.push(key) + expect(keys).toEqual([]) + expect((await warm.getStats()).items).toBe(0) + + await staged.commit() + expect(new TextDecoder().decode((await warm.get(KEY))!)).toBe('abcdef') + expect((await warm.getStats()).bytes).toBe(6) + expect(await stagingEntries()).toEqual([]) + }) + + test('abort removes the staged bytes and a later commit is a no-op', async () => { + const staged = await warm.stageWrite(KEY, metadata(KEY, 3)) + await staged.write(new TextEncoder().encode('abc')) + await staged.abort() + await staged.commit() + expect(await warm.exists(KEY)).toBe(false) + expect(await stagingEntries()).toEqual([]) + }) + + test('a commit larger than the tier cap removes the stale entry instead', async () => { + const capped = new DiskStorageTier({ directory, maxSizeBytes: 4 }) + await capped.set(KEY, bytes(2), metadata(KEY, 2)) + const staged = await capped.stageWrite(KEY, metadata(KEY, 8)) + await staged.write(bytes(8)) + await staged.commit() + expect(await capped.exists(KEY)).toBe(false) + expect(await stagingEntries()).toEqual([]) + }) + + test('startup removes staging files abandoned by a previous process', async () => { + const staged = await warm.stageWrite(KEY, metadata(KEY, 3)) + await staged.write(bytes(3)) + expect(await stagingEntries()).toHaveLength(1) + + const restarted = new DiskStorageTier({ directory }) + await restarted.getStats() + expect(await stagingEntries()).toEqual([]) + expect(await restarted.exists(KEY)).toBe(false) + }) +})