mirror of
https://github.com/mauriceboe/TREK.git
synced 2026-10-19 05:26:47 +00:00
feat(server): S3 driver put — bounded-peek routing, Upload inactivity deadline
This commit is contained in:
jubnl
@@ -1,10 +1,12 @@
|
||||
/// <reference types="@aws-lite/s3-types" />
|
||||
import awsLite from '@aws-lite/client';
|
||||
import type { Readable } from 'node:stream';
|
||||
import { PassThrough, type Readable } from 'node:stream';
|
||||
import fs from 'node:fs';
|
||||
import { assertValidKey, assertValidPrefix, isValidKey } from '../storage-keys';
|
||||
import {
|
||||
StorageBackendError,
|
||||
StorageNotFoundError,
|
||||
isLocalTempFile,
|
||||
type ByteRange,
|
||||
type LocalTempFile,
|
||||
type ObjectStat,
|
||||
@@ -129,8 +131,92 @@ export class S3Driver implements StorageDriver {
|
||||
return null;
|
||||
}
|
||||
|
||||
async put(_key: string, _source: Readable | LocalTempFile, _opts?: PutOptions): Promise<void> {
|
||||
throw new StorageBackendError('not implemented');
|
||||
async put(key: string, source: Readable | LocalTempFile, opts?: PutOptions): Promise<void> {
|
||||
assertValidKey(key);
|
||||
const s3 = await this.client();
|
||||
const base: Record<string, unknown> = { Bucket: this.bucket, Key: this.keyPrefix + key };
|
||||
if (opts?.contentType) base.ContentType = opts.contentType;
|
||||
|
||||
if (isLocalTempFile(source)) {
|
||||
// Ownership transfer: consume the tmp file on success AND failure.
|
||||
try {
|
||||
const { size } = await fs.promises.stat(source.tmpPath);
|
||||
if (size < MULTIPART_THRESHOLD) {
|
||||
await this.deadlined(s3.PutObject({ ...base, File: source.tmpPath }), `put '${key}'`);
|
||||
} else {
|
||||
await this.uploadMonitored(s3, base, fs.createReadStream(source.tmpPath), key);
|
||||
}
|
||||
} catch (err) {
|
||||
throw this.wrap(err, `put failed for '${key}'`);
|
||||
} finally {
|
||||
await fs.promises.rm(source.tmpPath, { force: true });
|
||||
}
|
||||
return;
|
||||
}
|
||||
|
||||
// Bounded peek (spec, Client): a stream ending within the threshold —
|
||||
// including zero bytes — becomes one PutObject with the buffered Body (the
|
||||
// same memory Upload would buffer for part 1, and one round trip instead
|
||||
// of 3+; aws-lite's Upload also breaks on empty streams). A larger stream
|
||||
// flows into Upload with the prefix re-attached. An unbounded stream never
|
||||
// reaches PutObject, which would buffer it wholesale.
|
||||
const peeked = await peekStream(source, MULTIPART_THRESHOLD); // source errors surface untouched
|
||||
// `readableEnded` guards a race: a source whose 'end' fires between the
|
||||
// peek resolving (size > threshold) and pipe() attaching would never end
|
||||
// the PassThrough — but ended means the chunks already hold everything.
|
||||
if (peeked.ended || source.readableEnded) {
|
||||
try {
|
||||
await this.deadlined(s3.PutObject({ ...base, Body: Buffer.concat(peeked.chunks) }), `put '${key}'`);
|
||||
} catch (err) {
|
||||
throw this.wrap(err, `put failed for '${key}'`);
|
||||
}
|
||||
return;
|
||||
}
|
||||
const combined = new PassThrough();
|
||||
for (const chunk of peeked.chunks) combined.write(chunk);
|
||||
source.pipe(combined);
|
||||
source.on('error', (err) => combined.destroy(err));
|
||||
try {
|
||||
await this.uploadMonitored(s3, base, combined, key, source);
|
||||
} catch (err) {
|
||||
throw this.wrap(err, `put failed for '${key}'`);
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* Multipart Upload under an INACTIVITY deadline: a fixed whole-operation
|
||||
* deadline would break large backups, so the timer resets on data movement.
|
||||
* Expiry destroys the monitored stream, driving Upload's own error path —
|
||||
* in-flight parts drain, then AbortMultipartUpload fires (verified upstream).
|
||||
*/
|
||||
private async uploadMonitored(
|
||||
s3: S3Api,
|
||||
base: Record<string, unknown>,
|
||||
body: Readable,
|
||||
key: string,
|
||||
sourceToDestroy?: Readable,
|
||||
): Promise<void> {
|
||||
let timer: NodeJS.Timeout | undefined;
|
||||
const expire = () => {
|
||||
const err = new StorageBackendError(
|
||||
`put '${key}' stalled for ${this.timeoutMs}ms on '${this.id}' — multipart upload aborted`,
|
||||
);
|
||||
body.destroy(err);
|
||||
sourceToDestroy?.destroy(err);
|
||||
};
|
||||
const reset = () => {
|
||||
clearTimeout(timer);
|
||||
timer = setTimeout(expire, this.timeoutMs);
|
||||
timer.unref?.();
|
||||
};
|
||||
body.on('data', reset); // body is consumed by Upload; extra listener observes only
|
||||
reset();
|
||||
try {
|
||||
await s3.Upload({ ...base, Body: body, ChunkSize: MULTIPART_THRESHOLD });
|
||||
} finally {
|
||||
clearTimeout(timer);
|
||||
body.off('data', reset);
|
||||
}
|
||||
}
|
||||
|
||||
async getStream(key: string, range?: ByteRange): Promise<{ stream: Readable; stat: ObjectStat }> {
|
||||
@@ -265,3 +351,50 @@ function totalFromContentRange(contentRange: string | undefined, key: string, id
|
||||
}
|
||||
return Number(total);
|
||||
}
|
||||
|
||||
/**
|
||||
* Reads up to `limit` bytes in paused mode. Resolves `ended: true` when the
|
||||
* stream finished within the limit (chunks = the whole payload), else
|
||||
* `ended: false` with the peeked prefix still unconsumed downstream. Rejects
|
||||
* with the source's own error untouched (LocalDriver parity).
|
||||
*/
|
||||
function peekStream(source: Readable, limit: number): Promise<{ chunks: Buffer[]; size: number; ended: boolean }> {
|
||||
return new Promise((resolve, reject) => {
|
||||
const chunks: Buffer[] = [];
|
||||
let size = 0;
|
||||
const cleanup = () => {
|
||||
source.off('readable', onReadable);
|
||||
source.off('end', onEnd);
|
||||
source.off('error', onError);
|
||||
};
|
||||
const onReadable = () => {
|
||||
for (;;) {
|
||||
const raw = source.read() as Buffer | string | null;
|
||||
if (raw === null) return;
|
||||
// A source not explicitly in binary mode (e.g. an object-mode
|
||||
// Readable.from(string), as used by callers/tests) can yield string
|
||||
// chunks; PutObject's Body and Buffer.concat both need real Buffers.
|
||||
const chunk = Buffer.isBuffer(raw) ? raw : Buffer.from(raw);
|
||||
chunks.push(chunk);
|
||||
size += chunk.length;
|
||||
if (size > limit) {
|
||||
cleanup();
|
||||
resolve({ chunks, size, ended: false });
|
||||
return;
|
||||
}
|
||||
}
|
||||
};
|
||||
const onEnd = () => {
|
||||
cleanup();
|
||||
resolve({ chunks, size, ended: true });
|
||||
};
|
||||
const onError = (err: Error) => {
|
||||
cleanup();
|
||||
reject(err);
|
||||
};
|
||||
source.on('readable', onReadable);
|
||||
source.on('end', onEnd);
|
||||
source.on('error', onError);
|
||||
onReadable();
|
||||
});
|
||||
}
|
||||
|
||||
@@ -1,5 +1,8 @@
|
||||
import { describe, it, expect, vi } from 'vitest';
|
||||
import { Readable } from 'node:stream';
|
||||
import fs from 'node:fs';
|
||||
import os from 'node:os';
|
||||
import path from 'node:path';
|
||||
import {
|
||||
S3Driver,
|
||||
MULTIPART_THRESHOLD,
|
||||
@@ -233,3 +236,119 @@ describe('defaultClientFactory', () => {
|
||||
}
|
||||
});
|
||||
});
|
||||
|
||||
describe('S3Driver put — bounded peek routing', () => {
|
||||
it('routes a small stream (ends within threshold) to a single PutObject with buffered Body', async () => {
|
||||
const api = makeMockApi({ PutObject: vi.fn().mockResolvedValue({}) });
|
||||
await makeDriver(api).put('a/x.bin', Readable.from('hello'), { contentType: 'text/plain' });
|
||||
const input = (api.PutObject as ReturnType<typeof vi.fn>).mock.calls[0][0];
|
||||
expect(Buffer.isBuffer(input.Body) && input.Body.toString()).toBe('hello');
|
||||
expect(input).toMatchObject({ Key: 'a/x.bin', ContentType: 'text/plain' });
|
||||
expect(api.Upload).not.toHaveBeenCalled();
|
||||
});
|
||||
it('routes a zero-byte stream to PutObject (Upload breaks on empty streams upstream)', async () => {
|
||||
const api = makeMockApi({ PutObject: vi.fn().mockResolvedValue({}) });
|
||||
await makeDriver(api).put('a/empty.bin', Readable.from([]));
|
||||
const input = (api.PutObject as ReturnType<typeof vi.fn>).mock.calls[0][0];
|
||||
expect(Buffer.isBuffer(input.Body) && input.Body.length).toBe(0);
|
||||
});
|
||||
it('routes a stream exceeding the threshold to Upload with the peeked prefix re-attached', async () => {
|
||||
let uploaded: Buffer | null = null;
|
||||
const api = makeMockApi({
|
||||
Upload: vi.fn().mockImplementation(async ({ Body }: { Body: Readable }) => {
|
||||
uploaded = await drain(Body);
|
||||
}),
|
||||
});
|
||||
// Three chunks so the source cannot end during the peek (the peek stops
|
||||
// after chunk 2 crosses the threshold; chunk 3 is still pending) — the
|
||||
// Upload path is deterministic, not timing-dependent.
|
||||
const chunk = Buffer.alloc(6 * 1024 * 1024, 7);
|
||||
const big = Buffer.concat([chunk, chunk, chunk]);
|
||||
await makeDriver(api).put('a/big.bin', Readable.from([chunk, chunk, chunk]));
|
||||
expect(api.PutObject).not.toHaveBeenCalled();
|
||||
expect(uploaded!.length).toBe(big.length);
|
||||
expect(uploaded!.equals(big)).toBe(true);
|
||||
expect((api.Upload as ReturnType<typeof vi.fn>).mock.calls[0][0]).toMatchObject({
|
||||
Key: 'a/big.bin',
|
||||
ChunkSize: MULTIPART_THRESHOLD,
|
||||
});
|
||||
});
|
||||
it('routes a single-chunk oversized stream that ends at the peek boundary to PutObject (readableEnded guard)', async () => {
|
||||
// One huge chunk: 'end' can fire between the peek resolving and pipe()
|
||||
// attaching — the driver must detect readableEnded and PutObject the
|
||||
// already-buffered payload rather than hang a never-ending PassThrough.
|
||||
const api = makeMockApi({
|
||||
PutObject: vi.fn().mockResolvedValue({}),
|
||||
Upload: vi.fn().mockResolvedValue({}),
|
||||
});
|
||||
const big = Buffer.alloc(MULTIPART_THRESHOLD + 3, 7);
|
||||
const oneChunk = new Readable({ read() {} });
|
||||
oneChunk.push(big);
|
||||
oneChunk.push(null);
|
||||
await makeDriver(api).put('a/one-chunk.bin', oneChunk);
|
||||
// Either path must complete with the full payload; no hang, no data loss.
|
||||
const putCalls = (api.PutObject as ReturnType<typeof vi.fn>).mock.calls;
|
||||
const uploadCalls = (api.Upload as ReturnType<typeof vi.fn>).mock.calls;
|
||||
expect(putCalls.length + uploadCalls.length).toBe(1);
|
||||
if (putCalls.length === 1) {
|
||||
expect((putCalls[0][0].Body as Buffer).length).toBe(big.length);
|
||||
}
|
||||
});
|
||||
it('surfaces a source-stream error before any client call', async () => {
|
||||
const api = makeMockApi();
|
||||
const boom = new Readable({
|
||||
read() {
|
||||
this.destroy(new Error('source exploded'));
|
||||
},
|
||||
});
|
||||
await expect(makeDriver(api).put('a/broken.bin', boom)).rejects.toThrow('source exploded');
|
||||
expect(api.PutObject).not.toHaveBeenCalled();
|
||||
expect(api.Upload).not.toHaveBeenCalled();
|
||||
});
|
||||
it('expires a stalled Upload via the inactivity deadline (PassThrough destroyed)', async () => {
|
||||
const api = makeMockApi({
|
||||
Upload: vi.fn().mockImplementation(
|
||||
({ Body }: { Body: Readable }) =>
|
||||
new Promise((_, reject) => {
|
||||
Body.on('error', reject); // aws-lite Upload fails when its source errors
|
||||
}),
|
||||
),
|
||||
});
|
||||
const stalled = new Readable({ read() {} });
|
||||
stalled.push(Buffer.alloc(MULTIPART_THRESHOLD + 1)); // exceed the peek, then go silent
|
||||
await expect(makeDriver(api, { timeoutMs: 60 }).put('a/stall.bin', stalled)).rejects.toBeInstanceOf(
|
||||
StorageBackendError,
|
||||
);
|
||||
});
|
||||
});
|
||||
|
||||
describe('S3Driver put — LocalTempFile ownership', () => {
|
||||
async function makeTmp(bytes: number): Promise<string> {
|
||||
const dir = await fs.promises.mkdtemp(path.join(os.tmpdir(), 'trek-s3-'));
|
||||
const p = path.join(dir, 'src.bin');
|
||||
await fs.promises.writeFile(p, Buffer.alloc(bytes, 1));
|
||||
return p;
|
||||
}
|
||||
it('uploads a small temp file via PutObject File and consumes it', async () => {
|
||||
const api = makeMockApi({ PutObject: vi.fn().mockResolvedValue({}) });
|
||||
const tmp = await makeTmp(10);
|
||||
await makeDriver(api).put('a/t.bin', { tmpPath: tmp }, { contentType: 'image/png' });
|
||||
expect(api.PutObject).toHaveBeenCalledWith(
|
||||
expect.objectContaining({ File: tmp, Key: 'a/t.bin', ContentType: 'image/png' }),
|
||||
);
|
||||
expect(fs.existsSync(tmp)).toBe(false);
|
||||
});
|
||||
it('routes a threshold-sized temp file to Upload', async () => {
|
||||
const api = makeMockApi({ Upload: vi.fn().mockImplementation(async ({ Body }: { Body: Readable }) => drain(Body)) });
|
||||
const tmp = await makeTmp(MULTIPART_THRESHOLD);
|
||||
await makeDriver(api).put('a/t.bin', { tmpPath: tmp });
|
||||
expect(api.Upload).toHaveBeenCalled();
|
||||
expect(fs.existsSync(tmp)).toBe(false);
|
||||
});
|
||||
it('consumes the temp file on the failure path too (MirrorDriver precedent)', async () => {
|
||||
const api = makeMockApi({ PutObject: vi.fn().mockRejectedValue(awsError(500)) });
|
||||
const tmp = await makeTmp(10);
|
||||
await expect(makeDriver(api).put('a/t.bin', { tmpPath: tmp })).rejects.toBeInstanceOf(StorageBackendError);
|
||||
expect(fs.existsSync(tmp)).toBe(false);
|
||||
});
|
||||
});
|
||||
|
||||
Reference in New Issue
Block a user