Skip to content
Open
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
125 changes: 44 additions & 81 deletions src/runtime/webcore/s3/multipart.rs
Original file line number Diff line number Diff line change
Expand Up @@ -92,6 +92,7 @@
use core::cell::Cell;
use core::ffi::c_void;
use std::io::Write as _;
use std::rc::Rc;

use bstr::BStr;

Expand Down Expand Up @@ -245,11 +246,9 @@ pub(crate) enum PartState {
}

pub(crate) struct UploadPart {
/// Raw owned slice; backing allocation length is `allocated_size` (may exceed `data.len()`).
/// Freed via `free_allocated_slice`. Default is a static empty slice.
pub(crate) data: Cell<*const [u8]>,
/// Counted so that `perform` can keep the bytes across a request that frees the part.
pub(crate) data: JsCell<Option<Rc<Vec<u8>>>>,
pub ctx: bun_ptr::BackRef<MultiPartUpload, bun_ptr::Mut>, // BACKREF (LIFETIMES.tsv)
pub(crate) allocated_size: Cell<usize>,
pub(crate) state: Cell<PartState>,
pub(crate) part_number: Cell<u16>, // max is 10,000
pub(crate) retry: Cell<u8>, // auto retry, decrement until 0 and fail after this
Expand All @@ -261,26 +260,19 @@ pub(crate) struct UploadPartResult {
pub(crate) etag: Box<[u8]>,
}

impl UploadPart {
fn free_allocated_slice(&self) {
let allocated_size = self.allocated_size.get();
if allocated_size > 0 {
// SAFETY: `data.ptr` was allocated by the global allocator with capacity == allocated_size
// (either via `to_vec().into_boxed_slice()` where len==cap, or by taking ownership of
// StreamBuffer's backing allocation). Reconstruct and drop.
unsafe {
let ptr = (*self.data.get()).as_ptr().cast_mut();
drop(Vec::from_raw_parts(ptr, allocated_size, allocated_size));
}
}
self.data.set(std::ptr::from_ref::<[u8]>(b"" as &[u8]));
self.allocated_size.set(0);
}
/// Which bytes of `buffered` a new part takes.
#[derive(Clone, Copy)]
enum PartBytes {
/// All of them: the part takes the buffer's allocation.
Whole,
/// A copy of the first `usize` bytes.
First(usize),
}

#[inline]
fn data(&self) -> &[u8] {
// SAFETY: data is either a static empty slice or a live heap slice owned by this part
unsafe { &*self.data.get() }
impl UploadPart {
/// Release the bytes. A part that has sent them does not need them again.
fn free_data(&self) {
self.data.set(None);
}

fn on_part_response(result: S3PartResult, this: *mut c_void) -> bun_jsc::JsResult<()> {
Expand All @@ -293,7 +285,7 @@ impl UploadPart {

if this.state.get() == PartState::Canceled || ctx.state.get() == State::Finished {
scoped_log!(S3MultiPartUpload, "onPartResponse {} canceled", part_number);
this.free_allocated_slice();
this.free_data();
MultiPartUpload::deref_(ctx_ptr);
return Ok(());
}
Expand All @@ -311,7 +303,7 @@ impl UploadPart {
} else {
scoped_log!(S3MultiPartUpload, "onPartResponse {} failed", part_number);
this.state.set(PartState::NotAssigned);
this.free_allocated_slice();
this.free_data();
// The ctx deref must run after fail():
let r = ctx.fail(err);
MultiPartUpload::deref_(ctx_ptr);
Expand All @@ -320,8 +312,8 @@ impl UploadPart {
}
S3PartResult::Etag(etag) => {
scoped_log!(S3MultiPartUpload, "onPartResponse {} success", part_number);
let sent = this.data().len();
this.free_allocated_slice();
let sent = this.data.get().as_deref().map_or(0, Vec::len);
this.free_data();
// we will need to order this
ctx.multipart_etags.with_mut(|etags| {
etags.push(UploadPartResult {
Expand All @@ -345,6 +337,7 @@ impl UploadPart {

fn perform(&self) -> bun_jsc::JsResult<()> {
let ctx = self.ctx.get();
let data = self.data.get().clone();
let mut params_buffer = [0u8; 2048];
let written = {
let mut w: &mut [u8] = &mut params_buffer[..];
Expand All @@ -365,7 +358,7 @@ impl UploadPart {
path: &ctx.path,
method: bun_http::Method::PUT,
proxy_url: ctx.proxy_url(),
body: self.data(),
body: data.as_deref().map_or(&[][..], Vec::as_slice),
search_params: Some(search_params),
request_payer: ctx.request_payer,
..Default::default()
Expand All @@ -390,7 +383,7 @@ impl UploadPart {

match state {
PartState::Pending => {
self.free_allocated_slice();
self.free_data();
}
// if is not pending we will free later or is already freed
_ => {}
Expand All @@ -401,7 +394,7 @@ impl UploadPart {
impl Drop for MultiPartUpload {
fn drop(&mut self) {
scoped_log!(S3MultiPartUpload, "deinit");
// queue: Box<[UploadPart]> — dropped automatically (parts' raw `data` already freed during lifecycle)
// queue: Box<[UploadPart]> — dropped automatically, with any `data` a part still owns
// KeepAlive::unref takes an `EventLoopCtx` (aio cycle-break vtable),
// not `&VirtualMachine`. Route through the global hook like simple_request does.
let _ = self.vm;
Expand Down Expand Up @@ -479,12 +472,8 @@ impl MultiPartUpload {
}

/// This is the only place we allocate the queue or the parts, this is responsible for the flow of parts and the max allowed concurrency
fn get_create_part(
&self,
chunk: &[u8],
allocated_size: usize,
needs_clone: bool,
) -> Option<&UploadPart> {
/// The bytes leave `buffered` only when the queue has a slot, and the part owns them.
fn get_create_part(&self, bytes: PartBytes) -> Option<&UploadPart> {
let mut available = self.available.get();
let Some(index) = available.find_first_set() else {
// this means that the queue is full and we cannot flush it
Expand All @@ -507,8 +496,7 @@ impl MultiPartUpload {
// zero set just in case
for _ in 0..queue_size {
queue.push(UploadPart {
data: Cell::new(std::ptr::from_ref::<[u8]>(b"" as &[u8])),
allocated_size: Cell::new(0),
data: JsCell::new(None),
part_number: Cell::new(0),
ctx: self_ref,
index: Cell::new(0),
Expand All @@ -518,21 +506,16 @@ impl MultiPartUpload {
}
self.queue.set(Some(queue.into_boxed_slice()));
}
let (data, allocated_len): (*const [u8], usize) = if needs_clone {
let owned = Box::<[u8]>::from(chunk);
let len = owned.len();
(bun_core::heap::into_raw(owned).cast_const(), len)
} else {
(std::ptr::from_ref::<[u8]>(chunk), allocated_size)
};

let part_number = self.current_part_number.get();
self.current_part_number.set(part_number + 1);

let queue = self.queue.get().as_deref().expect("queue allocated above");
let queue_item = &queue[index];
queue_item.data.set(data);
queue_item.allocated_size.set(allocated_len);
let data = match bytes {
PartBytes::Whole => self.buffered.take().list,
PartBytes::First(len) => self.buffered.get().slice()[..len].to_vec(),
};
queue_item.data.set(Some(Rc::new(data)));
queue_item.part_number.set(part_number);
queue_item.index.set(index as u8); // @truncate
queue_item.retry.set(self.options.get().retry);
Expand Down Expand Up @@ -888,13 +871,9 @@ impl MultiPartUpload {
)
}

fn enqueue_part(
&self,
chunk: &[u8],
allocated_size: usize,
needs_clone: bool,
) -> bun_jsc::JsResult<bool> {
let Some(part) = self.get_create_part(chunk, allocated_size, needs_clone) else {
/// `Ok(false)`: the queue is full and `buffered` is untouched. On `Err` a part has the bytes.
fn enqueue_part(&self, bytes: PartBytes) -> bun_jsc::JsResult<bool> {
let Some(part) = self.get_create_part(bytes) else {
return Ok(false);
};

Expand Down Expand Up @@ -958,42 +937,26 @@ impl MultiPartUpload {
}
// if is one big chunk we can pass ownership and avoid dupe
if self.buffered.get().cursor == 0 && self.buffered.get().size() == len {
let owned = self.buffered.replace(StreamBuffer::default());
// we need to know the allocated size to free the memory later
let allocated_size = owned.memory_cost();
let slice_len = owned.slice().len();

// we dont care about the result because we are sending everything
if self.enqueue_part(owned.slice(), allocated_size, false)? {
if self.enqueue_part(PartBytes::Whole)? {
scoped_log!(
S3MultiPartUpload,
"processMultiPart {} {} full buffer enqueued",
BStr::new(&self.path),
slice_len
len
);
} else {
scoped_log!(
S3MultiPartUpload,
"processMultiPart {} {} queue full",
BStr::new(&self.path),
len
);
let _ = core::mem::ManuallyDrop::new(owned);
return Ok(());
}
let appended = self.buffered.replace(owned);
if appended.is_not_empty() {
self.buffered
.with_mut(|b| b.write(appended.slice()).map(|_| ()))
.unwrap_or(());
}
scoped_log!(
S3MultiPartUpload,
"processMultiPart {} {} queue full",
BStr::new(&self.path),
slice_len
);

return Ok(());
}

let slice_ptr = std::ptr::from_ref::<[u8]>(&self.buffered.get().slice()[..len]);
// allocated size is the slice len because we dupe the buffer
// SAFETY: slice_ptr points at self.buffered's storage which is not mutated until after enqueue_part dupes it
if self.enqueue_part(unsafe { &*slice_ptr }, len, true)? {
if self.enqueue_part(PartBytes::First(len))? {
scoped_log!(
S3MultiPartUpload,
"processMultiPart {} {} slice enqueued",
Expand Down
112 changes: 112 additions & 0 deletions test/js/bun/s3/s3-upload-terminate.test.ts
Original file line number Diff line number Diff line change
@@ -0,0 +1,112 @@
// A streaming multipart upload buffers bytes until it has a part, then hands
// that part the bytes to send. When the buffer held exactly one part it handed
// the part the buffer's own allocation and gave the buffer up only after the
// part's request had been started, so a request that failed inside the enqueue
// left two owners of those bytes: the part, which the failure frees, and the
// local buffer, freed again when the error unwinds.
//
// terminate() is what makes the enqueue fail that way. The S3 client sends
// nothing new for a VM that is stopping, so it fails the request in place, and
// the upload reports that failure to its promise. Each worker below starts its
// uploads and is terminated as its first parts are being enqueued. A run that
// frees a part's bytes twice aborts the process under ASAN, so a clean exit is
// the check.
import { expect, test } from "bun:test";
import { bunEnv, bunExe } from "harness";

// The terminate has to land while a worker enqueues a part, so each worker
// starts several uploads and several workers run at once.
const WORKERS = 4;
const UPLOADS_PER_WORKER = 8;
// The minimum S3 part size. One chunk is one part, so the first chunk of each
// upload fills the buffer with exactly one part.
const PART_SIZE = 5 * 1024 * 1024;

const worker = /* js */ `
const { parentPort, workerData } = require("node:worker_threads");
const client = new Bun.S3Client({
endpoint: workerData.endpoint,
bucket: "bucket",
accessKeyId: "key",
secretAccessKey: "secret",
region: "us-east-1",
});
const chunk = new Uint8Array(${PART_SIZE}).fill(97);
for (let i = 0; i < ${UPLOADS_PER_WORKER}; i++) {
const stream = new ReadableStream({ pull(controller) { controller.enqueue(chunk); } });
client
.write(workerData.key + "-" + i, new Response(stream), { partSize: ${PART_SIZE}, queueSize: 1, retry: 0 })
.catch(() => {});
}
parentPort.postMessage("started");
`;

const host = /* js */ `
const { Worker } = require("node:worker_threads");
const server = Bun.serve({
port: 0,
async fetch(req) {
const url = new URL(req.url);
if (req.method === "POST" && url.searchParams.has("uploads")) {
return new Response(
"<InitiateMultipartUploadResult><UploadId>upload-1</UploadId></InitiateMultipartUploadResult>",
{ headers: { "content-type": "application/xml" } },
);
}
if (req.method === "PUT") {
try { await req.arrayBuffer(); } catch {}
return new Response("", { headers: { ETag: '"etag"' } });
}
try { await req.text(); } catch {}
return new Response(
'<CompleteMultipartUploadResult><ETag>"etag-1"</ETag></CompleteMultipartUploadResult>',
{ headers: { "content-type": "application/xml" } },
);
},
});
await Promise.all(
Array.from({ length: ${WORKERS} }, (_, i) => {
const w = new Worker(${JSON.stringify(worker)}, {
eval: true,
workerData: { endpoint: server.url.origin, key: "key-" + i },
});
// Terminate as soon as this worker reports that its uploads are started.
// A worker that fails before that fails the run instead of hanging it.
return new Promise((resolve, reject) => {
w.once("message", resolve);
w.once("error", reject);
w.once("exit", code => reject(new Error("worker exited with code " + code + " before it started its uploads")));
}).then(() => w.terminate());
}),
).catch(error => {
console.error(error);
process.exit(1);
});
server.stop(true);
console.log("ok");
`;

test("terminate() while a multipart upload enqueues a part frees the part's bytes once", async () => {
await using proc = Bun.spawn({
cmd: [bunExe(), "-e", host],
env: {
...bunEnv,
// An inherited proxy would take the requests away from the loopback
// stand-in.
HTTP_PROXY: undefined,
Comment thread
claude[bot] marked this conversation as resolved.
HTTPS_PROXY: undefined,
ALL_PROXY: undefined,
http_proxy: undefined,
https_proxy: undefined,
all_proxy: undefined,
// An upload still open when its worker goes is not freed (#39692). That
// leak is not what this checks, and LeakSanitizer would report it.
ASAN_OPTIONS: [bunEnv.ASAN_OPTIONS, "detect_leaks=0"].filter(Boolean).join(":"),
},
stdout: "pipe",
stderr: "pipe",
});
const [stdout, stderr, exitCode] = await Promise.all([proc.stdout.text(), proc.stderr.text(), proc.exited]);
expect({ stdout: stdout.trim(), stderr: stderr.trim() }).toEqual({ stdout: "ok", stderr: "" });
expect(exitCode).toBe(0);
});
Loading