diff --git a/core/services/gcs/src/backend.rs b/core/services/gcs/src/backend.rs index c038fb5de58c..bfd9fee06d07 100644 --- a/core/services/gcs/src/backend.rs +++ b/core/services/gcs/src/backend.rs @@ -45,6 +45,7 @@ use super::core::*; use super::deleter::GcsDeleter; use super::lister::GcsLister; use super::reader::*; +use super::writer::GcsResumableWriter; use super::writer::GcsWriter; use super::writer::GcsWriters; use opendal_core::raw::*; @@ -478,13 +479,24 @@ impl Service for GcsBackend { fn write(&self, ctx: &OperationContext, path: &str, args: OpWrite) -> Result { let output: GcsWriters = { - let concurrent = args.concurrent(); - let w = GcsWriter::new(self.core.clone(), ctx.clone(), path, args); - // Multipart uploads schedule work through the operation executor - // supplied by the caller. - let w = oio::MultipartWriter::new(ctx.executor().clone(), w, concurrent); - - Ok(w) + if args.if_not_exists() { + Ok(TwoWays::Two(GcsResumableWriter::new( + self.core.clone(), + ctx.clone(), + path, + args, + ))) + } else { + let concurrent = args.concurrent(); + let w = GcsWriter::new(self.core.clone(), ctx.clone(), path, args); + // Multipart uploads schedule work through the operation executor + // supplied by the caller. + Ok(TwoWays::One(oio::MultipartWriter::new( + ctx.executor().clone(), + w, + concurrent, + ))) + } }?; Ok(output) diff --git a/core/services/gcs/src/core.rs b/core/services/gcs/src/core.rs index 5faf9ac04874..a74fa2d5c2e6 100644 --- a/core/services/gcs/src/core.rs +++ b/core/services/gcs/src/core.rs @@ -28,6 +28,7 @@ use http::header::CACHE_CONTROL; use http::header::CONTENT_DISPOSITION; use http::header::CONTENT_ENCODING; use http::header::CONTENT_LENGTH; +use http::header::CONTENT_RANGE; use http::header::CONTENT_TYPE; use http::header::HOST; use http::header::IF_MATCH; @@ -54,6 +55,12 @@ pub mod constants { pub const X_GOOG_ACL: &str = "x-goog-acl"; pub const X_GOOG_STORAGE_CLASS: &str = "x-goog-storage-class"; pub const X_GOOG_META_PREFIX: &str = "x-goog-meta-"; + pub const X_UPLOAD_CONTENT_TYPE: &str = "x-upload-content-type"; + + /// Intermediate resumable upload chunks must be multiples of 256 KiB. + /// + /// ref: + pub const GCS_RESUMABLE_CHUNK_SIZE: usize = 256 * 1024; } pub struct GcsCore { @@ -329,6 +336,108 @@ impl GcsCore { } } + pub async fn gcs_initiate_resumable_upload( + &self, + ctx: &OperationContext, + path: &str, + op: &OpWrite, + ) -> Result> { + let p = build_abs_path(&self.root, path); + + let request_metadata = InsertRequestMetadata { + storage_class: self.default_storage_class.as_deref(), + cache_control: op.cache_control(), + content_type: op.content_type(), + content_encoding: op.content_encoding(), + metadata: op.user_metadata(), + }; + + let mut url = format!( + "{}/upload/storage/v1/b/{}/o?uploadType=resumable&name={}", + self.endpoint, + self.bucket, + gcs_percent_encode_path(&p) + ); + + if let Some(acl) = &self.predefined_acl { + write!(&mut url, "&predefinedAcl={acl}").unwrap(); + } + if op.if_not_exists() { + write!(&mut url, "&ifGenerationMatch=0").unwrap(); + } + + let content_type = op.content_type().unwrap_or("application/octet-stream"); + let mut req = Request::post(&url) + .header(X_UPLOAD_CONTENT_TYPE, content_type) + .extension(Operation::Write) + .extension(ServiceOperation("InitiateResumableUpload")); + + let req = if request_metadata.is_empty() { + req.header(CONTENT_LENGTH, 0) + .body(Buffer::new()) + .map_err(new_request_build_error)? + } else { + let body = Buffer::from( + serde_json::to_vec(&request_metadata) + .expect("metadata serialization should succeed"), + ); + req.header(CONTENT_TYPE, "application/json; charset=UTF-8") + .header(CONTENT_LENGTH, body.len()) + .body(body) + .map_err(new_request_build_error)? + }; + + let req = self.sign(ctx, req).await?; + self.send(ctx, req).await + } + + pub async fn gcs_upload_resumable_chunk( + &self, + ctx: &OperationContext, + session_uri: &str, + start: u64, + total: Option, + body: Buffer, + ) -> Result> { + let size = body.len() as u64; + let content_range = match (size, total) { + (0, Some(total)) => format!("bytes */{total}"), + (0, None) => { + return Err(Error::new( + ErrorKind::Unexpected, + "empty resumable chunk requires a known total size", + )); + } + (_, Some(total)) => format!("bytes {start}-{}/{total}", start + size - 1), + (_, None) => format!("bytes {start}-{}/*", start + size - 1), + }; + + let req = Request::put(session_uri) + .header(CONTENT_LENGTH, size) + .header(CONTENT_RANGE, content_range) + .extension(Operation::Write) + .extension(ServiceOperation("UploadResumableChunk")) + .body(body) + .map_err(new_request_build_error)?; + + let req = self.sign(ctx, req).await?; + self.send(ctx, req).await + } + + pub async fn gcs_abort_resumable_upload( + &self, + ctx: &OperationContext, + session_uri: &str, + ) -> Result> { + let req = Request::delete(session_uri) + .extension(Operation::Write) + .extension(ServiceOperation("AbortResumableUpload")) + .body(Buffer::new()) + .map_err(new_request_build_error)?; + let req = self.sign(ctx, req).await?; + self.send(ctx, req).await + } + // It's for presign operation. Gcs only supports query sign over XML API. pub fn gcs_insert_object_xml_request( &self, diff --git a/core/services/gcs/src/writer.rs b/core/services/gcs/src/writer.rs index 9e50e3b9e0c5..b03950483c21 100644 --- a/core/services/gcs/src/writer.rs +++ b/core/services/gcs/src/writer.rs @@ -15,6 +15,7 @@ // specific language governing permissions and limitations // under the License. +use std::collections::VecDeque; use std::sync::Arc; use bytes::Buf; @@ -23,11 +24,12 @@ use http::StatusCode; use super::core::CompleteMultipartUploadRequestPart; use super::core::GcsCore; use super::core::InitiateMultipartUploadResult; +use super::core::constants::GCS_RESUMABLE_CHUNK_SIZE; use super::core::parse_error; use opendal_core::raw::*; use opendal_core::*; -pub type GcsWriters = oio::MultipartWriter; +pub type GcsWriters = TwoWays, GcsResumableWriter>; pub struct GcsWriter { core: Arc, @@ -49,26 +51,7 @@ impl GcsWriter { impl oio::MultipartWrite for GcsWriter { async fn write_once(&self, _: u64, body: Buffer) -> Result { - let size = body.len() as u64; - // Request builders own percent-encoding; writers pass logical paths unchanged. - let req = self - .core - .gcs_insert_object_request(&self.path, Some(size), &self.op, body)?; - - let req = self.core.sign(&self.ctx, req).await?; - - let resp = self.core.send(&self.ctx, req).await?; - - let status = resp.status(); - - match status { - StatusCode::CREATED | StatusCode::OK => { - let metadata = - GcsCore::build_metadata_from_object_response(&self.path, resp.into_body())?; - Ok(metadata) - } - _ => Err(parse_error(resp)), - } + write_object_once(&self.core, &self.ctx, &self.path, &self.op, body).await } async fn initiate_part(&self) -> Result { @@ -162,3 +145,162 @@ impl oio::MultipartWrite for GcsWriter { } } } + +/// JSON API resumable uploads can carry `ifGenerationMatch=0`. XML multipart +/// uploads cannot, so chunked `if_not_exists` writes use this path instead. +pub struct GcsResumableWriter { + core: Arc, + ctx: OperationContext, + path: String, + op: OpWrite, + session_uri: Option, + written: u64, + buffer: Buffer, +} + +impl GcsResumableWriter { + pub fn new(core: Arc, ctx: OperationContext, path: &str, op: OpWrite) -> Self { + GcsResumableWriter { + core, + ctx, + path: path.to_string(), + op, + session_uri: None, + written: 0, + buffer: Buffer::new(), + } + } + + fn push_buffer(&mut self, bs: Buffer) { + if bs.is_empty() { + return; + } + if self.buffer.is_empty() { + self.buffer = bs; + return; + } + + let mut parts = VecDeque::new(); + parts.extend(std::mem::replace(&mut self.buffer, Buffer::new())); + parts.extend(bs); + self.buffer = Buffer::from(parts); + } + + async fn ensure_session(&mut self) -> Result { + if let Some(uri) = &self.session_uri { + return Ok(uri.clone()); + } + + let resp = self + .core + .gcs_initiate_resumable_upload(&self.ctx, &self.path, &self.op) + .await?; + if !resp.status().is_success() { + return Err(parse_error(resp)); + } + + let uri = parse_location(resp.headers())? + .ok_or_else(|| { + Error::new( + ErrorKind::Unexpected, + "Location not present in resumable upload response", + ) + })? + .to_string(); + self.session_uri = Some(uri.clone()); + Ok(uri) + } + + async fn upload_chunk(&mut self, body: Buffer, total: Option) -> Result> { + let uri = self.ensure_session().await?; + let size = body.len() as u64; + let start = self.written; + let resp = self + .core + .gcs_upload_resumable_chunk(&self.ctx, &uri, start, total, body) + .await?; + + match resp.status() { + StatusCode::OK | StatusCode::CREATED => { + self.written = start + size; + let metadata = + GcsCore::build_metadata_from_object_response(&self.path, resp.into_body()) + .unwrap_or_else(|_| Metadata::default()); + Ok(Some(metadata)) + } + StatusCode::PERMANENT_REDIRECT => { + self.written = start + size; + Ok(None) + } + _ => Err(parse_error(resp)), + } + } +} + +impl oio::Write for GcsResumableWriter { + async fn write(&mut self, bs: Buffer) -> Result<()> { + if self.session_uri.is_none() && self.buffer.is_empty() { + self.buffer = bs; + return Ok(()); + } + + self.push_buffer(bs); + while self.buffer.len() >= GCS_RESUMABLE_CHUNK_SIZE { + let aligned = self.buffer.len() - (self.buffer.len() % GCS_RESUMABLE_CHUNK_SIZE); + let chunk = self.buffer.split_to(aligned); + self.upload_chunk(chunk, None).await?; + } + Ok(()) + } + + async fn close(&mut self) -> Result { + if self.session_uri.is_none() { + let body = std::mem::replace(&mut self.buffer, Buffer::new()); + return write_object_once(&self.core, &self.ctx, &self.path, &self.op, body).await; + } + + let remaining = std::mem::replace(&mut self.buffer, Buffer::new()); + let total = self.written + remaining.len() as u64; + match self.upload_chunk(remaining, Some(total)).await? { + Some(meta) => Ok(meta), + None => Ok(Metadata::default()), + } + } + + async fn abort(&mut self) -> Result<()> { + self.buffer = Buffer::new(); + let Some(uri) = self.session_uri.take() else { + return Ok(()); + }; + + let resp = self + .core + .gcs_abort_resumable_upload(&self.ctx, &uri) + .await?; + match resp.status() { + s if s.is_success() || s.as_u16() == 499 => Ok(()), + _ => Err(parse_error(resp)), + } + } +} + +async fn write_object_once( + core: &GcsCore, + ctx: &OperationContext, + path: &str, + op: &OpWrite, + body: Buffer, +) -> Result { + let size = body.len() as u64; + // Request builders own percent-encoding; writers pass logical paths unchanged. + let req = core.gcs_insert_object_request(path, Some(size), op, body)?; + let req = core.sign(ctx, req).await?; + let resp = core.send(ctx, req).await?; + + match resp.status() { + StatusCode::CREATED | StatusCode::OK => { + GcsCore::build_metadata_from_object_response(path, resp.into_body()) + } + _ => Err(parse_error(resp)), + } +} diff --git a/core/tests/behavior/async_write.rs b/core/tests/behavior/async_write.rs index 5a17dc69cf82..267bcd4c3150 100644 --- a/core/tests/behavior/async_write.rs +++ b/core/tests/behavior/async_write.rs @@ -821,14 +821,6 @@ pub async fn test_writer_write_with_if_not_exists(op: Operator) -> Result<()> { return Ok(()); } - // GCS XML API multipart uploads do not support preconditions, so the multipart - // writer path cannot honor if_not_exists. Tracked in - // https://github.com/apache/opendal/issues/8040 - #[cfg(feature = "services-gcs")] - if op.info().scheme() == services::GCS_SCHEME { - return Ok(()); - } - let path = TEST_FIXTURE.new_file_path(); let content = gen_fixed_bytes(cap.write_multi_min_size.unwrap_or(1));