From 2e4c48ae7a1704850d53ad54195a7241a60bedf7 Mon Sep 17 00:00:00 2001 From: Zhenchi Date: Fri, 19 Jan 2024 10:57:07 +0800 Subject: [PATCH] fix(index): S3 `EntityTooSmall` error (#3192) * fix(index): S3 `EntityTooSmall` error Signed-off-by: Zhenchi * fix: config api Signed-off-by: Zhenchi --------- Signed-off-by: Zhenchi --- src/mito2/src/access_layer.rs | 3 ++ src/mito2/src/cache/write_cache.rs | 2 + src/mito2/src/compaction/twcs.rs | 7 ++++ src/mito2/src/config.rs | 3 ++ src/mito2/src/flush.rs | 7 ++++ src/mito2/src/sst/index.rs | 53 +++++++++++++++++++------ src/mito2/src/sst/index/creator.rs | 24 ++++++++++- src/mito2/src/sst/index/intermediate.rs | 6 +++ src/mito2/src/sst/index/store.rs | 23 ++++++++++- src/mito2/src/sst/parquet/reader.rs | 18 +++++++-- src/mito2/src/worker.rs | 8 +++- tests-integration/tests/http.rs | 1 + 12 files changed, 133 insertions(+), 22 deletions(-) diff --git a/src/mito2/src/access_layer.rs b/src/mito2/src/access_layer.rs index 634dfc2162e..630a27f4364 100644 --- a/src/mito2/src/access_layer.rs +++ b/src/mito2/src/access_layer.rs @@ -135,6 +135,7 @@ impl AccessLayer { let indexer = IndexerBuilder { create_inverted_index: request.create_inverted_index, mem_threshold_index_create: request.mem_threshold_index_create, + write_buffer_size: request.index_write_buffer_size, file_id, file_path: index_file_path, metadata: &request.metadata, @@ -183,6 +184,8 @@ pub(crate) struct SstWriteRequest { pub(crate) create_inverted_index: bool, /// The threshold of memory size to create inverted index. pub(crate) mem_threshold_index_create: Option, + /// The size of write buffer for index. + pub(crate) index_write_buffer_size: Option, } /// Creates a fs object store with atomic write dir. diff --git a/src/mito2/src/cache/write_cache.rs b/src/mito2/src/cache/write_cache.rs index b02ce8f7f3e..0840fa22fe6 100644 --- a/src/mito2/src/cache/write_cache.rs +++ b/src/mito2/src/cache/write_cache.rs @@ -110,6 +110,7 @@ impl WriteCache { let indexer = IndexerBuilder { create_inverted_index: write_request.create_inverted_index, mem_threshold_index_create: write_request.mem_threshold_index_create, + write_buffer_size: write_request.index_write_buffer_size, file_id, file_path: self.file_cache.cache_file_path(puffin_key), metadata: &write_request.metadata, @@ -281,6 +282,7 @@ mod tests { storage: None, create_inverted_index: true, mem_threshold_index_create: None, + index_write_buffer_size: None, cache_manager: Default::default(), }; diff --git a/src/mito2/src/compaction/twcs.rs b/src/mito2/src/compaction/twcs.rs index b0adc897b6a..0a7b727bb4e 100644 --- a/src/mito2/src/compaction/twcs.rs +++ b/src/mito2/src/compaction/twcs.rs @@ -314,6 +314,12 @@ impl TwcsCompactionTask { .inverted_index .mem_threshold_on_create .map(|m| m.as_bytes() as _); + let index_write_buffer_size = Some( + self.engine_config + .inverted_index + .write_buffer_size + .as_bytes() as usize, + ); let metadata = self.metadata.clone(); let sst_layer = self.sst_layer.clone(); @@ -334,6 +340,7 @@ impl TwcsCompactionTask { storage, create_inverted_index, mem_threshold_index_create, + index_write_buffer_size, }, &write_opts, ) diff --git a/src/mito2/src/config.rs b/src/mito2/src/config.rs index dbd8989e0ac..093b8c06d0e 100644 --- a/src/mito2/src/config.rs +++ b/src/mito2/src/config.rs @@ -252,6 +252,8 @@ pub struct InvertedIndexConfig { pub create_on_compaction: Mode, /// Whether to apply the index on query: automatically or never. pub apply_on_query: Mode, + /// Write buffer size for creating the index. + pub write_buffer_size: ReadableSize, /// Memory threshold for performing an external sort during index creation. /// `None` means all sorting will happen in memory. #[serde_as(as = "NoneAsEmptyString")] @@ -266,6 +268,7 @@ impl Default for InvertedIndexConfig { create_on_flush: Mode::Auto, create_on_compaction: Mode::Auto, apply_on_query: Mode::Auto, + write_buffer_size: ReadableSize::mb(8), mem_threshold_on_create: Some(ReadableSize::mb(64)), intermediate_path: String::new(), } diff --git a/src/mito2/src/flush.rs b/src/mito2/src/flush.rs index 6ee585e40c2..150b2fb1b68 100644 --- a/src/mito2/src/flush.rs +++ b/src/mito2/src/flush.rs @@ -321,6 +321,12 @@ impl RegionFlushTask { .inverted_index .mem_threshold_on_create .map(|m| m.as_bytes() as _); + let index_write_buffer_size = Some( + self.engine_config + .inverted_index + .write_buffer_size + .as_bytes() as usize, + ); // Flush to level 0. let write_request = SstWriteRequest { @@ -331,6 +337,7 @@ impl RegionFlushTask { storage: version.options.storage.clone(), create_inverted_index, mem_threshold_index_create, + index_write_buffer_size, }; let Some(sst_info) = self .access_layer diff --git a/src/mito2/src/sst/index.rs b/src/mito2/src/sst/index.rs index 3e18cabe62a..2c26f638952 100644 --- a/src/mito2/src/sst/index.rs +++ b/src/mito2/src/sst/index.rs @@ -52,10 +52,17 @@ impl Indexer { pub async fn update(&mut self, batch: &Batch) { if let Some(creator) = self.inner.as_mut() { if let Err(err) = creator.update(batch).await { - warn!( - err; "Failed to update index, skip creating index, region_id: {}, file_id: {}", - self.region_id, self.file_id, - ); + if cfg!(any(test, feature = "test")) { + panic!( + "Failed to update index, region_id: {}, file_id: {}, err: {}", + self.region_id, self.file_id, err + ); + } else { + warn!( + err; "Failed to update index, skip creating index, region_id: {}, file_id: {}", + self.region_id, self.file_id, + ); + } // Skip index creation if error occurs. self.inner = None; @@ -76,10 +83,17 @@ impl Indexer { return Some(byte_count); } Err(err) => { - warn!( - err; "Failed to create index, region_id: {}, file_id: {}", - self.region_id, self.file_id, - ); + if cfg!(any(test, feature = "test")) { + panic!( + "Failed to create index, region_id: {}, file_id: {}, err: {}", + self.region_id, self.file_id, err + ); + } else { + warn!( + err; "Failed to create index, region_id: {}, file_id: {}", + self.region_id, self.file_id, + ); + } } } } @@ -91,10 +105,17 @@ impl Indexer { pub async fn abort(&mut self) { if let Some(mut creator) = self.inner.take() { if let Err(err) = creator.abort().await { - warn!( - err; "Failed to abort index, region_id: {}, file_id: {}", - self.region_id, self.file_id, - ); + if cfg!(any(test, feature = "test")) { + panic!( + "Failed to abort index, region_id: {}, file_id: {}, err: {}", + self.region_id, self.file_id, err + ); + } else { + warn!( + err; "Failed to abort index, region_id: {}, file_id: {}", + self.region_id, self.file_id, + ); + } } } } @@ -103,6 +124,7 @@ impl Indexer { pub(crate) struct IndexerBuilder<'a> { pub(crate) create_inverted_index: bool, pub(crate) mem_threshold_index_create: Option, + pub(crate) write_buffer_size: Option, pub(crate) file_id: FileId, pub(crate) file_path: String, pub(crate) metadata: &'a RegionMetadataRef, @@ -147,7 +169,8 @@ impl<'a> IndexerBuilder<'a> { self.intermediate_manager, self.mem_threshold_index_create, row_group_size, - ); + ) + .with_buffer_size(self.write_buffer_size); Indexer { file_id: self.file_id, @@ -236,6 +259,7 @@ mod tests { let indexer = IndexerBuilder { create_inverted_index: true, mem_threshold_index_create: Some(1024), + write_buffer_size: None, file_id: FileId::random(), file_path: "test".to_string(), metadata: &metadata, @@ -254,6 +278,7 @@ mod tests { let indexer = IndexerBuilder { create_inverted_index: false, mem_threshold_index_create: Some(1024), + write_buffer_size: None, file_id: FileId::random(), file_path: "test".to_string(), metadata: &metadata, @@ -272,6 +297,7 @@ mod tests { let indexer = IndexerBuilder { create_inverted_index: true, mem_threshold_index_create: Some(1024), + write_buffer_size: None, file_id: FileId::random(), file_path: "test".to_string(), metadata: &metadata, @@ -290,6 +316,7 @@ mod tests { let indexer = IndexerBuilder { create_inverted_index: true, mem_threshold_index_create: Some(1024), + write_buffer_size: None, file_id: FileId::random(), file_path: "test".to_string(), metadata: &metadata, diff --git a/src/mito2/src/sst/index/creator.rs b/src/mito2/src/sst/index/creator.rs index c143bc9aacd..bb4ee19f224 100644 --- a/src/mito2/src/sst/index/creator.rs +++ b/src/mito2/src/sst/index/creator.rs @@ -113,6 +113,12 @@ impl SstIndexCreator { } } + /// Sets the write buffer size of the store. + pub fn with_buffer_size(mut self, write_buffer_size: Option) -> Self { + self.store = self.store.with_write_buffer_size(write_buffer_size); + self + } + /// Updates index with a batch of rows. /// Garbage will be cleaned up if failed to update. pub async fn update(&mut self, batch: &Batch) -> Result<()> { @@ -125,7 +131,14 @@ impl SstIndexCreator { if let Err(update_err) = self.do_update(batch).await { // clean up garbage if failed to update if let Err(err) = self.do_cleanup().await { - warn!(err; "Failed to clean up index creator, file_path: {}", self.file_path); + if cfg!(any(test, feature = "test")) { + panic!( + "Failed to clean up index creator, file_path: {}, err: {}", + self.file_path, err + ); + } else { + warn!(err; "Failed to clean up index creator, file_path: {}", self.file_path); + } } return Err(update_err); } @@ -146,7 +159,14 @@ impl SstIndexCreator { let finish_res = self.do_finish().await; // clean up garbage no matter finish successfully or not if let Err(err) = self.do_cleanup().await { - warn!(err; "Failed to clean up index creator, file_path: {}", self.file_path); + if cfg!(any(test, feature = "test")) { + panic!( + "Failed to clean up index creator, file_path: {}, err: {}", + self.file_path, err + ); + } else { + warn!(err; "Failed to clean up index creator, file_path: {}", self.file_path); + } } finish_res.map(|_| (self.stats.row_count(), self.stats.byte_count())) diff --git a/src/mito2/src/sst/index/intermediate.rs b/src/mito2/src/sst/index/intermediate.rs index ab10b0e6803..66177ded735 100644 --- a/src/mito2/src/sst/index/intermediate.rs +++ b/src/mito2/src/sst/index/intermediate.rs @@ -45,6 +45,12 @@ impl IntermediateManager { Ok(Self { store }) } + /// Set the write buffer size for the store. + pub fn with_buffer_size(mut self, write_buffer_size: Option) -> Self { + self.store = self.store.with_write_buffer_size(write_buffer_size); + self + } + /// Returns the store to access to intermediate files. pub(crate) fn store(&self) -> &InstrumentedStore { &self.store diff --git a/src/mito2/src/sst/index/store.rs b/src/mito2/src/sst/index/store.rs index a5e66c708e6..23893cc861f 100644 --- a/src/mito2/src/sst/index/store.rs +++ b/src/mito2/src/sst/index/store.rs @@ -30,12 +30,23 @@ use crate::error::{OpenDalSnafu, Result}; pub(crate) struct InstrumentedStore { /// The underlying object store. object_store: ObjectStore, + /// The size of the write buffer. + write_buffer_size: Option, } impl InstrumentedStore { /// Create a new `InstrumentedStore`. pub fn new(object_store: ObjectStore) -> Self { - Self { object_store } + Self { + object_store, + write_buffer_size: None, + } + } + + /// Set the size of the write buffer. + pub fn with_write_buffer_size(mut self, write_buffer_size: Option) -> Self { + self.write_buffer_size = write_buffer_size.filter(|&size| size > 0); + self } /// Returns an [`InstrumentedAsyncRead`] for the given path. @@ -67,7 +78,15 @@ impl InstrumentedStore { write_count: &'a IntCounter, flush_count: &'a IntCounter, ) -> Result> { - let writer = self.object_store.writer(path).await.context(OpenDalSnafu)?; + let writer = match self.write_buffer_size { + Some(size) => self + .object_store + .writer_with(path) + .buffer(size) + .await + .context(OpenDalSnafu)?, + None => self.object_store.writer(path).await.context(OpenDalSnafu)?, + }; Ok(InstrumentedAsyncWrite::new( writer, write_byte_count, diff --git a/src/mito2/src/sst/parquet/reader.rs b/src/mito2/src/sst/parquet/reader.rs index 6ba5d63af9d..af1fcdc68a1 100644 --- a/src/mito2/src/sst/parquet/reader.rs +++ b/src/mito2/src/sst/parquet/reader.rs @@ -293,9 +293,21 @@ impl ParquetReaderBuilder { if self.file_handle.meta().inverted_index_available() { match index_applier.apply(self.file_handle.file_id()).await { Ok(row_groups) => row_group_ids = row_groups, - Err(err) => warn!( - err; "Failed to apply index, region_id: {}, file_id: {}", - self.file_handle.region_id(), self.file_handle.file_id()), + Err(err) => { + if cfg!(any(test, feature = "test")) { + panic!( + "Failed to apply index, region_id: {}, file_id: {}, err: {}", + self.file_handle.region_id(), + self.file_handle.file_id(), + err + ); + } else { + warn!( + err; "Failed to apply index, region_id: {}, file_id: {}", + self.file_handle.region_id(), self.file_handle.file_id() + ); + } + } } } } diff --git a/src/mito2/src/worker.rs b/src/mito2/src/worker.rs index 4e1ad569777..3f960f416a7 100644 --- a/src/mito2/src/worker.rs +++ b/src/mito2/src/worker.rs @@ -123,7 +123,9 @@ impl WorkerGroup { config.global_write_buffer_size.as_bytes() as usize, )); let intermediate_manager = - IntermediateManager::init_fs(&config.inverted_index.intermediate_path).await?; + IntermediateManager::init_fs(&config.inverted_index.intermediate_path) + .await? + .with_buffer_size(Some(config.inverted_index.write_buffer_size.as_bytes() as _)); let scheduler = Arc::new(LocalScheduler::new(config.max_background_jobs)); let write_cache = write_cache_from_config( &config, @@ -233,7 +235,9 @@ impl WorkerGroup { }); let scheduler = Arc::new(LocalScheduler::new(config.max_background_jobs)); let intermediate_manager = - IntermediateManager::init_fs(&config.inverted_index.intermediate_path).await?; + IntermediateManager::init_fs(&config.inverted_index.intermediate_path) + .await? + .with_buffer_size(Some(config.inverted_index.write_buffer_size.as_bytes() as _)); let write_cache = write_cache_from_config( &config, object_store_manager.clone(), diff --git a/tests-integration/tests/http.rs b/tests-integration/tests/http.rs index 8eb15c2f87f..7182c72c19b 100644 --- a/tests-integration/tests/http.rs +++ b/tests-integration/tests/http.rs @@ -777,6 +777,7 @@ allow_stale_entries = false create_on_flush = "auto" create_on_compaction = "auto" apply_on_query = "auto" +write_buffer_size = "8MiB" mem_threshold_on_create = "64.0MiB" intermediate_path = ""