fix(mito): decode cached manifest payloads (#10179)

* test(mito2): cover compressed manifest cache recovery

Signed-off-by: WenyXu <wenymedia@gmail.com>

* fix(mito2): decode cached manifest object payloads

Signed-off-by: WenyXu <wenymedia@gmail.com>

* fix(mito2): address manifest cache review feedback

Signed-off-by: WenyXu <wenymedia@gmail.com>

---------

Signed-off-by: WenyXu <wenymedia@gmail.com>
This commit is contained in:
Weny Xu
2026-11-05 18:19:50 +00:00
committed by GitHub
parent 8e4aa4c86f
commit 4263f5fc01
7 changed files with 435 additions and 49 deletions
+2 -2
View File
@@ -237,7 +237,7 @@ impl ManifestCache {
join_path(MANIFEST_DIR, key)
}
/// Gets a manifest file from cache.
/// Gets the encoded object-store payload from cache.
/// Returns the file data if found in cache, None otherwise.
pub(crate) async fn get_file(&self, key: &str) -> Option<Vec<u8>> {
if self.get(key).await.is_none() {
@@ -259,7 +259,7 @@ impl ManifestCache {
}
}
/// Puts a manifest file into cache.
/// Caches the object-store payload verbatim, before decompression.
pub(crate) async fn put_file(&self, key: String, data: Vec<u8>) {
let cache_file_path = self.cache_file_path(&key);
+63
View File
@@ -1366,3 +1366,66 @@ async fn test_region_hook_on_open() {
assert_eq!(hook.closed_count.load(Ordering::Relaxed), 1);
assert_eq!(hook.dropped_count.load(Ordering::Relaxed), 0);
}
#[tokio::test]
async fn test_reopen_region_with_compressed_manifest_cache() {
let mut env = TestEnv::new().await;
let engine = env
.create_engine(MitoConfig {
enable_write_cache: true,
compress_manifest: true,
manifest_checkpoint_distance: 1,
..Default::default()
})
.await;
let region_id = RegionId::new(1, 1);
let request = CreateRequestBuilder::new().build();
let table_dir = request.table_dir.clone();
let schema = rows_schema(&request);
engine
.handle_request(region_id, RegionRequest::Create(request))
.await
.unwrap();
for i in 0..3 {
put_rows(
&engine,
region_id,
Rows {
schema: schema.clone(),
rows: build_rows(i * 3, (i + 1) * 3),
},
)
.await;
flush_region(&engine, region_id, None).await;
}
let region = engine.get_region(region_id).unwrap();
let mut manager = region.manifest_ctx.manifest_manager.write().await;
tokio::time::timeout(
Duration::from_secs(10),
manager.wait_for_pending_checkpoint(),
)
.await
.unwrap();
assert!(manager.checkpointer().last_checkpoint_version() > 0);
drop(manager);
let stream = engine
.scan_to_stream(region_id, ScanRequest::default())
.await
.unwrap();
let expected = RecordBatches::try_collect(stream)
.await
.unwrap()
.pretty_print()
.unwrap();
reopen_region(&engine, region_id, table_dir, true, HashMap::new()).await;
let stream = engine
.scan_to_stream(region_id, ScanRequest::default())
.await
.unwrap();
let actual = RecordBatches::try_collect(stream)
.await
.unwrap()
.pretty_print()
.unwrap();
assert_eq!(expected, actual);
}
+239
View File
@@ -846,4 +846,243 @@ mod tests {
.unwrap();
assert!(entries.is_empty());
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
enum CacheSource {
WriteThrough,
RemoteFill,
LegacyDecoded,
Truncated,
}
#[derive(Debug, Clone, Copy)]
enum PayloadKind {
Delta,
Checkpoint { checksum: bool },
}
// Exercise both representations left by older versions, as well as new writes
// and remote fills. Truncation only applies to compression envelopes.
#[rstest::rstest]
#[case::write_through_uncompressed(CacheSource::WriteThrough, CompressionType::Uncompressed)]
#[case::write_through_gzip(CacheSource::WriteThrough, CompressionType::Gzip)]
#[case::write_through_zstd(CacheSource::WriteThrough, CompressionType::Zstd)]
#[case::write_through_xz(CacheSource::WriteThrough, CompressionType::Xz)]
#[case::write_through_bzip2(CacheSource::WriteThrough, CompressionType::Bzip2)]
#[case::remote_fill_uncompressed(CacheSource::RemoteFill, CompressionType::Uncompressed)]
#[case::remote_fill_gzip(CacheSource::RemoteFill, CompressionType::Gzip)]
#[case::remote_fill_zstd(CacheSource::RemoteFill, CompressionType::Zstd)]
#[case::remote_fill_xz(CacheSource::RemoteFill, CompressionType::Xz)]
#[case::remote_fill_bzip2(CacheSource::RemoteFill, CompressionType::Bzip2)]
#[case::legacy_decoded_uncompressed(CacheSource::LegacyDecoded, CompressionType::Uncompressed)]
#[case::legacy_decoded_gzip(CacheSource::LegacyDecoded, CompressionType::Gzip)]
#[case::legacy_decoded_zstd(CacheSource::LegacyDecoded, CompressionType::Zstd)]
#[case::legacy_decoded_xz(CacheSource::LegacyDecoded, CompressionType::Xz)]
#[case::legacy_decoded_bzip2(CacheSource::LegacyDecoded, CompressionType::Bzip2)]
#[case::truncated_gzip(CacheSource::Truncated, CompressionType::Gzip)]
#[case::truncated_zstd(CacheSource::Truncated, CompressionType::Zstd)]
#[case::truncated_xz(CacheSource::Truncated, CompressionType::Xz)]
#[case::truncated_bzip2(CacheSource::Truncated, CompressionType::Bzip2)]
#[tokio::test]
async fn test_manifest_cache_payload(
#[case] source: CacheSource,
#[case] compression: CompressionType,
#[values(PayloadKind::Delta, PayloadKind::Checkpoint { checksum: false }, PayloadKind::Checkpoint { checksum: true })]
payload: PayloadKind,
) {
use std::sync::atomic::{AtomicUsize, Ordering};
use common_base::readable_size::ReadableSize;
use object_store::layers::mock::MockLayerBuilder;
let remote_dir = create_temp_dir("manifest-cache-remote");
let cache_dir = create_temp_dir("manifest-cache-local");
let remote =
ObjectStore::new(Fs::default().root(remote_dir.path().to_str().unwrap())).unwrap();
let local =
ObjectStore::new(Fs::default().root(cache_dir.path().to_str().unwrap())).unwrap();
let reads = Arc::new(AtomicUsize::new(0));
let observed = reads.clone();
let remote = remote.layer(
MockLayerBuilder::default()
.reader_factory(Arc::new(move |_, _, reader| {
observed.fetch_add(1, Ordering::Relaxed);
reader
}))
.build()
.unwrap(),
);
let mut cache = ManifestCache::new(local.clone(), ReadableSize::mb(32), None, true).await;
let mut store = ManifestObjectStore::new(
"region/manifest/",
remote.clone(),
compression,
Default::default(),
(source != CacheSource::RemoteFill).then(|| cache.clone()),
);
let data = br#"{"manifest":"persistent cache payload"}"#;
let key = gen_path(
"region/manifest/",
&if matches!(payload, PayloadKind::Checkpoint { .. }) {
checkpoint_file(3)
} else {
delta_file(3)
},
compression,
);
if matches!(payload, PayloadKind::Checkpoint { .. }) {
store.save_checkpoint(3, data).await.unwrap();
} else {
store.save(3, data, false).await.unwrap();
}
let encoded = remote.read(&key).await.unwrap().to_vec();
if source == CacheSource::LegacyDecoded {
cache.put_file(key.clone(), data.to_vec()).await;
} else if source == CacheSource::Truncated {
cache
.put_file(key.clone(), encoded[..encoded.len() / 2].to_vec())
.await;
}
drop(store);
for iteration in 0..2 {
drop(cache);
cache = ManifestCache::new(local.clone(), ReadableSize::mb(32), None, true).await;
if source != CacheSource::RemoteFill || iteration > 0 {
assert!(
cache.get(&key).await.is_some(),
"cache recovery missed {key}"
);
}
store = ManifestObjectStore::new(
"region/manifest/",
remote.clone(),
compression,
Default::default(),
Some(cache.clone()),
);
let before = reads.load(Ordering::Relaxed);
let actual = if let PayloadKind::Checkpoint { checksum } = payload {
let metadata = CheckpointMetadata {
size: data.len(),
version: 3,
checksum: checksum.then(|| checkpoint_checksum(data)),
extend_metadata: Default::default(),
};
store
.checkpoint_storage
.load_checkpoint(metadata)
.await
.unwrap()
.unwrap()
} else {
let mut deltas = store.fetch_manifests(3, 4).await.unwrap();
assert_eq!(1, deltas.len());
deltas.remove(0)
};
assert_eq!((3, data.to_vec()), actual);
// A successful roundtrip alone could hide a decoded refill that causes
// repeated cache misses. Check both the bytes and remote read count.
assert_eq!(encoded, cache.get_file(&key).await.unwrap());
let must_refill = iteration == 0
&& (source == CacheSource::RemoteFill
|| (compression != CompressionType::Uncompressed
&& (source == CacheSource::LegacyDecoded
|| source == CacheSource::Truncated)));
assert_eq!(
usize::from(must_refill),
reads.load(Ordering::Relaxed) - before
);
drop(store);
}
}
#[rstest::rstest]
#[case::checksum(false)]
#[case::decompression(true)]
#[tokio::test]
async fn test_manifest_cache_corrupted_remote(#[case] invalid_encoding: bool) {
use crate::error::Error;
use common_base::readable_size::ReadableSize;
let remote_dir = create_temp_dir("manifest-cache-remote");
let cache_dir = create_temp_dir("manifest-cache-local");
let remote =
ObjectStore::new(Fs::default().root(remote_dir.path().to_str().unwrap())).unwrap();
let local =
ObjectStore::new(Fs::default().root(cache_dir.path().to_str().unwrap())).unwrap();
let cache = ManifestCache::new(local, ReadableSize::mb(32), None, true).await;
let mut store = ManifestObjectStore::new(
"region/",
remote.clone(),
CompressionType::Gzip,
Default::default(),
Some(cache.clone()),
);
store.save_checkpoint(3, b"checkpoint").await.unwrap();
let key = gen_path("region/", &checkpoint_file(3), CompressionType::Gzip);
let corrupted = if invalid_encoding {
b"invalid gzip".to_vec()
} else {
CompressionType::Gzip
.encode(b"other".as_slice())
.await
.unwrap()
};
remote.write(&key, corrupted).await.unwrap();
// A legacy decoded cache entry must not hide the remote error on refill.
cache.put_file(key.clone(), b"checkpoint".to_vec()).await;
let error = store.load_last_checkpoint().await.unwrap_err();
if invalid_encoding {
assert!(matches!(error, Error::DecompressObject { .. }), "{error}");
} else {
assert!(matches!(error, Error::ChecksumMismatch { .. }), "{error}");
}
assert!(cache.get(&key).await.is_none());
}
#[tokio::test]
async fn test_manifest_cache_uncompressed_fallback() {
use common_base::readable_size::ReadableSize;
let remote_dir = create_temp_dir("manifest-cache-remote");
let cache_dir = create_temp_dir("manifest-cache-local");
let remote =
ObjectStore::new(Fs::default().root(remote_dir.path().to_str().unwrap())).unwrap();
let local =
ObjectStore::new(Fs::default().root(cache_dir.path().to_str().unwrap())).unwrap();
let cache = ManifestCache::new(local.clone(), ReadableSize::mb(32), None, true).await;
let store = ManifestObjectStore::new(
"region/",
remote.clone(),
CompressionType::Uncompressed,
Default::default(),
Some(cache.clone()),
);
store.save_checkpoint(3, b"checkpoint").await.unwrap();
drop(store);
drop(cache);
let cache = ManifestCache::new(local, ReadableSize::mb(32), None, true).await;
assert!(
cache
.get("region/00000000000000000003.checkpoint")
.await
.is_some()
);
// Keep _last_checkpoint, but remove its payload to prove the fallback hit.
remote
.delete("region/00000000000000000003.checkpoint")
.await
.unwrap();
let mut store = ManifestObjectStore::new(
"region/",
remote,
CompressionType::Gzip,
Default::default(),
Some(cache),
);
assert_eq!(
Some((3, b"checkpoint".to_vec())),
store.load_last_checkpoint().await.unwrap()
);
}
}
+26 -21
View File
@@ -27,7 +27,7 @@ use crate::error::{
CompressObjectSnafu, DecompressObjectSnafu, OpenDalSnafu, Result, SerdeJsonSnafu, Utf8Snafu,
};
use crate::manifest::storage::size_tracker::Tracker;
use crate::manifest::storage::utils::{get_from_cache, put_to_cache, write_and_put_cache};
use crate::manifest::storage::utils::{get_from_cache_decoded, write_and_put_cache};
use crate::manifest::storage::{
FALL_BACK_COMPRESS_TYPE, LAST_CHECKPOINT_FILE, checkpoint_checksum, checkpoint_file, gen_path,
verify_checksum,
@@ -102,7 +102,9 @@ impl<T: Tracker> CheckpointStorage<T> {
let path = self.checkpoint_file_path(version);
// Try to get from cache first
if let Some(data) = get_from_cache(self.manifest_cache.as_ref(), &path).await {
if let Some(data) =
get_from_cache_decoded(self.manifest_cache.as_ref(), &path, self.compress_type).await
{
verify_checksum(&data, metadata.checksum)?;
self.record_cached_checkpoint_size(version, &path).await;
return Ok(Some((version, data)));
@@ -113,19 +115,21 @@ impl<T: Tracker> CheckpointStorage<T> {
let checkpoint_data = match self.object_store.read(&path).await {
Ok(checkpoint) => {
let checkpoint_size = checkpoint.len();
let decompress_data =
self.compress_type
.decode(checkpoint)
.await
.with_context(|_| DecompressObjectSnafu {
compress_type: self.compress_type,
path: path.clone(),
})?;
let decompress_data = self
.compress_type
.decode(checkpoint.clone())
.await
.with_context(|_| DecompressObjectSnafu {
compress_type: self.compress_type,
path: path.clone(),
})?;
verify_checksum(&decompress_data, metadata.checksum)?;
// set the checkpoint size
self.size_tracker.record(version, checkpoint_size as u64);
// Add to cache
put_to_cache(self.manifest_cache.as_ref(), path, &decompress_data).await;
if let Some(cache) = &self.manifest_cache {
cache.put_file(path, checkpoint.to_vec()).await;
}
Ok(Some(decompress_data))
}
Err(e) => {
@@ -142,8 +146,12 @@ impl<T: Tracker> CheckpointStorage<T> {
);
// Try to get fallback from cache first
if let Some(data) =
get_from_cache(self.manifest_cache.as_ref(), &fall_back_path).await
if let Some(data) = get_from_cache_decoded(
self.manifest_cache.as_ref(),
&fall_back_path,
FALL_BACK_COMPRESS_TYPE,
)
.await
{
verify_checksum(&data, metadata.checksum)?;
self.record_cached_checkpoint_size(version, &fall_back_path)
@@ -155,7 +163,7 @@ impl<T: Tracker> CheckpointStorage<T> {
Ok(checkpoint) => {
let checkpoint_size = checkpoint.len();
let decompress_data = FALL_BACK_COMPRESS_TYPE
.decode(checkpoint)
.decode(checkpoint.clone())
.await
.with_context(|_| DecompressObjectSnafu {
compress_type: FALL_BACK_COMPRESS_TYPE,
@@ -164,12 +172,9 @@ impl<T: Tracker> CheckpointStorage<T> {
verify_checksum(&decompress_data, metadata.checksum)?;
self.size_tracker.record(version, checkpoint_size as u64);
// Add fallback to cache
put_to_cache(
self.manifest_cache.as_ref(),
fall_back_path,
&decompress_data,
)
.await;
if let Some(cache) = &self.manifest_cache {
cache.put_file(fall_back_path, checkpoint.to_vec()).await;
}
Ok(Some(decompress_data))
}
Err(e) if e.kind() == ErrorKind::NotFound => Ok(None),
@@ -220,7 +225,7 @@ impl<T: Tracker> CheckpointStorage<T> {
Ok(Some(checkpoint_metadata))
}
/// A cache hit contains decoded bytes; accounting failures must not prevent recovery.
/// Records the remote object size; accounting failures must not prevent recovery.
async fn record_cached_checkpoint_size(&self, version: ManifestVersion, path: &str) {
match self.object_store.stat(path).await {
Ok(metadata) => self.size_tracker.record(version, metadata.content_length()),
+7 -5
View File
@@ -31,7 +31,7 @@ use crate::error::{
};
use crate::manifest::storage::size_tracker::Tracker;
use crate::manifest::storage::utils::{
get_from_cache, put_to_cache, sort_manifests, write_and_put_cache,
get_from_cache_decoded, sort_manifests, write_and_put_cache,
};
use crate::manifest::storage::{
FETCH_MANIFEST_PARALLELISM, delta_file, file_compress_type, file_version, gen_path,
@@ -206,8 +206,9 @@ impl<T: Tracker> DeltaStorage<T> {
let _permit = semaphore.acquire().await.unwrap();
let cache_key = entry.path();
let compress_type = file_compress_type(entry.name());
// Try to get from cache first
if let Some(data) = get_from_cache(self.manifest_cache.as_ref(), cache_key).await {
if let Some(data) = get_from_cache_decoded(self.manifest_cache.as_ref(), cache_key, compress_type).await {
let size = entry.metadata().content_length();
if size > 0 {
self.delta_tracker.record(*v, size);
@@ -221,7 +222,6 @@ impl<T: Tracker> DeltaStorage<T> {
}
// Fetch from remote object store
let compress_type = file_compress_type(entry.name());
let bytes = match self.object_store.read(entry.path()).await {
Ok(bytes) => bytes,
Err(error) if error.kind() == ErrorKind::NotFound => {
@@ -234,7 +234,7 @@ impl<T: Tracker> DeltaStorage<T> {
};
self.delta_tracker.record(*v, bytes.len() as u64);
let data = compress_type
.decode(bytes)
.decode(bytes.clone())
.await
.context(DecompressObjectSnafu {
compress_type,
@@ -242,7 +242,9 @@ impl<T: Tracker> DeltaStorage<T> {
})?;
// Add to cache
put_to_cache(self.manifest_cache.as_ref(), cache_key.to_string(), &data).await;
if let Some(cache) = &self.manifest_cache {
cache.put_file(cache_key.to_string(), bytes.to_vec()).await;
}
Ok((*v, data))
});
+29 -21
View File
@@ -12,6 +12,8 @@
// See the License for the specific language governing permissions and
// limitations under the License.
use common_datasource::compression::CompressionType;
use common_telemetry::debug;
use object_store::{Entry, ObjectStore};
use snafu::ResultExt;
use store_api::ManifestVersion;
@@ -21,19 +23,29 @@ use crate::error::{OpenDalSnafu, Result};
const CHUNK_SIZE: usize = 64 * 1024 * 1024; // 64MB
/// Gets a manifest file from cache.
/// Returns the file data if found in cache, None otherwise.
pub(crate) async fn get_from_cache(cache: Option<&ManifestCache>, key: &str) -> Option<Vec<u8>> {
/// Gets and decodes the object payload cached under `key`.
///
/// Older remote fills cached decoded content. If decoding fails, discard the
/// entry and let the caller refill it from object storage. Remote decode and
/// checksum errors are still returned by the caller.
pub(crate) async fn get_from_cache_decoded(
cache: Option<&ManifestCache>,
key: &str,
compress_type: CompressionType,
) -> Option<Vec<u8>> {
let cache = cache?;
cache.get_file(key).await
}
/// Puts a manifest file into cache.
pub(crate) async fn put_to_cache(cache: Option<&ManifestCache>, key: String, data: &[u8]) {
let Some(cache) = cache else {
return;
};
cache.put_file(key, data.to_vec()).await
let raw = cache.get_file(key).await?;
match compress_type.decode(raw.as_slice()).await {
Ok(data) => Some(data),
Err(error) => {
debug!(
"Failed to decode cached manifest {}, removing entry: {}",
key, error
);
cache.remove(key).await;
None
}
}
}
/// Removes a manifest file from cache.
@@ -44,19 +56,15 @@ pub(crate) async fn remove_from_cache(cache: Option<&ManifestCache>, key: &str)
cache.remove(key).await
}
/// Writes data to object store and puts it into cache.
/// Writes data to object store and caches the same encoded payload.
pub(crate) async fn write_and_put_cache(
object_store: &ObjectStore,
cache: Option<&ManifestCache>,
path: &str,
data: Vec<u8>,
) -> Result<()> {
// Clone data for cache before writing, only if cache is enabled.
let cache_data = if cache.is_some() {
Some(data.clone())
} else {
None
};
// Clone only when caching is enabled; retain the owner for the cache write.
let cache_data = cache.map(|cache| (cache, data.clone()));
// Write to object store
object_store
@@ -66,8 +74,8 @@ pub(crate) async fn write_and_put_cache(
.context(OpenDalSnafu)?;
// Put to cache if we cloned the data
if let Some(data) = cache_data {
put_to_cache(cache, path.to_string(), &data).await;
if let Some((cache, data)) = cache_data {
cache.put_file(path.to_string(), data).await;
}
Ok(())
@@ -1072,3 +1072,72 @@ impl oio::List for ZeroSizeLister {
Ok(Some(oio::Entry::new(&path, builder.build())))
}
}
#[rstest::rstest]
#[case(CompressionType::Uncompressed)]
#[case(CompressionType::Gzip)]
#[case(CompressionType::Zstd)]
#[tokio::test]
async fn manifest_cache_write_through_survives_reopen(#[case] compression: CompressionType) {
use common_base::readable_size::ReadableSize;
use common_test_util::temp_dir::create_temp_dir;
use object_store::ObjectStore;
use object_store::services::Fs;
use crate::cache::manifest_cache::ManifestCache;
use crate::manifest::manager::RegionManifestOptions;
use crate::sst::FormatType;
let remote_dir = create_temp_dir("manifest-manager-remote");
let cache_dir = create_temp_dir("manifest-manager-cache");
let remote = ObjectStore::new(Fs::default().root(remote_dir.path().to_str().unwrap())).unwrap();
let local = ObjectStore::new(Fs::default().root(cache_dir.path().to_str().unwrap())).unwrap();
let cache = ManifestCache::new(local.clone(), ReadableSize::mb(32), None, true).await;
// Use a relative prefix, matching production keys and the recovered index.
let mut options = RegionManifestOptions {
manifest_dir: "region/".to_string(),
object_store: remote,
compress_type: compression,
checkpoint_distance: 3,
remove_file_options: Default::default(),
manifest_cache: Some(cache.clone()),
};
let mut writer = RegionManifestManager::new(
Arc::new(basic_region_metadata()),
0,
options.clone(),
FormatType::PrimaryKey,
false,
&Default::default(),
)
.await
.unwrap();
for _ in 0..7 {
writer.update(nop_action(), false).await.unwrap();
wait_for_manifest_checkpoint(&writer).await;
}
let key = gen_path(
"region/",
&checkpoint_file(writer.checkpointer().last_checkpoint_version()),
compression,
);
assert!(cache.get(&key).await.is_some());
let version = writer.manifest().manifest_version;
let usage = writer.manifest_usage();
drop(writer);
drop(cache);
// Release the final old cache handle before recovering a new index.
drop(options.manifest_cache.take());
let cache = ManifestCache::new(local, ReadableSize::mb(32), None, true).await;
assert!(
cache.get(&key).await.is_some(),
"cache recovery missed {key}"
);
options.manifest_cache = Some(cache);
let reopened = RegionManifestManager::open(options, &Default::default())
.await
.unwrap()
.unwrap();
assert_eq!(version, reopened.manifest().manifest_version);
assert_eq!(usage, reopened.manifest_usage());
}