1use std::sync::Arc;
16use std::time::Duration;
17
18use bytes::{Bytes, BytesMut};
19use fail::fail_point;
20use futures::StreamExt;
21use opendal::layers::{RetryLayer, TimeoutLayer};
22use opendal::raw::BoxedStaticFuture;
23use opendal::services::Memory;
24use opendal::{Execute, Executor, Operator, Writer};
25use risingwave_common::config::ObjectStoreConfig;
26use risingwave_common::range::RangeBoundsExt;
27use thiserror_ext::AsReport;
28
29use crate::object::object_metrics::ObjectStoreMetrics;
30use crate::object::{
31 MonitoredObjectStore, ObjectDataStream, ObjectError, ObjectMetadata, ObjectMetadataIter,
32 ObjectRangeBounds, ObjectResult, ObjectStore, OperationType, StreamingUploader, prefix,
33};
34
35#[derive(Clone)]
37pub struct OpendalObjectStore {
38 pub(crate) op: Operator,
39 pub(crate) media_type: MediaType,
40 pub(crate) config: Arc<ObjectStoreConfig>,
41 pub(crate) metrics: Arc<ObjectStoreMetrics>,
42}
43
44#[derive(Clone)]
45pub enum MediaType {
46 Memory,
47 Hdfs,
48 Gcs,
49 Minio,
50 S3,
51 Obs,
52 Oss,
53 Webhdfs,
54 Azblob,
55 Fs,
56}
57
58impl MediaType {
59 pub fn as_str(&self) -> &'static str {
60 match self {
61 MediaType::Memory => "Memory",
62 MediaType::Hdfs => "Hdfs",
63 MediaType::Gcs => "Gcs",
64 MediaType::Minio => "Minio",
65 MediaType::S3 => "S3",
66 MediaType::Obs => "Obs",
67 MediaType::Oss => "Oss",
68 MediaType::Webhdfs => "Webhdfs",
69 MediaType::Azblob => "Azblob",
70 MediaType::Fs => "Fs",
71 }
72 }
73}
74
75impl OpendalObjectStore {
76 pub fn test_new_memory_engine() -> ObjectResult<Self> {
78 let builder = Memory::default();
80 let op: Operator = Operator::new(builder)?;
81
82 Ok(Self {
83 op,
84 media_type: MediaType::Memory,
85 config: Arc::new(ObjectStoreConfig::default()),
86 metrics: Arc::new(ObjectStoreMetrics::unused()),
87 })
88 }
89}
90
91#[inline]
92fn timestamp_to_secs(t: opendal::raw::Timestamp) -> f64 {
93 t.into_inner().as_second() as f64
94}
95
96#[async_trait::async_trait]
97impl ObjectStore for OpendalObjectStore {
98 type StreamingUploader = OpendalStreamingUploader;
99
100 fn get_object_prefix(&self, obj_id: u64, use_new_object_prefix_strategy: bool) -> String {
101 match self.media_type {
102 MediaType::S3 => prefix::s3::get_object_prefix(obj_id),
103 MediaType::Minio => prefix::s3::get_object_prefix(obj_id),
104 MediaType::Memory => String::default(),
105 MediaType::Hdfs
106 | MediaType::Gcs
107 | MediaType::Obs
108 | MediaType::Oss
109 | MediaType::Webhdfs
110 | MediaType::Azblob
111 | MediaType::Fs => {
112 prefix::opendal_engine::get_object_prefix(obj_id, use_new_object_prefix_strategy)
113 }
114 }
115 }
116
117 async fn upload(&self, path: &str, obj: Bytes) -> ObjectResult<()> {
118 if obj.is_empty() {
119 Err(ObjectError::internal("upload empty object"))
120 } else {
121 self.op.write(path, obj).await?;
122 Ok(())
123 }
124 }
125
126 async fn streaming_upload(&self, path: &str) -> ObjectResult<Self::StreamingUploader> {
127 if self.op.info().capability().write_can_multi {
128 Ok(OpendalStreamingUploader::Native(Box::new(
129 OpendalNativeStreamingUploader::new(
130 self.op.clone(),
131 path.to_owned(),
132 self.config.clone(),
133 self.metrics.clone(),
134 self.store_media_type(),
135 )
136 .await?,
137 )))
138 } else {
139 Ok(OpendalStreamingUploader::Buffered(
140 OpendalBufferedStreamingUploader::new(
141 self.clone()
142 .monitored(self.metrics.clone(), self.config.clone()),
143 path.to_owned(),
144 ),
145 ))
146 }
147 }
148
149 async fn read(&self, path: &str, range: impl ObjectRangeBounds) -> ObjectResult<Bytes> {
150 let data = if range.is_full() {
151 self.op.read(path).await?
152 } else {
153 self.op
154 .read_with(path)
155 .range(range.map(|v| *v as u64))
156 .await?
157 };
158
159 if let Some(len) = range.len()
160 && len != data.len()
161 {
162 return Err(ObjectError::internal(format!(
163 "mismatched size: expected {}, found {} when reading {} at {:?}",
164 len,
165 data.len(),
166 path,
167 range,
168 )));
169 }
170
171 Ok(data.to_bytes())
172 }
173
174 async fn streaming_read(
178 &self,
179 path: &str,
180 range: impl ObjectRangeBounds,
181 ) -> ObjectResult<ObjectDataStream> {
182 fail_point!("opendal_streaming_read_err", |_| Err(
183 ObjectError::internal("opendal streaming read error")
184 ));
185 let range = range.map(|v| *v as u64);
186
187 let reader = self
192 .op
193 .clone()
194 .layer(TimeoutLayer::new().with_io_timeout(Duration::from_millis(
195 self.config.retry.streaming_read_attempt_timeout_ms,
196 )))
197 .layer(
198 RetryLayer::new()
199 .with_min_delay(Duration::from_millis(
200 self.config.retry.req_backoff_interval_ms,
201 ))
202 .with_max_delay(Duration::from_millis(
203 self.config.retry.req_backoff_max_delay_ms,
204 ))
205 .with_max_times(self.config.retry.streaming_read_retry_attempts)
206 .with_factor(self.config.retry.req_backoff_factor as f32)
207 .with_jitter(),
208 )
209 .reader_with(path)
210 .await?;
211 let stream = reader.into_bytes_stream(range).await?.map(|item| {
212 item.map(|b| Bytes::copy_from_slice(b.as_ref()))
213 .map_err(|e| {
214 ObjectError::internal(format!("reader into_stream fail {}", e.as_report()))
215 })
216 });
217
218 Ok(Box::pin(stream))
219 }
220
221 async fn metadata(&self, path: &str) -> ObjectResult<ObjectMetadata> {
222 let opendal_metadata = self.op.stat(path).await?;
223 let key = path.to_owned();
224 let last_modified = match opendal_metadata.last_modified() {
225 Some(t) => timestamp_to_secs(t),
226 None => 0_f64,
227 };
228
229 let total_size = opendal_metadata.content_length() as usize;
230 let metadata = ObjectMetadata {
231 key,
232 last_modified,
233 total_size,
234 };
235 Ok(metadata)
236 }
237
238 async fn delete(&self, path: &str) -> ObjectResult<()> {
239 self.op.delete(path).await?;
240 Ok(())
241 }
242
243 async fn delete_objects(&self, paths: &[String]) -> ObjectResult<()> {
246 self.op.delete_iter(paths.to_vec()).await?;
247 Ok(())
248 }
249
250 async fn list(
251 &self,
252 prefix: &str,
253 start_after: Option<String>,
254 limit: Option<usize>,
255 ) -> ObjectResult<ObjectMetadataIter> {
256 let mut object_lister = self.op.lister_with(prefix).recursive(true);
257 if let Some(start_after) = start_after {
258 object_lister = object_lister.start_after(&start_after);
259 }
260 let object_lister = object_lister.await?;
261
262 let op = self.op.clone();
263 let stream = object_lister.then(move |object| {
264 let op = op.clone();
265 async move { Self::listed_object_metadata(&op, object?).await }
266 });
267
268 Ok(stream.take(limit.unwrap_or(usize::MAX)).boxed())
269 }
270
271 fn store_media_type(&self) -> &'static str {
272 self.media_type.as_str()
273 }
274}
275
276impl OpendalObjectStore {
277 async fn listed_object_metadata(
278 op: &Operator,
279 object: opendal::Entry,
280 ) -> ObjectResult<ObjectMetadata> {
281 let key = object.path().to_owned();
282 let meta = object.metadata();
284 let mut last_modified = meta.last_modified().map(timestamp_to_secs);
285 let mut total_size = meta.content_length() as usize;
286 if last_modified.is_none() || total_size == 0 {
287 let stat_meta = op.stat(&key).await?;
289 last_modified = stat_meta.last_modified().map(timestamp_to_secs);
290 total_size = stat_meta.content_length() as usize;
291 }
292 Ok(ObjectMetadata {
293 key,
294 last_modified: last_modified.unwrap_or(0_f64),
295 total_size,
296 })
297 }
298
299 pub async fn copy(&self, from_path: &str, to_path: &str) -> ObjectResult<()> {
300 self.op.copy(from_path, to_path).await?;
301 Ok(())
302 }
303}
304
305struct OpendalStreamingUploaderExecute {
306 metrics: Arc<ObjectStoreMetrics>,
308 media_type: &'static str,
309}
310
311impl OpendalStreamingUploaderExecute {
312 const STREAMING_UPLOAD_TYPE: OperationType = OperationType::StreamingUpload;
313
314 fn new(metrics: Arc<ObjectStoreMetrics>, media_type: &'static str) -> Self {
315 Self {
316 metrics,
317 media_type,
318 }
319 }
320}
321
322impl Execute for OpendalStreamingUploaderExecute {
323 fn execute(&self, f: BoxedStaticFuture<()>) {
324 let operation_type_str = Self::STREAMING_UPLOAD_TYPE.as_str();
325 let media_type = self.media_type;
326
327 let metrics = self.metrics.clone();
328 let _handle = tokio::spawn(async move {
329 let _timer = metrics
330 .operation_latency
331 .with_label_values(&[media_type, operation_type_str])
332 .start_timer();
333
334 f.await
335 });
336 }
337}
338
339pub enum OpendalStreamingUploader {
340 Native(Box<OpendalNativeStreamingUploader>),
341 Buffered(OpendalBufferedStreamingUploader),
342}
343
344pub struct OpendalNativeStreamingUploader {
346 writer: Writer,
347 buf: Vec<Bytes>,
350 not_uploaded_len: usize,
352 is_valid: bool,
354
355 abort_on_err: bool,
356
357 upload_part_size: usize,
358}
359
360impl OpendalNativeStreamingUploader {
361 pub async fn new(
362 op: Operator,
363 path: String,
364 config: Arc<ObjectStoreConfig>,
365 metrics: Arc<ObjectStoreMetrics>,
366 media_type: &'static str,
367 ) -> ObjectResult<Self> {
368 let monitored_execute = OpendalStreamingUploaderExecute::new(metrics, media_type);
369 let executor = Executor::with(monitored_execute);
370 let ctx = op.base_context().with_executor(executor);
374 let op = op.with_context(ctx);
375
376 let writer = op
377 .clone()
378 .layer(TimeoutLayer::new().with_io_timeout(Duration::from_millis(
379 config.retry.streaming_upload_attempt_timeout_ms,
380 )))
381 .layer(
382 RetryLayer::new()
383 .with_min_delay(Duration::from_millis(config.retry.req_backoff_interval_ms))
384 .with_max_delay(Duration::from_millis(config.retry.req_backoff_max_delay_ms))
385 .with_max_times(config.retry.streaming_upload_retry_attempts)
386 .with_factor(config.retry.req_backoff_factor as f32)
387 .with_jitter(),
388 )
389 .writer_with(&path)
390 .concurrent(config.opendal_upload_concurrency)
391 .await?;
392 Ok(Self {
393 writer,
394 buf: vec![],
395 not_uploaded_len: 0,
396 is_valid: true,
397 abort_on_err: config.opendal_writer_abort_on_err,
398 upload_part_size: config.upload_part_size,
399 })
400 }
401
402 async fn flush(&mut self) -> ObjectResult<()> {
403 let data: Vec<Bytes> = self.buf.drain(..).collect();
404 debug_assert_eq!(
405 data.iter().map(|b| b.len()).sum::<usize>(),
406 self.not_uploaded_len
407 );
408 if let Err(err) = self.writer.write(data).await {
409 self.is_valid = false;
410 if self.abort_on_err {
411 self.writer.abort().await?;
412 }
413 return Err(err.into());
414 }
415 self.not_uploaded_len = 0;
416 Ok(())
417 }
418}
419
420impl StreamingUploader for OpendalNativeStreamingUploader {
421 async fn write_bytes(&mut self, data: Bytes) -> ObjectResult<()> {
422 assert!(self.is_valid);
423 self.not_uploaded_len += data.len();
424 self.buf.push(data);
425 if self.not_uploaded_len >= self.upload_part_size {
426 self.flush().await?;
427 }
428 Ok(())
429 }
430
431 async fn finish(mut self) -> ObjectResult<()> {
432 assert!(self.is_valid);
433 if self.not_uploaded_len > 0 {
434 self.flush().await?;
435 }
436
437 assert!(self.buf.is_empty());
438 assert_eq!(self.not_uploaded_len, 0);
439
440 self.is_valid = false;
441 match self.writer.close().await {
442 Ok(_) => (),
443 Err(err) => {
444 if self.abort_on_err {
445 self.writer.abort().await?;
446 }
447 return Err(err.into());
448 }
449 };
450
451 Ok(())
452 }
453
454 fn get_memory_usage(&self) -> u64 {
456 self.not_uploaded_len as u64
457 }
458}
459
460pub struct OpendalBufferedStreamingUploader {
461 store: MonitoredObjectStore<OpendalObjectStore>,
462 path: String,
463 buf: BytesMut,
464}
465
466impl OpendalBufferedStreamingUploader {
467 fn new(store: MonitoredObjectStore<OpendalObjectStore>, path: String) -> Self {
468 Self {
469 store,
470 path,
471 buf: BytesMut::new(),
472 }
473 }
474}
475
476impl StreamingUploader for OpendalBufferedStreamingUploader {
477 async fn write_bytes(&mut self, data: Bytes) -> ObjectResult<()> {
478 self.buf.extend_from_slice(&data);
479 Ok(())
480 }
481
482 async fn finish(self) -> ObjectResult<()> {
483 self.store.upload(&self.path, self.buf.freeze()).await
484 }
485
486 fn get_memory_usage(&self) -> u64 {
487 self.buf.len() as u64
488 }
489}
490
491impl StreamingUploader for OpendalStreamingUploader {
492 async fn write_bytes(&mut self, data: Bytes) -> ObjectResult<()> {
493 match self {
494 Self::Native(uploader) => uploader.write_bytes(data).await,
495 Self::Buffered(uploader) => uploader.write_bytes(data).await,
496 }
497 }
498
499 async fn finish(self) -> ObjectResult<()> {
500 match self {
501 Self::Native(uploader) => (*uploader).finish().await,
502 Self::Buffered(uploader) => uploader.finish().await,
503 }
504 }
505
506 fn get_memory_usage(&self) -> u64 {
507 match self {
508 Self::Native(uploader) => uploader.get_memory_usage(),
509 Self::Buffered(uploader) => uploader.get_memory_usage(),
510 }
511 }
512}
513
514#[cfg(test)]
515mod tests {
516 use futures::TryStreamExt;
517
518 use super::*;
519
520 async fn list_all(prefix: &str, store: &OpendalObjectStore) -> Vec<ObjectMetadata> {
521 store
522 .list(prefix, None, None)
523 .await
524 .unwrap()
525 .try_collect::<Vec<_>>()
526 .await
527 .unwrap()
528 }
529
530 #[tokio::test]
531 async fn test_memory_upload() {
532 let block = Bytes::from("123456");
533 let store = OpendalObjectStore::test_new_memory_engine().unwrap();
534 store.upload("/abc", block).await.unwrap();
535
536 store.read("/ab", 0..3).await.unwrap_err();
538
539 let bytes = store.read("/abc", 4..6).await.unwrap();
540 assert_eq!(String::from_utf8(bytes.to_vec()).unwrap(), "56".to_owned());
541
542 store.delete("/abc").await.unwrap();
543
544 store.read("/abc", 0..3).await.unwrap_err();
546 }
547
548 #[tokio::test]
549 async fn test_memory_read_out_of_range_returns_error() {
550 let block = Bytes::from("123456");
551 let store = OpendalObjectStore::test_new_memory_engine().unwrap();
552 store.upload("/abc", block).await.unwrap();
553
554 store.read("/abc", 4..44).await.unwrap_err();
556 }
557
558 #[tokio::test]
559 async fn test_memory_metadata() {
560 let block = Bytes::from("123456");
561 let path = "/abc".to_owned();
562 let obj_store = OpendalObjectStore::test_new_memory_engine().unwrap();
563 obj_store.upload("/abc", block).await.unwrap();
564
565 let err = obj_store.metadata("/not_exist").await.unwrap_err();
566 assert!(err.is_object_not_found_error());
567 let metadata = obj_store.metadata("/abc").await.unwrap();
568 assert_eq!(metadata.total_size, 6);
569 obj_store.delete(&path).await.unwrap();
570 }
571
572 #[tokio::test]
573 async fn test_buffered_streaming_upload() {
574 let store = OpendalObjectStore::test_new_memory_engine().unwrap();
575 let mut uploader = OpendalBufferedStreamingUploader::new(
576 store
577 .clone()
578 .monitored(store.metrics.clone(), store.config.clone()),
579 "/abc".to_owned(),
580 );
581
582 uploader.write_bytes(Bytes::from("123")).await.unwrap();
583 assert_eq!(uploader.get_memory_usage(), 3);
584 uploader.write_bytes(Bytes::from("456")).await.unwrap();
585 assert_eq!(uploader.get_memory_usage(), 6);
586 uploader.finish().await.unwrap();
587
588 let read_obj = store.read("/abc", ..).await.unwrap();
589 assert_eq!(read_obj, Bytes::from("123456"));
590 }
591
592 #[tokio::test]
593 async fn test_empty_buffered_streaming_upload() {
594 let store = OpendalObjectStore::test_new_memory_engine().unwrap();
595 let uploader = OpendalBufferedStreamingUploader::new(
596 store
597 .clone()
598 .monitored(store.metrics.clone(), store.config.clone()),
599 "/abc".to_owned(),
600 );
601
602 uploader.finish().await.unwrap_err();
603 }
604
605 #[tokio::test]
606 async fn test_listed_object_metadata_propagates_stat_failure() {
607 let store = OpendalObjectStore::test_new_memory_engine().unwrap();
608 store
609 .upload("sst", Bytes::from_static(b"data"))
610 .await
611 .unwrap();
612 let mut objects = store.op.lister_with("").recursive(true).await.unwrap();
615 let object = objects.next().await.unwrap().unwrap();
616 store.delete("sst").await.unwrap();
617 let error = OpendalObjectStore::listed_object_metadata(&store.op, object)
618 .await
619 .unwrap_err();
620 assert!(error.is_object_not_found_error());
621 }
622
623 #[tokio::test]
624 async fn test_memory_delete_objects_and_list_object() {
625 let block1 = Bytes::from("123456");
626 let block2 = Bytes::from("987654");
627 let store = OpendalObjectStore::test_new_memory_engine().unwrap();
628 store.upload("abc", Bytes::from("123456")).await.unwrap();
629 store.upload("prefix/abc", block1).await.unwrap();
630
631 store.upload("prefix/xyz", block2).await.unwrap();
632
633 assert_eq!(list_all("", &store).await.len(), 3);
634 assert_eq!(list_all("prefix/", &store).await.len(), 2);
635 let str_list = [String::from("prefix/abc"), String::from("prefix/xyz")];
636
637 store.delete_objects(&str_list).await.unwrap();
638
639 assert!(store.read("prefix/abc/", ..).await.is_err());
640 assert!(store.read("prefix/xyz/", ..).await.is_err());
641 assert_eq!(list_all("", &store).await.len(), 1);
642 assert_eq!(list_all("prefix/", &store).await.len(), 0);
643 }
644}