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

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
117 changes: 76 additions & 41 deletions service/src/service/router.ts
Original file line number Diff line number Diff line change
Expand Up @@ -23,6 +23,7 @@ import { jobsSubmitted } from '../metrics';
import { captureTraceCarrier, withSpan } from '../telemetry';
import { Jobs, Languages } from '../enum';
import { FileRefAuthorizationError, authorizeRequestedFiles } from './file-authorization';
import { createUploadSessionRegistrar } from './upload-session';
import { prepareSandboxJobSecurity } from '../sandbox-egress';
import logger from '../logger';

Expand Down Expand Up @@ -431,9 +432,6 @@ router.post('/upload', uploadLimiter, async (req: t.AuthenticatedRequest, res: R
reject(err instanceof Error ? err : new Error(String(err)));
return;
}
connection.set(`session:${session_id}`, sessionKey, 'EX', env.SESSION_CACHE_TTL);
logger.info(`[${INSTANCE_ID}] Upload: Session ID: ${session_id} | User ID: ${userId} | Session key: ${sessionKey}`);

const putHeaders: Record<string, string> = {
'Content-Type': mimeType,
/* file-server URL-decodes this header before storing metadata.
Expand All @@ -444,26 +442,35 @@ router.post('/upload', uploadLimiter, async (req: t.AuthenticatedRequest, res: R
if (readOnly) {
putHeaders['X-Read-Only'] = 'true';
}
axios.put<t.UploadResult>(
`${env.FILE_SERVER_URL}/sessions/${session_id}/objects/${fileId}`,
file,
{
headers: internalServiceHeaders(putHeaders),
maxBodyLength: planFileSize,
maxContentLength: planFileSize,
signal: abortController.signal,
}
)
connection.set(`session:${session_id}`, sessionKey, 'EX', env.SESSION_CACHE_TTL)
.then(() => {
logger.info(`[${INSTANCE_ID}] Upload: Session ID: ${session_id} | User ID: ${userId} | Session key: ${sessionKey}`);
return axios.put<t.UploadResult>(
`${env.FILE_SERVER_URL}/sessions/${session_id}/objects/${fileId}`,
file,
{
headers: internalServiceHeaders(putHeaders),
maxBodyLength: planFileSize,
maxContentLength: planFileSize,
signal: abortController.signal,
},
);
})
.then(response => {
clearTimeout(uploadTimeout);
resolve(response.data);
})
.catch(error => {
clearTimeout(uploadTimeout);
file.resume();
reject(error);
});
});

/* Busboy may take additional event-loop turns to drain the multipart
* body before `finish` aggregates this promise. Mark early failures as
* observed while preserving the original promise for Promise.all. */
void uploadPromise.catch(() => undefined);
uploadPromises.push(uploadPromise);
});

