1use std::sync::Arc;
18use std::time::{Duration, Instant};
19
20use common_base::readable_size::ReadableSize;
21use common_telemetry::{debug, info};
22use futures::AsyncWriteExt;
23use object_store::ObjectStore;
24use snafu::ResultExt;
25use store_api::storage::RegionId;
26
27use crate::access_layer::{
28 FilePathProvider, Metrics, RegionFilePathFactory, SstInfoArray, SstWriteRequest,
29 TempFileCleaner, WriteCachePathProvider, WriteType, new_fs_cache_store,
30};
31use crate::cache::file_cache::{FileCache, FileCacheRef, FileType, IndexKey, IndexValue};
32use crate::error::{self, Result};
33use crate::metrics::{
34 UPLOAD_BYTES_TOTAL, WRITE_CACHE_DOWNLOAD_BYTES_TOTAL, WRITE_CACHE_DOWNLOAD_ELAPSED,
35};
36use crate::sst::file::RegionFileId;
37use crate::sst::index::IndexerBuilderImpl;
38use crate::sst::index::intermediate::IntermediateManager;
39use crate::sst::index::puffin_manager::{PuffinManagerFactory, SstPuffinManager};
40use crate::sst::parquet::writer::ParquetWriter;
41use crate::sst::parquet::{SstInfo, WriteOptions};
42use crate::sst::{DEFAULT_WRITE_BUFFER_SIZE, DEFAULT_WRITE_CONCURRENCY};
43
44pub struct WriteCache {
48 file_cache: FileCacheRef,
50 puffin_manager_factory: PuffinManagerFactory,
52 intermediate_manager: IntermediateManager,
54}
55
56pub type WriteCacheRef = Arc<WriteCache>;
57
58impl WriteCache {
59 pub async fn new(
62 local_store: ObjectStore,
63 cache_capacity: ReadableSize,
64 ttl: Option<Duration>,
65 puffin_manager_factory: PuffinManagerFactory,
66 intermediate_manager: IntermediateManager,
67 ) -> Result<Self> {
68 let file_cache = Arc::new(FileCache::new(local_store, cache_capacity, ttl));
69 file_cache.recover(false).await;
70
71 Ok(Self {
72 file_cache,
73 puffin_manager_factory,
74 intermediate_manager,
75 })
76 }
77
78 pub async fn new_fs(
80 cache_dir: &str,
81 cache_capacity: ReadableSize,
82 ttl: Option<Duration>,
83 puffin_manager_factory: PuffinManagerFactory,
84 intermediate_manager: IntermediateManager,
85 ) -> Result<Self> {
86 info!("Init write cache on {cache_dir}, capacity: {cache_capacity}");
87
88 let local_store = new_fs_cache_store(cache_dir).await?;
89 Self::new(
90 local_store,
91 cache_capacity,
92 ttl,
93 puffin_manager_factory,
94 intermediate_manager,
95 )
96 .await
97 }
98
99 pub(crate) fn file_cache(&self) -> FileCacheRef {
101 self.file_cache.clone()
102 }
103
104 pub(crate) fn build_puffin_manager(&self) -> SstPuffinManager {
106 let store = self.file_cache.local_store();
107 let path_provider = WriteCachePathProvider::new(self.file_cache.clone());
108 self.puffin_manager_factory.build(store, path_provider)
109 }
110
111 pub(crate) async fn put_and_upload_sst(
113 &self,
114 data: &bytes::Bytes,
115 region_id: RegionId,
116 sst_info: &SstInfo,
117 upload_request: SstUploadRequest,
118 ) -> Result<Metrics> {
119 let file_id = sst_info.file_id;
120 let mut metrics = Metrics::new(WriteType::Flush);
121
122 let parquet_key = IndexKey::new(region_id, file_id, FileType::Parquet);
124
125 let cache_start = Instant::now();
127 let cache_path = self.file_cache.cache_file_path(parquet_key);
128 let store = self.file_cache.local_store();
129 let cleaner = TempFileCleaner::new(region_id, store.clone());
130 let write_res = store
131 .write(&cache_path, data.clone())
132 .await
133 .context(crate::error::OpenDalSnafu);
134 if let Err(e) = write_res {
135 cleaner.clean_by_file_id(file_id).await;
136 return Err(e);
137 }
138
139 metrics.write_batch = cache_start.elapsed();
140
141 let upload_start = Instant::now();
143 let region_file_id = RegionFileId::new(region_id, file_id);
144 let remote_path = upload_request
145 .dest_path_provider
146 .build_sst_file_path(region_file_id);
147
148 if let Err(e) = self
149 .upload(parquet_key, &remote_path, &upload_request.remote_store)
150 .await
151 {
152 self.remove(parquet_key).await;
154 return Err(e);
155 }
156
157 metrics.upload_parquet = upload_start.elapsed();
158 Ok(metrics)
159 }
160
161 pub(crate) fn intermediate_manager(&self) -> &IntermediateManager {
163 &self.intermediate_manager
164 }
165
166 pub(crate) async fn write_and_upload_sst(
168 &self,
169 write_request: SstWriteRequest,
170 upload_request: SstUploadRequest,
171 write_opts: &WriteOptions,
172 write_type: WriteType,
173 ) -> Result<(SstInfoArray, Metrics)> {
174 let region_id = write_request.metadata.region_id;
175
176 let store = self.file_cache.local_store();
177 let path_provider = WriteCachePathProvider::new(self.file_cache.clone());
178 let indexer = IndexerBuilderImpl {
179 build_type: write_request.op_type.into(),
180 metadata: write_request.metadata.clone(),
181 row_group_size: write_opts.row_group_size,
182 puffin_manager: self
183 .puffin_manager_factory
184 .build(store.clone(), path_provider.clone()),
185 intermediate_manager: self.intermediate_manager.clone(),
186 index_options: write_request.index_options,
187 inverted_index_config: write_request.inverted_index_config,
188 fulltext_index_config: write_request.fulltext_index_config,
189 bloom_filter_index_config: write_request.bloom_filter_index_config,
190 };
191
192 let cleaner = TempFileCleaner::new(region_id, store.clone());
193 let mut writer = ParquetWriter::new_with_object_store(
195 store.clone(),
196 write_request.metadata,
197 write_request.index_config,
198 indexer,
199 path_provider.clone(),
200 Metrics::new(write_type),
201 )
202 .await
203 .with_file_cleaner(cleaner);
204
205 let sst_info = match write_request.source {
206 either::Left(source) => {
207 writer
208 .write_all(source, write_request.max_sequence, write_opts)
209 .await?
210 }
211 either::Right(flat_source) => writer.write_all_flat(flat_source, write_opts).await?,
212 };
213 let mut metrics = writer.into_metrics();
214
215 if sst_info.is_empty() {
217 return Ok((sst_info, metrics));
218 }
219
220 let mut upload_tracker = UploadTracker::new(region_id);
221 let mut err = None;
222 let remote_store = &upload_request.remote_store;
223 for sst in &sst_info {
224 let parquet_key = IndexKey::new(region_id, sst.file_id, FileType::Parquet);
225 let parquet_path = upload_request
226 .dest_path_provider
227 .build_sst_file_path(RegionFileId::new(region_id, sst.file_id));
228 let start = Instant::now();
229 if let Err(e) = self.upload(parquet_key, &parquet_path, remote_store).await {
230 err = Some(e);
231 break;
232 }
233 metrics.upload_parquet += start.elapsed();
234 upload_tracker.push_uploaded_file(parquet_path);
235
236 if sst.index_metadata.file_size > 0 {
237 let puffin_key = IndexKey::new(region_id, sst.file_id, FileType::Puffin);
238 let puffin_path = upload_request
239 .dest_path_provider
240 .build_index_file_path(RegionFileId::new(region_id, sst.file_id));
241 let start = Instant::now();
242 if let Err(e) = self.upload(puffin_key, &puffin_path, remote_store).await {
243 err = Some(e);
244 break;
245 }
246 metrics.upload_puffin += start.elapsed();
247 upload_tracker.push_uploaded_file(puffin_path);
248 }
249 }
250
251 if let Some(err) = err {
252 upload_tracker
254 .clean(&sst_info, &self.file_cache, remote_store)
255 .await;
256 return Err(err);
257 }
258
259 Ok((sst_info, metrics))
260 }
261
262 pub(crate) async fn remove(&self, index_key: IndexKey) {
264 self.file_cache.remove(index_key).await
265 }
266
267 pub(crate) async fn download(
270 &self,
271 index_key: IndexKey,
272 remote_path: &str,
273 remote_store: &ObjectStore,
274 file_size: u64,
275 ) -> Result<()> {
276 if let Err(e) = self
277 .download_without_cleaning(index_key, remote_path, remote_store, file_size)
278 .await
279 {
280 let filename = index_key.to_string();
281 TempFileCleaner::clean_atomic_dir_files(&self.file_cache.local_store(), &[&filename])
282 .await;
283
284 return Err(e);
285 }
286 Ok(())
287 }
288
289 async fn download_without_cleaning(
290 &self,
291 index_key: IndexKey,
292 remote_path: &str,
293 remote_store: &ObjectStore,
294 file_size: u64,
295 ) -> Result<()> {
296 const DOWNLOAD_READER_CONCURRENCY: usize = 8;
297 const DOWNLOAD_READER_CHUNK_SIZE: ReadableSize = ReadableSize::mb(8);
298
299 let file_type = index_key.file_type;
300 let timer = WRITE_CACHE_DOWNLOAD_ELAPSED
301 .with_label_values(&[match file_type {
302 FileType::Parquet => "download_parquet",
303 FileType::Puffin => "download_puffin",
304 }])
305 .start_timer();
306
307 let reader = remote_store
308 .reader_with(remote_path)
309 .concurrent(DOWNLOAD_READER_CONCURRENCY)
310 .chunk(DOWNLOAD_READER_CHUNK_SIZE.as_bytes() as usize)
311 .await
312 .context(error::OpenDalSnafu)?
313 .into_futures_async_read(0..file_size)
314 .await
315 .context(error::OpenDalSnafu)?;
316
317 let cache_path = self.file_cache.cache_file_path(index_key);
318 let mut writer = self
319 .file_cache
320 .local_store()
321 .writer(&cache_path)
322 .await
323 .context(error::OpenDalSnafu)?
324 .into_futures_async_write();
325
326 let region_id = index_key.region_id;
327 let file_id = index_key.file_id;
328 let bytes_written =
329 futures::io::copy(reader, &mut writer)
330 .await
331 .context(error::DownloadSnafu {
332 region_id,
333 file_id,
334 file_type,
335 })?;
336 writer.close().await.context(error::DownloadSnafu {
337 region_id,
338 file_id,
339 file_type,
340 })?;
341
342 WRITE_CACHE_DOWNLOAD_BYTES_TOTAL.inc_by(bytes_written);
343
344 let elapsed = timer.stop_and_record();
345 debug!(
346 "Successfully download file '{}' to local '{}', file size: {}, region: {}, cost: {:?}s",
347 remote_path, cache_path, bytes_written, region_id, elapsed,
348 );
349
350 let index_value = IndexValue {
351 file_size: bytes_written as _,
352 };
353 self.file_cache.put(index_key, index_value).await;
354 Ok(())
355 }
356
357 pub(crate) async fn upload(
359 &self,
360 index_key: IndexKey,
361 upload_path: &str,
362 remote_store: &ObjectStore,
363 ) -> Result<()> {
364 let region_id = index_key.region_id;
365 let file_id = index_key.file_id;
366 let file_type = index_key.file_type;
367 let cache_path = self.file_cache.cache_file_path(index_key);
368
369 let start = Instant::now();
370 let cached_value = self
371 .file_cache
372 .local_store()
373 .stat(&cache_path)
374 .await
375 .context(error::OpenDalSnafu)?;
376 let reader = self
377 .file_cache
378 .local_store()
379 .reader(&cache_path)
380 .await
381 .context(error::OpenDalSnafu)?
382 .into_futures_async_read(0..cached_value.content_length())
383 .await
384 .context(error::OpenDalSnafu)?;
385
386 let mut writer = remote_store
387 .writer_with(upload_path)
388 .chunk(DEFAULT_WRITE_BUFFER_SIZE.as_bytes() as usize)
389 .concurrent(DEFAULT_WRITE_CONCURRENCY)
390 .await
391 .context(error::OpenDalSnafu)?
392 .into_futures_async_write();
393
394 let bytes_written =
395 futures::io::copy(reader, &mut writer)
396 .await
397 .context(error::UploadSnafu {
398 region_id,
399 file_id,
400 file_type,
401 })?;
402
403 writer.close().await.context(error::UploadSnafu {
405 region_id,
406 file_id,
407 file_type,
408 })?;
409
410 UPLOAD_BYTES_TOTAL.inc_by(bytes_written);
411
412 debug!(
413 "Successfully upload file to remote, region: {}, file: {}, upload_path: {}, cost: {:?}",
414 region_id,
415 file_id,
416 upload_path,
417 start.elapsed(),
418 );
419
420 let index_value = IndexValue {
421 file_size: bytes_written as _,
422 };
423 self.file_cache.put(index_key, index_value).await;
425
426 Ok(())
427 }
428}
429
430pub struct SstUploadRequest {
432 pub dest_path_provider: RegionFilePathFactory,
434 pub remote_store: ObjectStore,
436}
437
438pub(crate) struct UploadTracker {
440 region_id: RegionId,
442 files_uploaded: Vec<String>,
444}
445
446impl UploadTracker {
447 pub(crate) fn new(region_id: RegionId) -> Self {
449 Self {
450 region_id,
451 files_uploaded: Vec::new(),
452 }
453 }
454
455 pub(crate) fn push_uploaded_file(&mut self, path: String) {
457 self.files_uploaded.push(path);
458 }
459
460 pub(crate) async fn clean(
462 &self,
463 sst_info: &SstInfoArray,
464 file_cache: &FileCacheRef,
465 remote_store: &ObjectStore,
466 ) {
467 common_telemetry::info!(
468 "Start cleaning files on upload failure, region: {}, num_ssts: {}",
469 self.region_id,
470 sst_info.len()
471 );
472
473 for sst in sst_info {
475 let parquet_key = IndexKey::new(self.region_id, sst.file_id, FileType::Parquet);
476 file_cache.remove(parquet_key).await;
477
478 if sst.index_metadata.file_size > 0 {
479 let puffin_key = IndexKey::new(self.region_id, sst.file_id, FileType::Puffin);
480 file_cache.remove(puffin_key).await;
481 }
482 }
483
484 for file_path in &self.files_uploaded {
486 if let Err(e) = remote_store.delete(file_path).await {
487 common_telemetry::error!(e; "Failed to delete file {}", file_path);
488 }
489 }
490 }
491}
492
493#[cfg(test)]
494mod tests {
495 use common_test_util::temp_dir::create_temp_dir;
496 use object_store::ATOMIC_WRITE_DIR;
497 use store_api::region_request::PathType;
498
499 use super::*;
500 use crate::access_layer::OperationType;
501 use crate::cache::test_util::new_fs_store;
502 use crate::cache::{CacheManager, CacheStrategy};
503 use crate::error::InvalidBatchSnafu;
504 use crate::read::Source;
505 use crate::region::options::IndexOptions;
506 use crate::sst::parquet::reader::ParquetReaderBuilder;
507 use crate::test_util::TestEnv;
508 use crate::test_util::sst_util::{
509 assert_parquet_metadata_eq, new_batch_by_range, new_source, sst_file_handle_with_file_id,
510 sst_region_metadata,
511 };
512
513 #[tokio::test]
514 async fn test_write_and_upload_sst() {
515 let mut env = TestEnv::new().await;
518 let mock_store = env.init_object_store_manager();
519 let path_provider = RegionFilePathFactory::new("test".to_string(), PathType::Bare);
520
521 let local_dir = create_temp_dir("");
522 let local_store = new_fs_store(local_dir.path().to_str().unwrap());
523
524 let write_cache = env
525 .create_write_cache(local_store.clone(), ReadableSize::mb(10))
526 .await;
527
528 let metadata = Arc::new(sst_region_metadata());
530 let region_id = metadata.region_id;
531 let source = new_source(&[
532 new_batch_by_range(&["a", "d"], 0, 60),
533 new_batch_by_range(&["b", "f"], 0, 40),
534 new_batch_by_range(&["b", "h"], 100, 200),
535 ]);
536
537 let write_request = SstWriteRequest {
538 op_type: OperationType::Flush,
539 metadata,
540 source: either::Left(source),
541 storage: None,
542 max_sequence: None,
543 cache_manager: Default::default(),
544 index_options: IndexOptions::default(),
545 index_config: Default::default(),
546 inverted_index_config: Default::default(),
547 fulltext_index_config: Default::default(),
548 bloom_filter_index_config: Default::default(),
549 };
550
551 let upload_request = SstUploadRequest {
552 dest_path_provider: path_provider.clone(),
553 remote_store: mock_store.clone(),
554 };
555
556 let write_opts = WriteOptions {
557 row_group_size: 512,
558 ..Default::default()
559 };
560
561 let (mut sst_infos, _) = write_cache
563 .write_and_upload_sst(write_request, upload_request, &write_opts, WriteType::Flush)
564 .await
565 .unwrap();
566 let sst_info = sst_infos.remove(0);
567
568 let file_id = sst_info.file_id;
569 let sst_upload_path =
570 path_provider.build_sst_file_path(RegionFileId::new(region_id, file_id));
571 let index_upload_path =
572 path_provider.build_index_file_path(RegionFileId::new(region_id, file_id));
573
574 let key = IndexKey::new(region_id, file_id, FileType::Parquet);
576 assert!(write_cache.file_cache.contains_key(&key));
577
578 let remote_data = mock_store.read(&sst_upload_path).await.unwrap();
580 let cache_data = local_store
581 .read(&write_cache.file_cache.cache_file_path(key))
582 .await
583 .unwrap();
584 assert_eq!(remote_data.to_vec(), cache_data.to_vec());
585
586 let index_key = IndexKey::new(region_id, file_id, FileType::Puffin);
588 assert!(write_cache.file_cache.contains_key(&index_key));
589
590 let remote_index_data = mock_store.read(&index_upload_path).await.unwrap();
591 let cache_index_data = local_store
592 .read(&write_cache.file_cache.cache_file_path(index_key))
593 .await
594 .unwrap();
595 assert_eq!(remote_index_data.to_vec(), cache_index_data.to_vec());
596
597 let sst_index_key = IndexKey::new(region_id, file_id, FileType::Parquet);
599 write_cache.remove(sst_index_key).await;
600 assert!(!write_cache.file_cache.contains_key(&sst_index_key));
601 write_cache.remove(index_key).await;
602 assert!(!write_cache.file_cache.contains_key(&index_key));
603 }
604
605 #[tokio::test]
606 async fn test_read_metadata_from_write_cache() {
607 common_telemetry::init_default_ut_logging();
608 let mut env = TestEnv::new().await;
609 let data_home = env.data_home().display().to_string();
610 let mock_store = env.init_object_store_manager();
611
612 let local_dir = create_temp_dir("");
613 let local_path = local_dir.path().to_str().unwrap();
614 let local_store = new_fs_store(local_path);
615
616 let write_cache = env
618 .create_write_cache(local_store.clone(), ReadableSize::mb(10))
619 .await;
620 let cache_manager = Arc::new(
621 CacheManager::builder()
622 .write_cache(Some(write_cache.clone()))
623 .build(),
624 );
625
626 let metadata = Arc::new(sst_region_metadata());
628
629 let source = new_source(&[
630 new_batch_by_range(&["a", "d"], 0, 60),
631 new_batch_by_range(&["b", "f"], 0, 40),
632 new_batch_by_range(&["b", "h"], 100, 200),
633 ]);
634
635 let write_request = SstWriteRequest {
637 op_type: OperationType::Flush,
638 metadata,
639 source: either::Left(source),
640 storage: None,
641 max_sequence: None,
642 cache_manager: cache_manager.clone(),
643 index_options: IndexOptions::default(),
644 index_config: Default::default(),
645 inverted_index_config: Default::default(),
646 fulltext_index_config: Default::default(),
647 bloom_filter_index_config: Default::default(),
648 };
649 let write_opts = WriteOptions {
650 row_group_size: 512,
651 ..Default::default()
652 };
653 let upload_request = SstUploadRequest {
654 dest_path_provider: RegionFilePathFactory::new(data_home.clone(), PathType::Bare),
655 remote_store: mock_store.clone(),
656 };
657
658 let (mut sst_infos, _) = write_cache
659 .write_and_upload_sst(write_request, upload_request, &write_opts, WriteType::Flush)
660 .await
661 .unwrap();
662 let sst_info = sst_infos.remove(0);
663 let write_parquet_metadata = sst_info.file_metadata.unwrap();
664
665 let handle = sst_file_handle_with_file_id(sst_info.file_id, 0, 1000);
667 let builder = ParquetReaderBuilder::new(
668 data_home,
669 PathType::Bare,
670 handle.clone(),
671 mock_store.clone(),
672 )
673 .cache(CacheStrategy::EnableAll(cache_manager.clone()));
674 let reader = builder.build().await.unwrap();
675
676 assert_parquet_metadata_eq(write_parquet_metadata, reader.parquet_metadata());
678 }
679
680 #[tokio::test]
681 async fn test_write_cache_clean_tmp_files() {
682 common_telemetry::init_default_ut_logging();
683 let mut env = TestEnv::new().await;
684 let data_home = env.data_home().display().to_string();
685 let mock_store = env.init_object_store_manager();
686
687 let write_cache_dir = create_temp_dir("");
688 let write_cache_path = write_cache_dir.path().to_str().unwrap();
689 let write_cache = env
690 .create_write_cache_from_path(write_cache_path, ReadableSize::mb(10))
691 .await;
692
693 let cache_manager = Arc::new(
695 CacheManager::builder()
696 .write_cache(Some(write_cache.clone()))
697 .build(),
698 );
699
700 let metadata = Arc::new(sst_region_metadata());
702
703 let source = Source::Iter(Box::new(
705 [
706 Ok(new_batch_by_range(&["a", "d"], 0, 60)),
707 InvalidBatchSnafu {
708 reason: "Abort the writer",
709 }
710 .fail(),
711 ]
712 .into_iter(),
713 ));
714
715 let write_request = SstWriteRequest {
717 op_type: OperationType::Flush,
718 metadata,
719 source: either::Left(source),
720 storage: None,
721 max_sequence: None,
722 cache_manager: cache_manager.clone(),
723 index_options: IndexOptions::default(),
724 index_config: Default::default(),
725 inverted_index_config: Default::default(),
726 fulltext_index_config: Default::default(),
727 bloom_filter_index_config: Default::default(),
728 };
729 let write_opts = WriteOptions {
730 row_group_size: 512,
731 ..Default::default()
732 };
733 let upload_request = SstUploadRequest {
734 dest_path_provider: RegionFilePathFactory::new(data_home.clone(), PathType::Bare),
735 remote_store: mock_store.clone(),
736 };
737
738 write_cache
739 .write_and_upload_sst(write_request, upload_request, &write_opts, WriteType::Flush)
740 .await
741 .unwrap_err();
742 let atomic_write_dir = write_cache_dir.path().join(ATOMIC_WRITE_DIR);
743 let mut entries = tokio::fs::read_dir(&atomic_write_dir).await.unwrap();
744 let mut has_files = false;
745 while let Some(entry) = entries.next_entry().await.unwrap() {
746 if entry.file_type().await.unwrap().is_dir() {
747 continue;
748 }
749 has_files = true;
750 common_telemetry::warn!(
751 "Found remaining temporary file in atomic dir: {}",
752 entry.path().display()
753 );
754 }
755
756 assert!(!has_files);
757 }
758}