andygrove commented on code in PR #6031: URL: https://github.com/apache/datafusion-comet/pull/6031#discussion_r4053636286
########## native/core/src/parquet/objectstore/retry.rs: ########## @@ -0,0 +1,628 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! Correctness safety net for the scope-hint SPI. +//! +//! # Why +//! +//! `CometS3ScopedCredentialProvider::getPolicyLocationsFor` is *advisory*: vendors are +//! encouraged to report narrower scopes than the policy grants, so a request outside the +//! reported scope can still legitimately fail with 403 at S3. The scope hint is a latency +//! optimization that lets Comet share a bridge across paths inside one scope, saving one +//! JVM round-trip per get. S3 itself remains authoritative. +//! +//! This wrapper is the mechanism that keeps that split honest: on a 403 from a cached +//! (scope-bound) store, we invoke a caller-provided rebuild closure exactly once, passing +//! it the `Path` of the request that failed. The rebuild constructs a fresh bridge bound +//! to that path, re-fires the SPI against it (so the vendor answers `getPolicyLocationsFor` +//! *for the actual failing request*), and appends the resulting scope entry to the outer +//! `object_store` registry cache — alongside any pre-existing entries, not replacing them. +//! We then retry the same operation against the rebuilt store; a second 403 propagates +//! unchanged. Because the new entry is inserted with the vendor's fresh (narrower) scope +//! rather than an all-covering catchall, disjoint scoped stores on the same bucket can +//! continue to coexist after a 403 recovery. +//! +//! # Non-goals +//! +//! - Not a generic retry policy. Only `Error::PermissionDenied` (403) is intercepted; every +//! other error (including 401/`Unauthenticated`) passes through as-is. +//! - Not an infinite retry. Exactly one rebuild per wrapper instance, exactly one retry per +//! operation. +//! - Not a stream-level retry. `list`, `list_with_offset`, and `delete_stream` return +//! `BoxStream`s whose per-item errors are surfaced as-is; wrapping them would require +//! materializing the stream. Comet's parquet path first hits 403 at `get_opts`/`get_ranges` +//! which are covered; the rebuild there populates the shared cache and subsequent stream +//! requests use the newly-inserted scope entry. +//! - Not applied to `put_multipart_opts`: a partial multipart upload cannot be transparently +//! retried, so a 403 mid-upload propagates for the caller to handle. + +use std::fmt; +use std::ops::Range; +use std::sync::Arc; + +use async_trait::async_trait; +use bytes::Bytes; +use futures::stream::BoxStream; +use log::debug; +use object_store::path::Path; +use object_store::{ + CopyOptions, Error, GetOptions, GetResult, ListResult, MultipartUpload, ObjectMeta, + ObjectStore, PutMultipartOptions, PutOptions, PutPayload, PutResult, RenameOptions, Result, +}; +use once_cell::sync::OnceCell; + +/// Rebuild function contract: construct a fresh backing `ObjectStore` scoped for the failing +/// request and register it in the outer `object_store` registry cache. Called at most once +/// per wrapper. +/// +/// The `Option<&Path>` argument is the location that triggered the 403 (source path for +/// copy/rename). The rebuild closure passes this into the fresh bridge so the SPI's +/// `getPolicyLocationsFor` reflects the actual failing request rather than the path baked in +/// at construction time. `None` is reserved for callers that intercept a 403 without a path +/// (none of the current retry sites) and lets the closure fall back to its pre-baked path. +pub type RebuildFn = + Arc<dyn Fn(Option<&Path>) -> Result<Arc<dyn ObjectStore>> + Send + Sync + 'static>; + +/// Wraps an `Arc<dyn ObjectStore>` so a single 403 rebuilds the store once and retries. +/// +/// See the module-level doc-comment for rationale and non-goals. +pub struct RetryOn403ObjectStore { + inner: Arc<dyn ObjectStore>, + rebuild: RebuildFn, + /// Populated on the first 403 we successfully recover from. Cached so a subsequent + /// request against this same wrapper skips straight to the rebuilt store, and a 403 there + /// is treated as authoritative. + rebuilt: OnceCell<Arc<dyn ObjectStore>>, +} + +impl RetryOn403ObjectStore { + pub fn new(inner: Arc<dyn ObjectStore>, rebuild: RebuildFn) -> Self { + Self { + inner, + rebuild, + rebuilt: OnceCell::new(), + } + } + + /// Return the store that should service *this* request: the post-rebuild store when we + /// have one, else the original. + fn current(&self) -> Arc<dyn ObjectStore> { + self.rebuilt + .get() + .cloned() + .unwrap_or_else(|| Arc::clone(&self.inner)) Review Comment: Agreed, and I think this constrains where the fix can live. A path-aware selector inside the cache wouldn't be consulted at read time: `prepare_object_store_with_configs` registers one store per `ObjectStoreUrl`, and that URL carries scheme, config hash, backend and host but no scope. `DefaultObjectStoreRegistry::register_store` is a plain `insert`, so two `NativeScan` operators on the same bucket under different scopes overwrite each other, and `FileScanConfig::open` resolves from that single URL at execute time. That's also why the one-file-per-partition case you mention can't be fixed in the cache -- a file group holds many paths but only one URL. Would it work to make the registered store itself the selector? One wrapper per `(bucket, config_hash, backend)` holding the scoped stores and dispatching on `location` inside `get_opts` and `get_ranges` would put routing in one place, give one retry per operation with no process-lifetime latch, and handle a partition that spans scopes. The alternative is splitting file groups by scope and registering a URL per scope, which looks considerably more invasive. Either way, the comment above `register_object_store` currently documents the invariant that `Vec<ScopeEntry>` breaks, so that needs updating too. -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected] --------------------------------------------------------------------- To unsubscribe, e-mail: [email protected] For additional commands, e-mail: [email protected]