Expand Down Expand Up @@ -541,7 +548,6 @@ router.post('/upload/batch', uploadLimiter, async (req: t.AuthenticatedRequest,
let uploadId: string | undefined;
let uploadVersionRaw: string | undefined;
let readOnly = false;
let sessionKeySet = false;
let hasResponded = false;
let filesLimitReached = false;
/* `SessionKeyResolutionError.status` spans 400 | 500 — 400 is a
Expand All @@ -552,6 +558,14 @@ router.post('/upload/batch', uploadLimiter, async (req: t.AuthenticatedRequest,
* 500 we see and convert it into a single 500 batch response on
* `bb.on('finish')`. */
let serverError: SessionKeyResolutionError | undefined;
/* Redis registration is also a batch-level dependency fault. Keep file
* promises fulfilled while Busboy drains, then surface one 500. */
let sessionRegistrationError: Error | undefined;

const ensureSessionRegistered = createUploadSessionRegistrar((sessionKey) => {
logger.info(`[${INSTANCE_ID}] Batch upload: Session ID: ${session_id} | User ID: ${userId} | Session key: ${sessionKey}`);
return connection.set(`session:${session_id}`, sessionKey, 'EX', env.SESSION_CACHE_TTL);
});

const planFileSize = planLimits[req.planId ?? '']?.max_file_size ?? planLimits.default.max_file_size;
/* See note on the single-upload busboy above for why preservePath is set. */
Expand Down Expand Up @@ -599,6 +613,13 @@ router.post('/upload/batch', uploadLimiter, async (req: t.AuthenticatedRequest,
const uploadTimeout = setTimeout(() => {
abortController.abort('timeout');
file.resume();
/* If Redis is still pending, make that shared registration barrier
* fail before any file promise settles. This prevents Busboy from
* finishing with a 400 while a later Redis rejection arrives too
* late to be surfaced as the batch-level dependency failure. */
if (ensureSessionRegistered.rejectPending(
new Error('Upload session registration timed out'),
)) return;
resolve({ status: 'error', filename, error: 'Upload timeout' });
}, UPLOAD_TIMEOUT_MS);

Expand Down Expand Up @@ -635,12 +656,6 @@ router.post('/upload/batch', uploadLimiter, async (req: t.AuthenticatedRequest,
resolve({ status: 'error', filename, error: message });
return;
}
if (!sessionKeySet) {
connection.set(`session:${session_id}`, sessionKey, 'EX', env.SESSION_CACHE_TTL);
sessionKeySet = true;
logger.info(`[${INSTANCE_ID}] Batch upload: Session ID: ${session_id} | User ID: ${userId} | Session key: ${sessionKey}`);
}

const putHeaders: Record<string, string> = {
'Content-Type': mimeType,
/* file-server URL-decodes this header before storing metadata.
Expand All @@ -651,31 +666,42 @@ router.post('/upload/batch', uploadLimiter, async (req: t.AuthenticatedRequest,
if (readOnly) {
putHeaders['X-Read-Only'] = 'true';
}
axios.put<t.UploadResult>(
const failSessionRegistration = (error: unknown): void => {
clearTimeout(uploadTimeout);
file.resume();
const normalizedError = error instanceof Error ? error : new Error(String(error));
sessionRegistrationError ??= normalizedError;
resolve({ status: 'error', filename, error: 'Failed to register upload session' });
};
const resolveUploadFailure = (error: unknown): void => {
clearTimeout(uploadTimeout);
if (abortController.signal.aborted) {
const reason = abortController.signal.reason === 'timeout' ? 'Upload timeout' : 'File size limit exceeded';
resolve({ status: 'error', filename, error: reason });
return;
}
file.resume();
const message = error instanceof Error ? error.message : 'Unknown upload error';
logger.error(`[${INSTANCE_ID}] Batch upload file failed: ${filename} | Session: ${session_id}`, { error: message });
resolve({ status: 'error', filename, error: message });
};
const forwardFile = (): Promise<void> => axios.put<t.UploadResult>(
`${env.FILE_SERVER_URL}/sessions/${session_id}/objects/${fileId}`,
file,
{
headers: internalServiceHeaders(putHeaders),
maxBodyLength: planFileSize,
maxContentLength: planFileSize,
signal: abortController.signal,
}
)
.then(response => {
clearTimeout(uploadTimeout);
resolve({ status: 'success', filename: response.data.filename, fileId: response.data.fileId });
})
.catch(error => {
clearTimeout(uploadTimeout);
if (abortController.signal.aborted) {
const reason = abortController.signal.reason === 'timeout' ? 'Upload timeout' : 'File size limit exceeded';
resolve({ status: 'error', filename, error: reason });
return;
}
const message = error instanceof Error ? error.message : 'Unknown upload error';
logger.error(`[${INSTANCE_ID}] Batch upload file failed: ${filename} | Session: ${session_id}`, { error: message });
resolve({ status: 'error', filename, error: message });
});
},
).then(response => {
clearTimeout(uploadTimeout);
resolve({ status: 'success', filename: response.data.filename, fileId: response.data.fileId });
}, resolveUploadFailure);

void ensureSessionRegistered(sessionKey)
.then(forwardFile, failSessionRegistration)
.catch(failSessionRegistration);
Comment thread
danny-avila marked this conversation as resolved.
});

uploadPromises.push(uploadPromise);
Expand All @@ -701,6 +727,15 @@ router.post('/upload/batch', uploadLimiter, async (req: t.AuthenticatedRequest,
try {
const results = await Promise.all(uploadPromises);

if (sessionRegistrationError) {
logger.error(
`[${INSTANCE_ID}] Batch upload session registration failed for session ${session_id}:`,
sessionRegistrationError,
);
res.status(500).json({ error: 'Error registering upload session' });
return;
}

/* If sessionKey resolution faulted with a 500 status (server
* misconfiguration — see `serverError` declaration above),
* surface the fault as a single batch-level 500 instead of
Expand All @@ -721,9 +756,9 @@ router.post('/upload/batch', uploadLimiter, async (req: t.AuthenticatedRequest,
return;
}

/* SessionKey was set inline in the per-file handler under
* `sessionKeySet`. No batch-level fallback needed: if zero files
* succeeded, no session was created. */
/* SessionKey was set inline in the per-file handler through
* `ensureSessionRegistered`. No batch-level fallback is needed when
* no valid file reaches the forwarding step. */

let succeeded = 0;
let failed = 0;
Expand Down
107 changes: 107 additions & 0 deletions service/src/service/upload-session.test.ts
Original file line number Diff line number Diff line change
@@ -0,0 +1,107 @@
import { describe, expect, test } from 'bun:test';
import { createUploadSessionRegistrar } from './upload-session';

function deferred(): {
promise: Promise<void>;
resolve: () => void;
reject: (error: Error) => void;
} {
let resolve!: () => void;
let reject!: (error: Error) => void;
const promise = new Promise<void>((resolvePromise, rejectPromise) => {
resolve = resolvePromise;
reject = rejectPromise;
});
return { promise, resolve, reject };
}

describe('createUploadSessionRegistrar', () => {
test('waits for session registration before forwarding a file', async () => {
const registration = deferred();
const events: string[] = [];
const ensureSessionRegistered = createUploadSessionRegistrar((sessionKey) => {
events.push(`session:set:start:${sessionKey}`);
return registration.promise.then(() => {
events.push('session:set:done');
});
});

const result = ensureSessionRegistered('user:user-1').then(async () => {
events.push('file:put');
return 'uploaded';
});

await Promise.resolve();
expect(events).toEqual(['session:set:start:user:user-1']);

registration.resolve();
expect(await result).toBe('uploaded');
expect(events).toEqual(['session:set:start:user:user-1', 'session:set:done', 'file:put']);
});

test('shares one pending registration across every file in a batch', async () => {
const registration = deferred();
let registrations = 0;
const forwarded: string[] = [];
const ensureSessionRegistered = createUploadSessionRegistrar(() => {
registrations += 1;
return registration.promise;
});

const uploads = [
ensureSessionRegistered('user:user-1').then(async () => {
forwarded.push('first');
return 'first';
}),
ensureSessionRegistered('user:user-1').then(async () => {
forwarded.push('second');
return 'second';
}),
];

await Promise.resolve();
expect(registrations).toBe(1);
expect(forwarded).toEqual([]);

registration.resolve();
expect(await Promise.all(uploads)).toEqual(['first', 'second']);
expect(forwarded).toEqual(['first', 'second']);
});

test('does not forward files when session registration fails', async () => {
const registration = deferred();
let forwarded = false;
const ensureSessionRegistered = createUploadSessionRegistrar(() => registration.promise);
const result = ensureSessionRegistered('user:user-1').then(async () => {
forwarded = true;
return 'uploaded';
});

registration.reject(new Error('Redis unavailable'));

await expect(result).rejects.toThrow('Redis unavailable');
expect(forwarded).toBe(false);
});

test('keeps a pending registration timeout terminal when Redis rejects later', async () => {
const registration = deferred();
let forwarded = false;
const ensureSessionRegistered = createUploadSessionRegistrar(() => registration.promise);
const result = ensureSessionRegistered('user:user-1').then(async () => {
forwarded = true;
return 'uploaded';
});

expect(ensureSessionRegistered.rejectPending(
new Error('Upload session registration timed out'),
)).toBe(true);
await expect(result).rejects.toThrow('Upload session registration timed out');
expect(forwarded).toBe(false);

registration.reject(new Error('Redis unavailable after timeout'));
await Promise.resolve();

expect(forwarded).toBe(false);
expect(ensureSessionRegistered.rejectPending(new Error('too late'))).toBe(false);
});
});
67 changes: 67 additions & 0 deletions service/src/service/upload-session.ts
Original file line number Diff line number Diff line change
@@ -0,0 +1,67 @@
type RegisterUploadSession = (sessionKey: string) => Promise<unknown>;

export interface UploadSessionRegistrar {
(sessionKey: string): Promise<unknown>;
/**
* Rejects the shared registration barrier only while Redis is still
* pending. The underlying Redis promise remains observed, but its eventual
* outcome cannot reopen the barrier or forward a file after timeout.
*/
rejectPending(error: Error): boolean;
}

/**
* Creates a request-scoped session registrar. Every file shares the first
* registration promise, allowing callers to wait for the Redis write without
* issuing duplicate SETs for a batch.
*/
export function createUploadSessionRegistrar(
registerSession: RegisterUploadSession,
): UploadSessionRegistrar {
let sessionRegistered: Promise<unknown> | undefined;
let registrationPending = false;
let resolveRegistration!: (value: unknown) => void;
let rejectRegistration!: (error: unknown) => void;

const ensureSessionRegistered = (sessionKey: string): Promise<unknown> => {
if (!sessionRegistered) {
registrationPending = true;
sessionRegistered = new Promise((resolve, reject) => {
resolveRegistration = resolve;
rejectRegistration = reject;
});

let registration: Promise<unknown>;
try {
registration = registerSession(sessionKey);
} catch (error) {
registrationPending = false;
rejectRegistration(error);
return sessionRegistered;
}

void registration.then(
(value) => {
if (!registrationPending) return;
registrationPending = false;
resolveRegistration(value);
},
(error: unknown) => {
if (!registrationPending) return;
registrationPending = false;
rejectRegistration(error);
},
);
}
return sessionRegistered;
};

ensureSessionRegistered.rejectPending = (error: Error): boolean => {
if (!registrationPending) return false;
registrationPending = false;
rejectRegistration(error);
return true;
};

return ensureSessionRegistered;
}