From c78f251c3fbabbe91b269bd495ebcd330dc5df3b Mon Sep 17 00:00:00 2001 From: gaurav-amz Date: Thu, 11 Jun 2026 13:30:59 +0530 Subject: [PATCH 01/14] [analytics-engine] Wipe spill directory contents instead of the directory itself (#22085) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * [analytics-engine] Wipe spill directory contents instead of the directory itself create_global_runtime previously called fs::remove_dir_all(spill_dir), which rmdirs the spill root as its final step. That requires write permission on the spill directory parent. On managed deployments where spill_dir is a mount point and the JVM user does not own its parent, this fails with EACCES at boot: Failed to clear leaked spill files in /mnt/analytics/backend/spill (io kind=PermissionDenied): Permission denied (os error 13) Wipe the directory children only — top-level files and directories are removed, but the spill root itself stays in place. DataFusion create_local_dirs sees the existing root, skips std::fs::create_dir, and provisions a fresh datafusion-XXXXXX via tempdir_in, all of which only require write on the spill directory itself. Top-level symlinks are unlinked rather than descended into so a stray symlink cannot redirect the wipe outside the spill mount. Tests: - create_global_runtime_succeeds_when_jvm_does_not_own_spill_parent: chmods parent to 0o555 to reproduce the mount-point topology. - create_global_runtime_unlinks_top_level_symlink_without_following: pins the symlink-defense behavior end-to-end. - create_global_runtime_clears_leaked_spill_files_recursively: updated assertion to reflect that the root is now preserved (previously asserted DataFusion recreated it). - SpillCleanupOnBootIT (new): pre-seeds leaked entries before cluster boot; verifies cluster comes up healthy, seeded entries are gone, and the spill root still exists. Signed-off-by: Gaurav Singh * [analytics-engine] Defer spill cleanup recursion to background thread Follow-up to the prior commit. The previous fix wiped spill directory contents synchronously on the boot thread via fs::read_dir + fs::remove_dir_all per child. For typical workloads this completes in under a second, but the boot path stays sub-second only as long as orphan counts and sizes stay small. Split the cleanup into two phases so boot is fast even when prior orphans hold tens of GB. Phase 1 (sync, on boot thread): unlink top-level files and symlinks inline; rename top-level subdirectories to .stale siblings. Renames are metadata-only — sub-millisecond per entry, contents not touched. Failure aborts boot loudly with full error context. Phase 2 (async, background thread): re-scan the spill directory and recursively delete every *.stale entry. Filtering by suffix means the fresh datafusion-XXXXXX/ provisioned by DataFusion tempdir_in is never a deletion target. Re-scanning also picks up *.stale leftovers from a prior boot whose cleanup thread did not finish. Failures are logged only — the cluster has joined; stragglers are reaped on the next boot. Spawn failure (rare — EAGAIN/ENOMEM under extreme system pressure) is also fail-soft: by phase 1 the spill directory is already in a valid state for DataFusion to start, and *.stale entries persist until a future boot reaps them. Aborting boot here would convert a transient thread-library failure into a node-level outage with no benefit. Race-free with DataFusion runtime build because phase 1 completes before RuntimeEnvBuilder::build() runs. By the time tempdir_in creates spill/datafusion-/, all orphans are renamed to *.stale; the cleanup thread filters by .stale suffix and never sees the live dir. Tests added: - create_global_runtime_renames_orphan_subdirs_to_stale_then_async_removes: pins the rename-then-async-delete contract (originals gone inline, *.stale entries cleaned by the background thread within 2s). - create_global_runtime_cleans_prior_boot_stale_entries: pins recovery from an interrupted prior phase 2 (no double-suffix). Tests updated: - create_global_runtime_clears_leaked_spill_files_recursively: waits on phase 2 via a wait_until helper. - create_global_runtime_succeeds_when_jvm_does_not_own_spill_parent: waits on phase 2 for the leaked subdir; loose file is still inline. Signed-off-by: Gaurav Singh * [analytics-engine] Unify spill cleanup phase 1 to a single rename branch Follow-up to the prior commit. Phase 1 previously had two branches: files and symlinks were unlinked inline; directories were renamed to *.stale. The asymmetry was historical — files were cheap to unlink synchronously, so we did. But it bought nothing structurally and introduced two real costs: 1. Stuck top-level files (immutable bit, ENOENT race, etc.) failed boot loudly. Async-renaming those files lets boot succeed and converts the failure into a phase-2 warning, which is the same fail-soft behavior we already use for everything else in phase 2. 2. Phase 1 had to call entry.file_type() (an lstat syscall) on every entry to dispatch between unlink-vs-rename. Switch phase 1 to a single uniform rename branch: every immediate child that does not already end in .stale is renamed to .stale. fs::rename is metadata-only and does not follow symlinks, so a symlink at the leaf is renamed (the link itself, not the target) just like any other entry. Phase 2 now does the dispatch instead — file_type via lstat-semantics to decide between fs::remove_file (files and symlinks) and fs::remove_dir_all (directories). The dispatch latency moves off the boot path; total syscall count is unchanged. Tests updated: every test that previously asserted "X removed inline" now asserts "X renamed to X.stale and removed by phase 2" via the existing wait_until helper. Symlink test gains an explicit phase 2 wait and re-confirms the symlink target outside the spill dir is preserved verbatim across both phases. Signed-off-by: Gaurav Singh * [analytics-engine] Surface stale_entry_count in spill stats The two-phase boot cleanup renames orphan spill entries to *.stale and removes them in a background thread. If that thread fails (or never spawns due to system pressure), or if individual entries cannot be removed (immutable bit, broken perms, EBUSY, etc.), the .stale entries persist on disk. Until now there was no operational signal for this state — only per-node log warnings. Add a stale_entry_count field to SpillStats, computed live by SpillStatsCollector via Files.list(spillDirectory) on each stats request (filters by .stale suffix, counts). The count is rendered in the disk_spill JSON block and surfaced through the existing /_plugins/_analytics_backend_datafusion/stats/disk_spill endpoint. Operators can alarm on stale_entry_count > 0 sustained across boots to detect cleanup failures that would otherwise only show up as slow disk-space leak. Cost: one extra read_dir per stats request (~10-50us). Stats endpoint is queried at most every few seconds, so the cost is negligible. On collection error the count falls back to 0 — disk usage fields remain the primary signal. Tests: - SpillStatsTests: wire-format round-trip carries the new field; toXContent renders stale_entry_count; equals/hashCode include it. - SpillStatsCollectorTests: testCollectCountsStaleEntries seeds a .stale dir + .stale file + .stale symlink alongside non-.stale entries; asserts the count reports exactly the .stale set. - WriteableRoundTripPropertyTests: 6-arg combinator covers the new field in property-based round-trip generation. - DataFusionStatsTests, TransportDataFusionStatsActionTests: ctor-site updates to 6-arg form (no behavioral change). Signed-off-by: Gaurav Singh --------- Signed-off-by: Gaurav Singh --- .../rust/src/api.rs | 457 ++++++++++++++++-- .../be/datafusion/stats/SpillStats.java | 30 +- .../datafusion/stats/SpillStatsCollector.java | 37 +- .../WriteableRoundTripPropertyTests.java | 3 +- .../TransportDataFusionStatsActionTests.java | 2 +- .../stats/DataFusionStatsTests.java | 2 +- .../stats/SpillStatsCollectorTests.java | 31 ++ .../be/datafusion/stats/SpillStatsTests.java | 19 +- sandbox/qa/analytics-engine-rest/build.gradle | 55 +++ .../analytics/qa/SpillCleanupOnBootIT.java | 118 +++++ 10 files changed, 684 insertions(+), 70 deletions(-) create mode 100644 sandbox/qa/analytics-engine-rest/src/test/java/org/opensearch/analytics/qa/SpillCleanupOnBootIT.java diff --git a/sandbox/plugins/analytics-backend-datafusion/rust/src/api.rs b/sandbox/plugins/analytics-backend-datafusion/rust/src/api.rs index 3be2e91b9c609..3c90320a30037 100644 --- a/sandbox/plugins/analytics-backend-datafusion/rust/src/api.rs +++ b/sandbox/plugins/analytics-backend-datafusion/rust/src/api.rs @@ -348,32 +348,52 @@ pub struct ShardView { /// Returns a heap-allocated pointer (as i64) to `DataFusionRuntime`. /// Caller must call `close_global_runtime` exactly once to free it. /// -/// # Side effect: spill directory is wiped +/// # Side effect: spill directory contents are wiped (two-phase) /// -/// When `spill_dir` is non-empty, the directory is cleared (`remove_dir_all`) -/// before the `DiskManager` is built. DataFusion does not sweep stale entries -/// itself — `create_local_dirs` in `datafusion-execution` only creates a fresh -/// `datafusion-XXXXXX/` subdirectory and never touches siblings — so a prior -/// non-graceful shutdown (kill -9, OOM-kill, container restart) leaves orphaned -/// `datafusion-*/` trees that accumulate forever. The directory is -/// OpenSearch-owned by contract, so wiping is safe. +/// When `spill_dir` is non-empty, every immediate child is removed: +/// * **Phase 1 (sync, on this thread):** every immediate child (files, +/// symlinks, and subdirectories alike) is renamed to a `.stale` +/// sibling. Renames are metadata-only — one syscall per entry, contents +/// are not touched. +/// * **Phase 2 (async, background thread):** every `*.stale` entry is +/// removed off the boot path. Files and symlinks via `remove_file` +/// (which unlinks the symlink itself without following it); directories +/// via `remove_dir_all`. /// -/// We do not recreate the directory here. DataFusion's `create_local_dirs` will -/// recreate the root via `std::fs::create_dir` and provision a fresh `TempDir` -/// inside it during `DiskManagerBuilder::build()`. +/// The split keeps boot fast even with tens of GB of orphan spill data: the +/// boot thread pays only the rename count (≈10s of µs per entry) and the +/// recursive unlinks happen in parallel with cluster join. Re-scanning by +/// suffix in phase 2 also cleans up `*.stale` leftovers from any prior boot +/// whose cleanup thread did not finish. /// -/// Any cleanup failure (permission denied, read-only filesystem, I/O error, -/// etc.) is operator-actionable and aborts boot with `DataFusionError::Configuration` -/// — the spill directory being in an unrecoverable state implies queries that -/// need to spill would also fail later, so we surface the problem at boot time -/// with full error context instead of letting orphans accumulate silently and -/// failing later mid-query. +/// This sweeps `datafusion-*/` orphans left by a non-graceful shutdown +/// (kill -9, OOM-kill, container restart) which would otherwise accumulate +/// forever — `create_local_dirs` in `datafusion-execution` only creates a +/// fresh `datafusion-XXXXXX/` and never touches siblings. The fresh dir's +/// name has no `.stale` suffix, so phase 2 cannot collide with it. /// -/// Safe today because (a) `datafusion.spill_directory` is `NodeScope + Final`, so -/// `DataFusionPlugin.createComponents` calls this exactly once per JVM in -/// production, and (b) Rust unit tests pass a fresh `tempdir()` per call. Anyone -/// adding a new caller (hot-reload, multiple runtimes sharing a directory, etc.) -/// must rethink this — calling mid-flight will nuke active spill state. +/// The spill directory itself is preserved. It may be a pre-existing mount +/// point whose parent the JVM user does not own; rmdir'ing the root would +/// fail with `EACCES` in that topology. Top-level symlinks are renamed in +/// phase 1 (which renames the link, never the target) and unlinked in +/// phase 2 via `remove_file`, so a stray symlink can never redirect the +/// wipe outside the spill mount. +/// +/// Phase 1 failure aborts boot with `DataFusionError::Configuration` — if +/// we cannot rename orphans now, queries that need to spill would fail +/// later mid-query anyway. Phase 2 failures (including thread-spawn) are +/// logged only; the cluster has either joined or is about to, and +/// stragglers are reaped on the next boot. +/// +/// # Operator contract +/// +/// The OpenSearch process must own `spill_dir` with `rwx`. Write on the +/// parent is **not** required — `spill_dir` is treated as a pre-existing +/// mount point. Callers wanting auto-creation must grant write on the parent. +/// +/// Safe today because `datafusion.spill_directory` is `NodeScope + Final`, so +/// this runs exactly once per JVM. A future caller invoking it mid-flight +/// would nuke active spill state — rethink before adding one. pub fn create_global_runtime( memory_pool_limit: i64, cache_manager_ptr: i64, @@ -410,24 +430,143 @@ pub fn create_global_runtime( spill_limit as u64 }; - // Clear leaked spill files from a prior non-graceful shutdown. Any failure - // here (permission denied, read-only filesystem, I/O error, etc.) is - // operator-actionable and signals that the spill directory is in a state - // where DataFusion cannot reliably write either — fail boot loudly with - // full error context so the operator fixes it before the node joins the - // cluster, instead of accumulating orphan files and inflating - // disk_used_bytes silently. + // Wipe leaked entries from a prior non-graceful shutdown. + // + // Two-phase to keep boot fast even when prior orphans hold tens of GB: + // Phase 1 (sync, on boot thread): rename every immediate child to a + // .stale sibling. One uniform branch — files, symlinks, and + // directories are all renamed identically. Renames are metadata-only + // (one syscall per entry, contents not touched). Failure aborts + // boot loudly with full error context. + // Phase 2 (async, background thread): re-scan and remove every + // *.stale entry — `remove_file` for files and symlinks (so we + // never follow a symlink through to its target), `remove_dir_all` + // for directories. Filtering by the .stale suffix means the fresh + // datafusion-XXXXXX/ that DataFusion provisions next is never a + // deletion target. + // + // The spill directory itself is never removed (mount-point safe). + // Phase 1 only requires write on spill_dir, never on its parent. let spill_path = PathBuf::from(spill_dir); if spill_path.exists() { - if let Err(e) = fs::remove_dir_all(&spill_path) { + let entries = fs::read_dir(&spill_path).map_err(|e| { let msg = format!( - "Failed to clear leaked spill files in {} (io kind={:?}): {}. \ - Verify the OpenSearch process owns the spill directory and the volume \ - is mounted read-write before restarting.", + "Failed to enumerate spill directory {} (io kind={:?}): {}. \ + Verify the process owns the spill directory before restarting.", spill_dir, e.kind(), e ); log::error!("{}", msg); - return Err(DataFusionError::Configuration(msg)); + DataFusionError::Configuration(msg) + })?; + for entry in entries { + let entry = entry.map_err(|e| { + let msg = format!( + "Failed to read spill directory entry in {} (io kind={:?}): {}", + spill_dir, e.kind(), e + ); + log::error!("{}", msg); + DataFusionError::Configuration(msg) + })?; + let path = entry.path(); + let name = entry.file_name(); + let name_str = name.to_string_lossy(); + + // Skip already-renamed leftovers from a prior boot whose async + // cleanup didn't finish — phase 2 below picks them up. + if name_str.ends_with(".stale") { + continue; + } + + // Rename to .stale within the same parent. Metadata-only; + // contents not touched. Same parent (spill_path), so requires + // only write on spill_path itself. Works uniformly for files, + // symlinks, and directories — fs::rename does not follow + // symlinks (it renames the link itself). + let stale = spill_path.join(format!("{}.stale", name_str)); + if let Err(e) = fs::rename(&path, &stale) { + let msg = format!( + "Failed to rename leaked spill entry {} -> {} (io kind={:?}): {}. \ + Verify the process owns the spill directory before restarting.", + path.display(), stale.display(), e.kind(), e + ); + log::error!("{}", msg); + return Err(DataFusionError::Configuration(msg)); + } + } + + // Phase 2: spawn the recursive deletion of all *.stale entries. + // Re-scan inside the thread so we pick up entries from this boot + // AND any *.stale leftovers from prior boots whose cleanup didn't + // finish. Errors inside the thread are logged only — the cluster + // has not yet joined, but stragglers are recoverable next boot. + // + // Spawn failure itself is rare (EAGAIN/ENOMEM under extreme system + // pressure) and not boot-fatal: the disk is already in a valid + // state for DataFusion to start; the *.stale entries just persist + // until a future boot with a healthy thread library reaps them. + // Failing boot here would be more disruptive than the wasted disk. + let spill_path_for_gc = spill_path.clone(); + if let Err(e) = std::thread::Builder::new() + .name("datafusion-spill-gc".to_string()) + .spawn(move || { + let entries = match fs::read_dir(&spill_path_for_gc) { + Ok(e) => e, + Err(e) => { + log::warn!( + "Background spill cleanup: failed to enumerate {} (io kind={:?}): {}", + spill_path_for_gc.display(), e.kind(), e + ); + return; + } + }; + for entry in entries { + let entry = match entry { + Ok(e) => e, + Err(e) => { + log::warn!("Background spill cleanup: read_dir error: {}", e); + continue; + } + }; + let name = entry.file_name(); + if !name.to_string_lossy().ends_with(".stale") { + continue; + } + let path = entry.path(); + // Use file_type (lstat semantics, does not follow symlinks) + // to dispatch: directories use remove_dir_all, everything + // else (regular files, symlinks) uses remove_file. This + // keeps the symlink defense — fs::remove_file on a symlink + // unlinks the link itself without following it. + let result = match entry.file_type() { + Ok(ft) => { + if ft.is_dir() && !ft.is_symlink() { + fs::remove_dir_all(&path) + } else { + fs::remove_file(&path) + } + } + Err(e) => { + log::warn!( + "Background spill cleanup: failed to stat {} (io kind={:?}): {}", + path.display(), e.kind(), e + ); + continue; + } + }; + if let Err(e) = result { + log::warn!( + "Background spill cleanup: failed to remove {} (io kind={:?}): {}", + path.display(), e.kind(), e + ); + } + } + }) + { + log::warn!( + "Failed to spawn background spill cleanup thread: {}. \ + Renamed *.stale entries will accumulate until the next successful boot.", + e + ); } } @@ -1726,6 +1865,19 @@ mod tests { /// tests would race on these globals and produce flaky assertions. static SPILL_GLOBALS_LOCK: std::sync::Mutex<()> = std::sync::Mutex::new(()); + /// Test helper: poll until `predicate` returns true or `timeout_ms` elapses. + /// Used to wait on the background spill-cleanup thread without an arbitrary sleep. + fn wait_until bool>(timeout_ms: u64, predicate: F) -> bool { + let deadline = std::time::Instant::now() + std::time::Duration::from_millis(timeout_ms); + while std::time::Instant::now() < deadline { + if predicate() { + return true; + } + std::thread::sleep(std::time::Duration::from_millis(20)); + } + predicate() + } + #[test] fn create_global_runtime_with_empty_spill_dir_disables_disk_manager() { // Empty spill_dir is the "disabled" sentinel from Java. The runtime must build @@ -1773,11 +1925,13 @@ mod tests { let ptr = create_global_runtime(64 * 1024 * 1024, 0, spill_path, 0).expect("runtime build"); assert!(ptr > 0); - // The startup cleanup must have removed the leaked file. - assert!( - !sentinel.exists(), - "leaked spill file must be removed by create_global_runtime startup cleanup" - ); + // Phase 1 renames the sentinel file to leaked_from_prior_run.tmp.stale + // synchronously; phase 2 unlinks it asynchronously. The original name + // is gone immediately; wait for the .stale name to disappear too. + assert!(!sentinel.exists(), "sentinel original name must be gone (renamed)"); + let stale = tmp.path().join("leaked_from_prior_run.tmp.stale"); + let cleaned = wait_until(2000, || !stale.exists()); + assert!(cleaned, "background cleanup must remove the .stale sentinel within 2s"); let runtime = unsafe { &*(ptr as *const DataFusionRuntime) }; assert!( @@ -1797,7 +1951,13 @@ mod tests { // Operator-confirmed contract: the spill directory is OpenSearch-owned and any // contents present at startup are leaked from a prior non-graceful shutdown. // create_global_runtime must clear the directory recursively (files AND - // subdirectories) before constructing the DiskManager. + // subdirectories) before / during constructing the DiskManager. + // + // Cleanup is two-phase: phase 1 renames every immediate child to *.stale on + // the boot thread (uniform across files, symlinks, dirs); phase 2 removes + // them in a background thread (remove_file for files/symlinks, remove_dir_all + // for dirs). The original names are gone immediately after phase 1; wait + // briefly for phase 2 to clear the *.stale entries. let _guard = SPILL_GLOBALS_LOCK.lock().unwrap_or_else(|e| e.into_inner()); let tmp = tempfile::tempdir().expect("tempdir"); let spill_path = tmp.path().to_str().expect("utf-8 path"); @@ -1810,21 +1970,28 @@ mod tests { fs::create_dir_all(&nested_dir).expect("seed nested subdirs"); let nested_file = nested_dir.join("deep.tmp"); fs::write(&nested_file, b"nested leak").expect("seed nested file"); + let outer_subdir = tmp.path().join("subdir"); assert!(top_file.exists()); assert!(nested_file.exists()); let ptr = create_global_runtime(64 * 1024 * 1024, 0, spill_path, 0).expect("runtime build"); assert!(ptr > 0); - // Both seeded entries must be gone. The directory itself ends up existing again - // because DataFusion's create_local_dirs recreates the root (and provisions a - // fresh datafusion-XXXXXX/ inside it) during DiskManagerBuilder::build() — our - // cleanup only wipes; DataFusion handles recreation. - assert!(!top_file.exists(), "top-level leaked file must be removed"); + // Phase 1: original names gone (renamed to *.stale). + assert!(!top_file.exists(), "top-level file original name must be gone (renamed)"); + assert!(!outer_subdir.exists(), "original subdir name must be gone (renamed)"); + + // Phase 2: *.stale entries cleaned by the background thread. + let stale_top = tmp.path().join("top.tmp.stale"); + let stale_dir = tmp.path().join("subdir.stale"); + let cleaned = wait_until(2000, || !stale_top.exists() && !stale_dir.exists()); + assert!(cleaned, "background cleanup must remove both .stale entries within 2s"); assert!(!nested_file.exists(), "nested leaked file must be removed"); assert!(!nested_dir.exists(), "nested leaked subdir must be removed"); - assert!(tmp.path().exists(), "DataFusion's create_local_dirs must have recreated the root"); - assert!(tmp.path().is_dir(), "spill root must be a directory after DataFusion provisioning"); + + // Spill root itself must remain (cleanup wipes children only). + assert!(tmp.path().exists(), "spill root must be preserved across cleanup"); + assert!(tmp.path().is_dir(), "spill root must remain a directory after cleanup"); unsafe { close_global_runtime(ptr) }; } @@ -1846,8 +2013,8 @@ mod tests { // Cleanup failure is operator-actionable (permissions, RO mount, I/O) and // implies the runtime would fail later mid-query anyway. Surface the failure // at boot with full context. Trigger the failure path by pointing spill_dir - // at a regular file: spill_path.exists() returns true, but remove_dir_all - // refuses to operate on a non-directory. + // at a regular file: spill_path.exists() returns true, but read_dir refuses + // to enumerate a non-directory and returns ErrorKind::NotADirectory. let _guard = SPILL_GLOBALS_LOCK.lock().unwrap_or_else(|e| e.into_inner()); let tmp = tempfile::tempdir().expect("tempdir"); let bad_path = tmp.path().join("regular_file"); @@ -1863,12 +2030,202 @@ mod tests { assert!(msg.contains(bad_path_str), "error must reference the offending path; got: {}", msg); assert!(msg.contains("io kind="), "error must include the io::ErrorKind for diagnosis; got: {}", msg); assert!( - msg.contains("Failed to clear leaked spill files"), + msg.contains("Failed to enumerate spill directory") + || msg.contains("Failed to clear leaked spill entry"), "error must identify it as a cleanup failure; got: {}", msg ); } + /// Spill directory is owned by the JVM but its parent isn't (mount-point + /// topology). `fs::remove_dir_all(spill_dir)` would fail at the final + /// rmdir; the fix wipes contents only. + /// + /// Skipped under EUID 0 — root has CAP_DAC_OVERRIDE (or the macOS equivalent) + /// and bypasses the chmod 0o555 we use to simulate the locked parent, so the + /// pre-fix code path would silently succeed and the test would assert nothing. + /// CI runs as a non-root user; if you need to verify under root, run the + /// test in a user namespace or container that drops the capability. + #[test] + #[cfg(unix)] + fn create_global_runtime_succeeds_when_jvm_does_not_own_spill_parent() { + use std::os::unix::fs::PermissionsExt; + + // SAFETY: geteuid() is async-signal-safe and has no preconditions. + let euid = unsafe { libc::geteuid() }; + if euid == 0 { + eprintln!( + "skipping create_global_runtime_succeeds_when_jvm_does_not_own_spill_parent: \ + EUID 0 bypasses chmod 0o555 via DAC override; the pre-fix code path \ + would silently succeed and the test would assert nothing" + ); + return; + } + + let _guard = SPILL_GLOBALS_LOCK.lock().unwrap_or_else(|e| e.into_inner()); + let parent = tempfile::tempdir().expect("parent tempdir"); + let spill_path = parent.path().join("spill"); + fs::create_dir(&spill_path).expect("create spill mount-point dir"); + + // Seed leaked entries: a datafusion-* subtree (kill -9 leftover) and a loose file. + let leaked_dir = spill_path.join("datafusion-aB3kF7"); + fs::create_dir(&leaked_dir).expect("create leaked subdir"); + let leaked_file = leaked_dir.join("tmp_spill_001.arrow"); + fs::write(&leaked_file, b"stale spill data").expect("seed leaked file"); + let loose_file = spill_path.join("stray.txt"); + fs::write(&loose_file, b"loose top-level file").expect("seed loose file"); + + // Lock parent to r+x only — simulates the mount-point parent the JVM doesn't own. + let original_parent_mode = fs::metadata(parent.path()).expect("stat parent").permissions().mode(); + let mut locked = fs::metadata(parent.path()).expect("stat parent").permissions(); + locked.set_mode(0o555); + fs::set_permissions(parent.path(), locked).expect("chmod parent 555"); + + let spill_str = spill_path.to_str().expect("utf-8 path"); + let result = create_global_runtime(64 * 1024 * 1024, 0, spill_str, 0); + + // Restore parent perms via RAII so tempdir cleanup runs even on assertion failure. + struct RestorePerms<'a> { + path: &'a std::path::Path, + mode: u32, + } + impl Drop for RestorePerms<'_> { + fn drop(&mut self) { + if let Ok(metadata) = fs::metadata(self.path) { + let mut perms = metadata.permissions(); + perms.set_mode(self.mode); + let _ = fs::set_permissions(self.path, perms); + } + } + } + let _restore = RestorePerms { path: parent.path(), mode: original_parent_mode }; + + let ptr = result.expect("runtime build must succeed when only the parent is read-only"); + assert!(ptr > 0); + + // Phase 1: both originals are renamed inline — gone immediately by the + // original name. Phase 2 unlinks the .stale entries asynchronously. + assert!(!loose_file.exists(), "loose top-level file original name must be gone (renamed)"); + assert!(!leaked_dir.exists(), "leaked datafusion-* original name must be gone (renamed)"); + + let stale_loose = spill_path.join("stray.txt.stale"); + let stale_dir = spill_path.join("datafusion-aB3kF7.stale"); + let cleaned = wait_until(2000, || !stale_loose.exists() && !stale_dir.exists()); + assert!(cleaned, "background cleanup must remove both .stale entries within 2s"); + assert!(!leaked_file.exists(), "leaked file under leaked subdir must be removed"); + + // Spill root itself preserved. + assert!(spill_path.exists(), "spill mount-point dir must be preserved"); + assert!(spill_path.is_dir(), "spill mount-point must remain a directory"); + + unsafe { close_global_runtime(ptr) }; + } + + /// Top-level symlinks must be unlinked, not followed — otherwise the wipe + /// could escape the spill directory and delete files elsewhere. + #[test] + #[cfg(unix)] + fn create_global_runtime_unlinks_top_level_symlink_without_following() { + let _guard = SPILL_GLOBALS_LOCK.lock().unwrap_or_else(|e| e.into_inner()); + let tmp = tempfile::tempdir().expect("tempdir"); + let spill_path = tmp.path().join("spill"); + fs::create_dir(&spill_path).expect("create spill dir"); + + // outside_target sits next to spill/. If cleanup followed the symlink, it would be deleted. + let outside = tmp.path().join("outside_target.txt"); + fs::write(&outside, b"must NOT be touched").expect("seed outside file"); + + let link = spill_path.join("escape_link"); + std::os::unix::fs::symlink(&outside, &link).expect("create symlink"); + assert!(link.is_symlink(), "precondition: link is a symlink"); + + let spill_str = spill_path.to_str().expect("utf-8 path"); + let ptr = create_global_runtime(64 * 1024 * 1024, 0, spill_str, 0).expect("runtime build"); + assert!(ptr > 0); + + // Phase 1: symlink is renamed inline (fs::rename does not follow symlinks). + assert!(!link.exists(), "symlink original name must be gone (renamed)"); + let stale_link = spill_path.join("escape_link.stale"); + + // Phase 2: remove_file on a symlink unlinks the link itself, never the target. + let cleaned = wait_until(2000, || !stale_link.exists()); + assert!(cleaned, "background cleanup must remove escape_link.stale within 2s"); + + // Critical: the target outside spill must be intact across both phases. + // If phase 1 had followed the symlink during rename, or phase 2 followed + // it during remove, the outside file would have been clobbered. + assert!( + outside.exists(), + "symlink target outside spill dir must NOT be touched (cleanup must not follow symlinks)" + ); + assert_eq!( + fs::read(&outside).expect("read outside file"), + b"must NOT be touched", + "symlink target contents must be preserved verbatim" + ); + + unsafe { close_global_runtime(ptr) }; + } + + /// Phase 1 (sync rename) must observably move each orphan subdir to a + /// *.stale name. Phase 2 then removes the *.stale entries asynchronously. + /// This pins the rename behavior so a future refactor that goes back to + /// inline recursive removal would be flagged. + #[test] + fn create_global_runtime_renames_orphan_subdirs_to_stale_then_async_removes() { + let _guard = SPILL_GLOBALS_LOCK.lock().unwrap_or_else(|e| e.into_inner()); + let tmp = tempfile::tempdir().expect("tempdir"); + let spill_path = tmp.path().to_str().expect("utf-8 path"); + + let leaked_a = tmp.path().join("datafusion-aB3kF7"); + fs::create_dir(&leaked_a).expect("create leaked a"); + fs::write(leaked_a.join("tmp_001.arrow"), b"data a").expect("write a"); + let leaked_b = tmp.path().join("datafusion-Xy9pQ2"); + fs::create_dir(&leaked_b).expect("create leaked b"); + fs::write(leaked_b.join("tmp_002.arrow"), b"data b").expect("write b"); + + let ptr = create_global_runtime(64 * 1024 * 1024, 0, spill_path, 0).expect("runtime build"); + assert!(ptr > 0); + + // Originals were renamed inline — gone immediately by the original name. + assert!(!leaked_a.exists(), "original orphan a must be renamed away"); + assert!(!leaked_b.exists(), "original orphan b must be renamed away"); + + // Phase 2 will eventually remove both *.stale entries. + let stale_a = tmp.path().join("datafusion-aB3kF7.stale"); + let stale_b = tmp.path().join("datafusion-Xy9pQ2.stale"); + let cleaned = wait_until(2000, || !stale_a.exists() && !stale_b.exists()); + assert!(cleaned, "background cleanup must remove both *.stale entries within 2s"); + + unsafe { close_global_runtime(ptr) }; + } + + /// Pre-existing *.stale entries from a prior boot whose phase 2 didn't finish + /// must be cleaned up by the next boot's phase 2, and phase 1 must not + /// double-suffix them (no datafusion-old.stale.stale). + #[test] + fn create_global_runtime_cleans_prior_boot_stale_entries() { + let _guard = SPILL_GLOBALS_LOCK.lock().unwrap_or_else(|e| e.into_inner()); + let tmp = tempfile::tempdir().expect("tempdir"); + let spill_path = tmp.path().to_str().expect("utf-8 path"); + + let leftover = tmp.path().join("datafusion-old.stale"); + fs::create_dir(&leftover).expect("create leftover"); + fs::write(leftover.join("residue.arrow"), b"prior boot data").expect("write residue"); + + let ptr = create_global_runtime(64 * 1024 * 1024, 0, spill_path, 0).expect("runtime build"); + assert!(ptr > 0); + + let cleaned = wait_until(2000, || !leftover.exists()); + assert!(cleaned, "prior-boot .stale leftover must be cleaned within 2s"); + assert!( + !tmp.path().join("datafusion-old.stale.stale").exists(), + "phase 1 must NOT double-suffix existing .stale entries" + ); + + unsafe { close_global_runtime(ptr) }; + } + #[test] fn stringview_gc_compacts_sliced_buffers() { let total_rows = 100_000usize; diff --git a/sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/stats/SpillStats.java b/sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/stats/SpillStats.java index fbc6fc4bbdc29..b0425af1e9a6e 100644 --- a/sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/stats/SpillStats.java +++ b/sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/stats/SpillStats.java @@ -36,13 +36,29 @@ public class SpillStats implements Writeable, ToXContentFragment { private final long diskAvailableBytes; private final long diskUsedBytes; private final long diskReservedBytes; - - public SpillStats(String directory, long diskTotalBytes, long diskAvailableBytes, long diskUsedBytes, long diskReservedBytes) { + /** + * Number of {@code *.stale} entries currently in the spill directory. + * Each entry is an orphan from a prior boot whose async cleanup has not yet + * removed it. A persistently non-zero value across boots indicates an + * operator-actionable failure (immutable bit, broken perms, etc.) preventing + * the background cleanup thread from completing. + */ + private final long staleEntryCount; + + public SpillStats( + String directory, + long diskTotalBytes, + long diskAvailableBytes, + long diskUsedBytes, + long diskReservedBytes, + long staleEntryCount + ) { this.directory = directory == null ? "" : directory; this.diskTotalBytes = diskTotalBytes; this.diskAvailableBytes = diskAvailableBytes; this.diskUsedBytes = diskUsedBytes; this.diskReservedBytes = diskReservedBytes; + this.staleEntryCount = staleEntryCount; } public SpillStats(StreamInput in) throws IOException { @@ -51,6 +67,7 @@ public SpillStats(StreamInput in) throws IOException { this.diskAvailableBytes = in.readVLong(); this.diskUsedBytes = in.readVLong(); this.diskReservedBytes = in.readVLong(); + this.staleEntryCount = in.readVLong(); } @Override @@ -60,6 +77,7 @@ public void writeTo(StreamOutput out) throws IOException { out.writeVLong(diskAvailableBytes); out.writeVLong(diskUsedBytes); out.writeVLong(diskReservedBytes); + out.writeVLong(staleEntryCount); } @Override @@ -70,6 +88,7 @@ public XContentBuilder toXContent(XContentBuilder builder, Params params) throws builder.field("disk_available_bytes", diskAvailableBytes); builder.field("disk_used_bytes", diskUsedBytes); builder.field("disk_reserved_bytes", diskReservedBytes); + builder.field("stale_entry_count", staleEntryCount); builder.endObject(); return builder; } @@ -94,6 +113,10 @@ public long getDiskReservedBytes() { return diskReservedBytes; } + public long getStaleEntryCount() { + return staleEntryCount; + } + @Override public boolean equals(Object o) { if (this == o) return true; @@ -103,11 +126,12 @@ public boolean equals(Object o) { && diskAvailableBytes == that.diskAvailableBytes && diskUsedBytes == that.diskUsedBytes && diskReservedBytes == that.diskReservedBytes + && staleEntryCount == that.staleEntryCount && Objects.equals(directory, that.directory); } @Override public int hashCode() { - return Objects.hash(directory, diskTotalBytes, diskAvailableBytes, diskUsedBytes, diskReservedBytes); + return Objects.hash(directory, diskTotalBytes, diskAvailableBytes, diskUsedBytes, diskReservedBytes, staleEntryCount); } } diff --git a/sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/stats/SpillStatsCollector.java b/sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/stats/SpillStatsCollector.java index 8262a2a9729fe..e9c44a37ee70f 100644 --- a/sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/stats/SpillStatsCollector.java +++ b/sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/stats/SpillStatsCollector.java @@ -19,6 +19,7 @@ import java.security.AccessController; import java.security.PrivilegedActionException; import java.security.PrivilegedExceptionAction; +import java.util.stream.Stream; /** * Builds a {@link SpillStats} snapshot from the configured spill directory and @@ -45,12 +46,12 @@ private SpillStatsCollector() {} * empty string means spill is disabled * @param spillMemoryLimit the resolved value of {@code datafusion.spill_memory_limit_bytes} * in bytes (the user-facing setting value) - * @return a populated {@link SpillStats}; on error, all byte fields are 0 but the - * {@code directory} and {@code disk_reserved_bytes} fields are preserved. + * @return a populated {@link SpillStats}; on error, byte fields and {@code stale_entry_count} + * are 0 but the {@code directory} and {@code disk_reserved_bytes} fields are preserved. */ public static SpillStats collect(String spillDirectory, long spillMemoryLimit) { if (spillDirectory == null || spillDirectory.isEmpty()) { - return new SpillStats("", 0L, 0L, 0L, 0L); + return new SpillStats("", 0L, 0L, 0L, 0L, 0L); } // SecurityManager: spill_directory is operator-configured and not on any path the @@ -64,7 +65,7 @@ public static SpillStats collect(String spillDirectory, long spillMemoryLimit) { // here would happily return volume stats for a file that was misconfigured into // place at the spill path. if (!Files.isDirectory(path)) { - return new SpillStats(spillDirectory, 0L, 0L, 0L, spillMemoryLimit); + return new SpillStats(spillDirectory, 0L, 0L, 0L, spillMemoryLimit, 0L); } FileStore fs = Environment.getFileStore(path); // Clamp negatives from FileStore — JDK-8162520: getTotalSpace/getUsableSpace can @@ -73,15 +74,37 @@ public static SpillStats collect(String spillDirectory, long spillMemoryLimit) { long total = clampNonNegative(fs.getTotalSpace()); long available = clampNonNegative(fs.getUsableSpace()); long used = Math.max(0L, total - available); - return new SpillStats(spillDirectory, total, available, used, spillMemoryLimit); + long staleCount = countStaleEntries(path); + return new SpillStats(spillDirectory, total, available, used, spillMemoryLimit, staleCount); }); } catch (PrivilegedActionException pae) { Throwable cause = pae.getCause() != null ? pae.getCause() : pae; logger.warn("Failed to read filesystem stats for spill directory [{}]: {}", spillDirectory, cause.getMessage()); - return new SpillStats(spillDirectory, 0L, 0L, 0L, spillMemoryLimit); + return new SpillStats(spillDirectory, 0L, 0L, 0L, spillMemoryLimit, 0L); } catch (RuntimeException e) { logger.warn("Unexpected error reading spill stats for [{}]: {}", spillDirectory, e.getMessage()); - return new SpillStats(spillDirectory, 0L, 0L, 0L, spillMemoryLimit); + return new SpillStats(spillDirectory, 0L, 0L, 0L, spillMemoryLimit, 0L); + } + } + + /** + * Count immediate children of {@code spillDirectory} whose names end in {@code .stale}. + * These are entries renamed by the boot-time cleanup that the background thread has + * not yet removed (or could not remove). A persistently non-zero value across boots + * indicates an operator-actionable issue. + * + *

Returns 0 on any I/O error rather than failing the whole stats request — the + * filesystem-level fields are always more important than this diagnostic counter. + */ + private static long countStaleEntries(Path spillDirectory) { + try (Stream children = Files.list(spillDirectory)) { + return children.filter(p -> { + Path fileName = p.getFileName(); + return fileName != null && fileName.toString().endsWith(".stale"); + }).count(); + } catch (Exception e) { + logger.warn("Failed to count stale spill entries in [{}]: {}", spillDirectory, e.getMessage()); + return 0L; } } diff --git a/sandbox/plugins/analytics-backend-datafusion/src/propertyTest/java/org/opensearch/be/datafusion/action/stats/WriteableRoundTripPropertyTests.java b/sandbox/plugins/analytics-backend-datafusion/src/propertyTest/java/org/opensearch/be/datafusion/action/stats/WriteableRoundTripPropertyTests.java index 6d48027aeb060..e83f76e5f6951 100644 --- a/sandbox/plugins/analytics-backend-datafusion/src/propertyTest/java/org/opensearch/be/datafusion/action/stats/WriteableRoundTripPropertyTests.java +++ b/sandbox/plugins/analytics-backend-datafusion/src/propertyTest/java/org/opensearch/be/datafusion/action/stats/WriteableRoundTripPropertyTests.java @@ -105,7 +105,8 @@ Arbitrary partitionGateStats() { Arbitrary spillStats() { Arbitrary directory = Arbitraries.strings().alpha().ofMinLength(3).ofMaxLength(20).map(s -> "/spill/" + s); Arbitrary nonNeg = Arbitraries.longs().between(0L, Long.MAX_VALUE / 2); - return Combinators.combine(directory, nonNeg, nonNeg, nonNeg, nonNeg).as(SpillStats::new); + Arbitrary staleCount = Arbitraries.longs().between(0L, 10_000L); + return Combinators.combine(directory, nonNeg, nonNeg, nonNeg, nonNeg, staleCount).as(SpillStats::new); } @Provide diff --git a/sandbox/plugins/analytics-backend-datafusion/src/test/java/org/opensearch/be/datafusion/action/stats/TransportDataFusionStatsActionTests.java b/sandbox/plugins/analytics-backend-datafusion/src/test/java/org/opensearch/be/datafusion/action/stats/TransportDataFusionStatsActionTests.java index d0ba89616aa62..b649f512a1e1e 100644 --- a/sandbox/plugins/analytics-backend-datafusion/src/test/java/org/opensearch/be/datafusion/action/stats/TransportDataFusionStatsActionTests.java +++ b/sandbox/plugins/analytics-backend-datafusion/src/test/java/org/opensearch/be/datafusion/action/stats/TransportDataFusionStatsActionTests.java @@ -216,7 +216,7 @@ public void testNodeOperationHandlesServiceNotStarted() { // ---- Helper: build DataFusionStats with only a populated SpillStats section ---- private static DataFusionStats statsWithSpill() { - return new DataFusionStats(null, null, null, new SpillStats("/mnt/spill", 100L, 60L, 40L, 80L)); + return new DataFusionStats(null, null, null, new SpillStats("/mnt/spill", 100L, 60L, 40L, 80L, 0L)); } // ---- Test: filter with "disk_spill" includes SpillStats ---- diff --git a/sandbox/plugins/analytics-backend-datafusion/src/test/java/org/opensearch/be/datafusion/stats/DataFusionStatsTests.java b/sandbox/plugins/analytics-backend-datafusion/src/test/java/org/opensearch/be/datafusion/stats/DataFusionStatsTests.java index 5622b164229fb..aa93548dd1028 100644 --- a/sandbox/plugins/analytics-backend-datafusion/src/test/java/org/opensearch/be/datafusion/stats/DataFusionStatsTests.java +++ b/sandbox/plugins/analytics-backend-datafusion/src/test/java/org/opensearch/be/datafusion/stats/DataFusionStatsTests.java @@ -38,7 +38,7 @@ private static DataFusionStats sequentialStats() { new NativeExecutorsStats(io, cpu, taskMonitors), new PartitionGateStats("datanode_gate", 12, 0, 0, 0, 0, 12), new PartitionGateStats("coordinator_gate", 12, 0, 0, 0, 0, 12), - new SpillStats("/mnt/spill", 100L, 60L, 40L, 80L) + new SpillStats("/mnt/spill", 100L, 60L, 40L, 80L, 0L) ); } diff --git a/sandbox/plugins/analytics-backend-datafusion/src/test/java/org/opensearch/be/datafusion/stats/SpillStatsCollectorTests.java b/sandbox/plugins/analytics-backend-datafusion/src/test/java/org/opensearch/be/datafusion/stats/SpillStatsCollectorTests.java index 70cf6fb75256a..532d0f668cdc8 100644 --- a/sandbox/plugins/analytics-backend-datafusion/src/test/java/org/opensearch/be/datafusion/stats/SpillStatsCollectorTests.java +++ b/sandbox/plugins/analytics-backend-datafusion/src/test/java/org/opensearch/be/datafusion/stats/SpillStatsCollectorTests.java @@ -24,6 +24,7 @@ public void testCollectReturnsZeroForEmptyDirectory() { assertEquals(0L, stats.getDiskAvailableBytes()); assertEquals(0L, stats.getDiskUsedBytes()); assertEquals(0L, stats.getDiskReservedBytes()); + assertEquals(0L, stats.getStaleEntryCount()); } public void testCollectReadsFilesystemForExistingDirectory() throws IOException { @@ -34,6 +35,7 @@ public void testCollectReadsFilesystemForExistingDirectory() throws IOException assertTrue(stats.getDiskAvailableBytes() > 0L); assertEquals(1_000_000L, stats.getDiskReservedBytes()); assertEquals(stats.getDiskTotalBytes(), stats.getDiskAvailableBytes() + stats.getDiskUsedBytes()); + assertEquals("empty spill dir has no .stale entries", 0L, stats.getStaleEntryCount()); } public void testCollectReturnsZeroBytesForMissingDirectoryButPreservesReserved() { @@ -43,5 +45,34 @@ public void testCollectReturnsZeroBytesForMissingDirectoryButPreservesReserved() assertEquals(missing.toString(), stats.getDirectory()); assertEquals(0L, stats.getDiskTotalBytes()); assertEquals(4242L, stats.getDiskReservedBytes()); + assertEquals(0L, stats.getStaleEntryCount()); + } + + /** + * Pre-seeds three {@code *.stale} entries (a dir, a regular file, and a symlink) + * plus two non-stale entries that must NOT be counted (a live datafusion-* dir + * and a regular file without the suffix). Verifies the count reports exactly 3. + */ + public void testCollectCountsStaleEntries() throws IOException { + Path tmp = createTempDir(); + // Three .stale entries — should be counted. + Files.createDirectory(tmp.resolve("datafusion-aB3kF7.stale")); + Files.writeString(tmp.resolve("loose.txt.stale"), "stale loose"); + // Symlinks may not be supported on every test platform; try and skip if unavailable. + try { + Files.createSymbolicLink(tmp.resolve("link.stale"), tmp.resolve("datafusion-aB3kF7.stale")); + } catch (UnsupportedOperationException | IOException e) { + // Symlinks unsupported here — accept the test on 2 entries instead of 3. + SpillStats statsNoLink = SpillStatsCollector.collect(tmp.toString(), 0L); + assertEquals("dir + file should be counted", 2L, statsNoLink.getStaleEntryCount()); + return; + } + + // Two NON-.stale entries — must be ignored. + Files.createDirectory(tmp.resolve("datafusion-newRandom")); + Files.writeString(tmp.resolve("README.md"), "live"); + + SpillStats stats = SpillStatsCollector.collect(tmp.toString(), 0L); + assertEquals("only the three .stale entries should be counted", 3L, stats.getStaleEntryCount()); } } diff --git a/sandbox/plugins/analytics-backend-datafusion/src/test/java/org/opensearch/be/datafusion/stats/SpillStatsTests.java b/sandbox/plugins/analytics-backend-datafusion/src/test/java/org/opensearch/be/datafusion/stats/SpillStatsTests.java index 227c841c89392..79df17bedbf9d 100644 --- a/sandbox/plugins/analytics-backend-datafusion/src/test/java/org/opensearch/be/datafusion/stats/SpillStatsTests.java +++ b/sandbox/plugins/analytics-backend-datafusion/src/test/java/org/opensearch/be/datafusion/stats/SpillStatsTests.java @@ -20,7 +20,7 @@ public class SpillStatsTests extends OpenSearchTestCase { public void testWireFormatRoundTrip() throws IOException { - SpillStats original = new SpillStats("/mnt/spill", 536_870_912_000L, 214_748_364_800L, 1_073_741_824L, 429_496_729_600L); + SpillStats original = new SpillStats("/mnt/spill", 536_870_912_000L, 214_748_364_800L, 1_073_741_824L, 429_496_729_600L, 7L); BytesStreamOutput out = new BytesStreamOutput(); original.writeTo(out); @@ -28,11 +28,12 @@ public void testWireFormatRoundTrip() throws IOException { try (StreamInput in = out.bytes().streamInput()) { SpillStats roundTripped = new SpillStats(in); assertEquals(original, roundTripped); + assertEquals(7L, roundTripped.getStaleEntryCount()); } } public void testToXContentRendersAllFields() throws IOException { - SpillStats stats = new SpillStats("/mnt/spill", 536_870_912_000L, 214_748_364_800L, 1_073_741_824L, 429_496_729_600L); + SpillStats stats = new SpillStats("/mnt/spill", 536_870_912_000L, 214_748_364_800L, 1_073_741_824L, 429_496_729_600L, 3L); XContentBuilder builder = XContentFactory.jsonBuilder().startObject(); stats.toXContent(builder, ToXContent.EMPTY_PARAMS); @@ -45,22 +46,25 @@ public void testToXContentRendersAllFields() throws IOException { assertTrue(json.contains("\"disk_available_bytes\":214748364800")); assertTrue(json.contains("\"disk_used_bytes\":1073741824")); assertTrue(json.contains("\"disk_reserved_bytes\":429496729600")); + assertTrue(json.contains("\"stale_entry_count\":3")); } public void testEquality() { - SpillStats a = new SpillStats("/mnt/spill", 100L, 50L, 25L, 80L); - SpillStats b = new SpillStats("/mnt/spill", 100L, 50L, 25L, 80L); - SpillStats differentDir = new SpillStats("/var/tmp", 100L, 50L, 25L, 80L); - SpillStats differentTotal = new SpillStats("/mnt/spill", 999L, 50L, 25L, 80L); + SpillStats a = new SpillStats("/mnt/spill", 100L, 50L, 25L, 80L, 0L); + SpillStats b = new SpillStats("/mnt/spill", 100L, 50L, 25L, 80L, 0L); + SpillStats differentDir = new SpillStats("/var/tmp", 100L, 50L, 25L, 80L, 0L); + SpillStats differentTotal = new SpillStats("/mnt/spill", 999L, 50L, 25L, 80L, 0L); + SpillStats differentStaleCount = new SpillStats("/mnt/spill", 100L, 50L, 25L, 80L, 5L); assertEquals(a, b); assertEquals(a.hashCode(), b.hashCode()); assertNotEquals(a, differentDir); assertNotEquals(a, differentTotal); + assertNotEquals("stale_entry_count must participate in equality", a, differentStaleCount); } public void testDisabledStateIsRepresentableWithEmptyDirectoryAndZeroBytes() throws IOException { - SpillStats disabled = new SpillStats("", 0L, 0L, 0L, 0L); + SpillStats disabled = new SpillStats("", 0L, 0L, 0L, 0L, 0L); XContentBuilder builder = XContentFactory.jsonBuilder().startObject(); disabled.toXContent(builder, ToXContent.EMPTY_PARAMS); @@ -69,5 +73,6 @@ public void testDisabledStateIsRepresentableWithEmptyDirectoryAndZeroBytes() thr assertTrue(json.contains("\"directory\":\"\"")); assertTrue(json.contains("\"disk_total_bytes\":0")); + assertTrue(json.contains("\"stale_entry_count\":0")); } } diff --git a/sandbox/qa/analytics-engine-rest/build.gradle b/sandbox/qa/analytics-engine-rest/build.gradle index 78fa9bb1f5d3d..68ed98ddcc648 100644 --- a/sandbox/qa/analytics-engine-rest/build.gradle +++ b/sandbox/qa/analytics-engine-rest/build.gradle @@ -127,6 +127,7 @@ integTest { exclude '**/StreamingCoordinatorReduceIT.class' exclude '**/QueryCacheIT.class' exclude '**/SpillStatsEnabledIT.class' + exclude '**/SpillCleanupOnBootIT.class' exclude '**/YmlOversamplingIT.class' // Note: parallel forks against the same 2-node testCluster slow the suite down @@ -223,6 +224,60 @@ testClusters.integTestSpillEnabled { setting 'datafusion.spill_memory_limit_bytes', '1073741824' } +// ── Spill-cleanup-on-boot variant: pre-seed leaked entries, verify boot wipe ─ +// Separate from integTestSpillEnabled so the stats IT isn't perturbed by seeded +// entries it doesn't expect. +// +// Scope: this variant verifies the end-to-end boot wipe + fresh datafusion-XXXXXX/ +// provisioning. It does NOT reproduce the EACCES boot failure that motivated the +// fix — the seeded directory lives under build/, whose parent the JVM owns. The +// EACCES regression is owned by the Rust unit test +// create_global_runtime_succeeds_when_jvm_does_not_own_spill_parent. +ext.spillCleanupItDir = layout.buildDirectory.dir('spill-cleanup-it').get().asFile.absolutePath + +task integTestSpillCleanup(type: RestIntegTestTask) { + description = 'Verifies boot-time spill cleanup wipes leaked entries and preserves the spill root (content wipe only; EACCES is covered by the Rust unit test).' + testClassesDirs = sourceSets.test.output.classesDirs + classpath = sourceSets.test.runtimeClasspath + filter { + includeTestsMatching 'org.opensearch.analytics.qa.SpillCleanupOnBootIT' + } + systemProperty 'tests.security.manager', 'false' +} +check.dependsOn(integTestSpillCleanup) + +// Seed the spill directory before cluster boot. Cluster start fires at execution +// time (TestClustersPlugin's TaskActionListener.beforeActions, which runs before +// any task's doFirst), so seeding inside doFirst would race the cluster boot. We +// gate on the task graph so ./gradlew tasks / assemble / etc. don't pay the I/O. +gradle.taskGraph.whenReady { graph -> + if (graph.hasTask(integTestSpillCleanup)) { + java.nio.file.Path spillCleanupRoot = java.nio.file.Paths.get(spillCleanupItDir) + java.nio.file.Files.createDirectories(spillCleanupRoot) + File[] spillCleanupExisting = spillCleanupRoot.toFile().listFiles() + if (spillCleanupExisting != null) { + for (File child : spillCleanupExisting) { + if (child.isDirectory()) { + child.deleteDir() + } else { + child.delete() + } + } + } + java.nio.file.Path leakedSubdir = spillCleanupRoot.resolve('datafusion-aB3kF7') + java.nio.file.Files.createDirectories(leakedSubdir) + java.nio.file.Files.writeString(leakedSubdir.resolve('tmp_spill_001.arrow'), 'stale spill data') + java.nio.file.Files.writeString(spillCleanupRoot.resolve('stray.txt'), 'loose top-level leak') + } +} + +testClusters.integTestSpillCleanup { + numberOfNodes = 2 + configureAnalyticsCluster(delegate) + setting 'datafusion.spill_directory', spillCleanupItDir + setting 'datafusion.spill_memory_limit_bytes', '1073741824' +} + // ── YML oversampling variant: verifies TopK fires from opensearch.yml setting ─ task integTestYmlOversampling(type: RestIntegTestTask) { description = 'Verifies TopK oversampling works when set via opensearch.yml (node settings)' diff --git a/sandbox/qa/analytics-engine-rest/src/test/java/org/opensearch/analytics/qa/SpillCleanupOnBootIT.java b/sandbox/qa/analytics-engine-rest/src/test/java/org/opensearch/analytics/qa/SpillCleanupOnBootIT.java new file mode 100644 index 0000000000000..3541d21a8541a --- /dev/null +++ b/sandbox/qa/analytics-engine-rest/src/test/java/org/opensearch/analytics/qa/SpillCleanupOnBootIT.java @@ -0,0 +1,118 @@ +/* + * SPDX-License-Identifier: Apache-2.0 + * + * The OpenSearch Contributors require contributions made to + * this file be licensed under the Apache-2.0 license or a + * compatible open source license. + */ + +package org.opensearch.analytics.qa; + +import com.fasterxml.jackson.databind.JsonNode; +import com.fasterxml.jackson.databind.ObjectMapper; + +import org.apache.hc.core5.http.ParseException; +import org.apache.hc.core5.http.io.entity.EntityUtils; +import org.opensearch.client.Request; +import org.opensearch.client.Response; +import org.opensearch.common.io.PathUtils; +import org.opensearch.test.rest.OpenSearchRestTestCase; + +import java.io.IOException; +import java.nio.file.Files; +import java.nio.file.Path; +import java.util.List; +import java.util.stream.Stream; + +/** + * End-to-end verification that boot-time cleanup wipes leaked spill entries and that + * DataFusion provisions a fresh {@code datafusion-XXXXXX/} on top. The + * {@code integTestSpillCleanup} Gradle task pre-seeds a {@code datafusion-aB3kF7/} + * subtree and a loose {@code stray.txt} into the spill directory before cluster boot; + * this IT asserts: + *

    + *
  1. Cluster booted healthy (200 from the spill stats endpoint).
    2. + *
  2. Seeded entries are gone — cleanup ran.
    3. + *
  3. Spill directory still exists — cleanup did not rmdir the root.
    4. + *
+ * + *

Scope note: this IT does not reproduce the EACCES boot failure + * that motivated the fix — the seeded spill directory lives under + * {@code build/spill-cleanup-it}, whose parent is writable to the test JVM, so the + * pre-fix {@code fs::remove_dir_all(spill_dir)} would also pass these assertions. + * The EACCES regression is owned by the Rust unit test + * {@code create_global_runtime_succeeds_when_jvm_does_not_own_spill_parent} in + * {@code analytics-backend-datafusion/rust/src/api.rs}, which chmods the parent to + * {@code 0o555}. + * + *

Default {@code integTest} excludes this class. + */ +public class SpillCleanupOnBootIT extends OpenSearchRestTestCase { + + private static final ObjectMapper MAPPER = new ObjectMapper(); + private static final String SPILL_ENDPOINT = "/_plugins/_analytics_backend_datafusion/stats/disk_spill"; + + /** Names of the leaked entries pre-seeded by the Gradle task. */ + private static final String LEAKED_SUBDIR = "datafusion-aB3kF7"; + private static final String LEAKED_LOOSE_FILE = "stray.txt"; + + @Override + protected boolean preserveClusterUponCompletion() { + return true; + } + + public void testLeakedEntriesWipedAndSpillRootPreserved() throws Exception { + // Cluster booted: stats endpoint responds. Proves create_global_runtime did not fail. + Response response = client().performRequest(new Request("GET", SPILL_ENDPOINT)); + assertEquals("spill stats endpoint must return 200", 200, response.getStatusLine().getStatusCode()); + + JsonNode root = parseResponse(response); + JsonNode nodes = root.get("nodes"); + assertNotNull("response must contain nodes section", nodes); + assertTrue("expected at least one node", nodes.size() > 0); + + // Pull the spill directory out of the per-node response and assert all nodes agree on it. + String reportedDir = null; + var nodeIds = nodes.fieldNames(); + while (nodeIds.hasNext()) { + String id = nodeIds.next(); + JsonNode spill = nodes.get(id).get("disk_spill"); + assertNotNull("node " + id + " missing spill section", spill); + String dir = spill.get("directory").asText(); + assertFalse("directory must not be empty (node " + id + ")", dir.isEmpty()); + if (reportedDir == null) { + reportedDir = dir; + } else { + assertEquals("all nodes in this IT share one spill directory", reportedDir, dir); + } + } + assertNotNull("at least one node must have reported a spill directory", reportedDir); + + Path spillRoot = PathUtils.get(reportedDir); + + // Spill root preserved, seeded entries gone. + assertTrue("spill directory must still exist after boot: " + spillRoot, Files.isDirectory(spillRoot)); + Path leakedSubdir = spillRoot.resolve(LEAKED_SUBDIR); + Path leakedLooseFile = spillRoot.resolve(LEAKED_LOOSE_FILE); + assertFalse("leaked datafusion-* subdir must be removed: " + leakedSubdir, Files.exists(leakedSubdir)); + assertFalse("leaked top-level file must be removed: " + leakedLooseFile, Files.exists(leakedLooseFile)); + + // DataFusion's tempdir_in provisions at least one fresh datafusion-XXXXXX/ per node. + // None should reuse the seeded leaked name — that would mean cleanup didn't run. + try (Stream children = Files.list(spillRoot)) { + List dfSubdirs = children.filter(Files::isDirectory) + .filter(p -> p.getFileName().toString().startsWith("datafusion-")) + .toList(); + assertFalse("expected at least one fresh datafusion-XXXXXX/ subdir; found none in " + spillRoot, + dfSubdirs.isEmpty()); + for (Path dfDir : dfSubdirs) { + assertNotEquals("fresh subdir must not collide with seeded leaked name: " + dfDir, + LEAKED_SUBDIR, dfDir.getFileName().toString()); + } + } + } + + private JsonNode parseResponse(Response response) throws IOException, ParseException { + return MAPPER.readTree(EntityUtils.toString(response.getEntity())); + } +} From 40f2c9d07a4f3b1245abfa2a80f8d8c999d01b41 Mon Sep 17 00:00:00 2001 From: rayshrey <121871912+rayshrey@users.noreply.github.com> Date: Thu, 11 Jun 2026 13:48:40 +0530 Subject: [PATCH 02/14] Adding virtual memory pools (#21885) * Adding virtual memory pools Signed-off-by: rayshrey * Node stats BWC and other fixes Signed-off-by: rayshrey --------- Signed-off-by: rayshrey --- libs/arrow-spi/build.gradle | 4 +- .../opensearch/arrow/spi/NativeAllocator.java | 122 +++- .../arrow/spi/NativeAllocatorPoolConfig.java | 3 - .../org/opensearch/arrow/spi/PoolGroup.java | 37 ++ .../spi/NativeAllocatorPoolConfigTests.java | 4 - .../arrow/allocator/ArrowBasePlugin.java | 504 +++++++------- .../arrow/allocator/ArrowBaseStatsAction.java | 75 +++ .../arrow/allocator/ArrowNativeAllocator.java | 616 ++++++++++++------ .../allocator/NativeMemoryRebalancer.java | 241 +++++++ .../arrow/allocator/ArrowBasePluginTests.java | 363 ++--------- .../allocator/ArrowNativeAllocatorTests.java | 172 ++--- .../NativeMemoryRebalancerTests.java | 152 +++++ .../arrow/flight/BackpressureProducerIT.java | 3 +- .../flight/NativeAllocatorBoundaryIT.java | 125 +--- .../flight/NativeMemoryRebalancerIT.java | 100 +++ .../flight/UnifiedNativeMemoryStatsIT.java | 95 +++ .../transport/FlightTransportTestBase.java | 4 +- .../dataformat-native/rust/common/src/lib.rs | 1 + .../rust/common/src/memory_pool.rs | 370 +++++++++++ .../be/datafusion/DataFusionPlugin.java | 99 ++- .../be/datafusion/DatafusionSettings.java | 33 + .../DataFusionPluginSettingsTests.java | 15 +- .../datafusion/DatafusionSettingsTests.java | 2 +- .../composite/CompositeMergeIT.java | 5 - .../UnifiedNativeMemoryFullStackIT.java | 84 +++ .../benchmark/VSRRotationBenchmark.java | 9 +- .../parquet/ParquetDataFormatPlugin.java | 60 +- .../opensearch/parquet/ParquetSettings.java | 95 +++ .../opensearch/parquet/bridge/RustBridge.java | 40 ++ .../src/main/rust/src/ffm.rs | 32 + .../src/main/rust/src/lib.rs | 1 + .../src/main/rust/src/memory.rs | 57 ++ .../ParquetDataFormatAwareEngineTests.java | 8 +- .../engine/ParquetIndexingEngineTests.java | 6 +- .../parquet/memory/ArrowBufferPoolTests.java | 4 +- .../parquet/vsr/VSRManagerTests.java | 5 +- .../opensearch/parquet/vsr/VSRPoolTests.java | 4 +- .../parquet/writer/ParquetWriterTests.java | 5 +- server/build.gradle | 1 + .../admin/cluster/node/stats/NodeStats.java | 51 +- .../cluster/node/stats/NodesStatsRequest.java | 2 + .../node/stats/TransportNodesStatsAction.java | 3 +- .../stats/TransportClusterStatsAction.java | 1 - .../indices/IndexingMemoryController.java | 16 +- .../opensearch/indices/IndicesService.java | 8 + .../main/java/org/opensearch/node/Node.java | 22 +- .../java/org/opensearch/node/NodeService.java | 4 +- .../stats/NativeAllocatorPoolStats.java | 179 +++-- .../plugins/SearchBackEndPlugin.java | 39 ++ ...kendNativeMemoryStatsVersionGateTests.java | 166 ----- .../cluster/node/stats/NodeStatsTests.java | 35 +- .../cluster/stats/ClusterStatsNodesTests.java | 1 - .../stats/ClusterStatsResponseTests.java | 1 - .../opensearch/cluster/DiskUsageTests.java | 7 - .../IndexingMemoryControllerTests.java | 19 + .../node/NodeServiceNativeMemoryTests.java | 294 +-------- .../stats/NativeAllocatorPoolStatsTests.java | 33 +- .../MockInternalClusterInfoService.java | 1 - .../opensearch/test/InternalTestCluster.java | 1 - 59 files changed, 2790 insertions(+), 1649 deletions(-) create mode 100644 libs/arrow-spi/src/main/java/org/opensearch/arrow/spi/PoolGroup.java create mode 100644 plugins/arrow-base/src/main/java/org/opensearch/arrow/allocator/ArrowBaseStatsAction.java create mode 100644 plugins/arrow-base/src/main/java/org/opensearch/arrow/allocator/NativeMemoryRebalancer.java create mode 100644 plugins/arrow-base/src/test/java/org/opensearch/arrow/allocator/NativeMemoryRebalancerTests.java create mode 100644 plugins/arrow-flight-rpc/src/internalClusterTest/java/org/opensearch/arrow/flight/NativeMemoryRebalancerIT.java create mode 100644 plugins/arrow-flight-rpc/src/internalClusterTest/java/org/opensearch/arrow/flight/UnifiedNativeMemoryStatsIT.java create mode 100644 sandbox/libs/dataformat-native/rust/common/src/memory_pool.rs create mode 100644 sandbox/plugins/composite-engine/src/internalClusterTest/java/org/opensearch/composite/UnifiedNativeMemoryFullStackIT.java create mode 100644 sandbox/plugins/parquet-data-format/src/main/rust/src/memory.rs delete mode 100644 server/src/test/java/org/opensearch/action/admin/cluster/node/stats/AnalyticsBackendNativeMemoryStatsVersionGateTests.java diff --git a/libs/arrow-spi/build.gradle b/libs/arrow-spi/build.gradle index abf7eecf84c77..c1a716c6f72c4 100644 --- a/libs/arrow-spi/build.gradle +++ b/libs/arrow-spi/build.gradle @@ -11,7 +11,9 @@ apply plugin: 'opensearch.publish' dependencies { api project(':libs:opensearch-core') api project(':libs:opensearch-common') - testImplementation project(':test:framework') + testImplementation(project(':test:framework')) { + exclude group: 'org.opensearch', module: 'opensearch-arrow-spi' + } } tasks.named('forbiddenApisMain').configure { diff --git a/libs/arrow-spi/src/main/java/org/opensearch/arrow/spi/NativeAllocator.java b/libs/arrow-spi/src/main/java/org/opensearch/arrow/spi/NativeAllocator.java index 89d6866da2c89..874d016cc40bb 100644 --- a/libs/arrow-spi/src/main/java/org/opensearch/arrow/spi/NativeAllocator.java +++ b/libs/arrow-spi/src/main/java/org/opensearch/arrow/spi/NativeAllocator.java @@ -9,36 +9,35 @@ package org.opensearch.arrow.spi; import java.io.Closeable; +import java.util.Set; +import java.util.function.Consumer; +import java.util.function.Supplier; /** - * Arrow-agnostic interface for a hierarchical native memory allocator. + * Unified native memory allocator interface. * - *

The implementation (backed by Arrow's {@code RootAllocator}) is provided by - * a plugin. The SPI allows other subsystems to interact with the allocator - * without depending on Arrow classes. - * - *

Plugins that need Arrow allocators obtain the implementation via - * service lookup or plugin extension and call {@link #getOrCreatePool} to - * register their pool. + *

Manages memory pools under a shared budget. Each pool has a minimum + * guaranteed allocation and a maximum burst limit. Implementations may + * redistribute unused capacity across pools. * * @opensearch.api */ public interface NativeAllocator extends Closeable { /** - * Returns the named pool, creating it on first access with the given limit. - * Subsequent calls with the same name return the same pool (first-call limit wins). + * Returns the named pool, creating it on first access. + * Subsequent calls with the same name return the existing pool (first-call config wins). * - * @param poolName logical pool name (e.g., "query", "flight") - * @param limit maximum bytes this pool can allocate in aggregate + * @param poolName logical pool name + * @param min minimum guaranteed bytes + * @param max maximum bytes this pool can allocate + * @param group the group this pool belongs to for aggregated stats, or null * @return an opaque pool handle */ - PoolHandle getOrCreatePool(String poolName, long limit); + PoolHandle getOrCreatePool(String poolName, long min, long max, PoolGroup group); /** - * Updates the limit of an existing pool. Children of the pool allocator - * inherit the change automatically via Arrow's parent-cap check at - * allocation time — no notification SPI is needed. + * Updates the effective limit of an existing pool. * * @param poolName logical pool name * @param newLimit new maximum bytes for the pool @@ -46,15 +45,78 @@ public interface NativeAllocator extends Closeable { void setPoolLimit(String poolName, long newLimit); /** - * Sets the root-level memory limit for the entire allocator. + * Registers a virtual pool with initial min/max and a callback + * invoked when the pool's limit changes. + * + * @param poolName logical pool name + * @param min minimum guaranteed bytes + * @param max initial maximum bytes (the pool's starting limit) + * @param group the group this pool belongs to for aggregated stats + * @param limitSetter callback invoked when the pool limit changes + * @return a handle to update stats from the native layer + */ + VirtualPoolHandle registerVirtualPool(String poolName, long min, long max, PoolGroup group, Consumer limitSetter); + + /** + * Updates the minimum guaranteed bytes for a pool. + * + * @param poolName logical pool name + * @param newMin new minimum bytes + */ + void setPoolMin(String poolName, long newMin); + + /** + * Returns all registered pool names. + */ + Set getAllPoolNames(); + + /** + * Adds a callback invoked before stats collection to refresh pool usage data. + * + * @param refresher runnable that updates pool stats + */ + void addStatsRefresher(Runnable refresher); + + /** + * Sets the supplier for process-wide native memory stats. + * + * @param supplier returns [allocatedBytes, residentBytes] + */ + void setNativeMemoryStatsSupplier(Supplier supplier); + + /** + * Registers a listener invoked when the effective limit of a pool group changes. + * The consumer receives the new total effective limit (sum of all pools in the group). * - * @param limit new maximum bytes for the root allocator + * @param group the pool group to listen on + * @param listener consumer that receives the new grouped limit in bytes + */ + void addPoolGroupLimitListener(PoolGroup group, Consumer listener); + + /** + * Handle for a virtual pool. Plugins update stats via this handle. */ - void setRootLimit(long limit); + interface VirtualPoolHandle { + /** + * Update the current usage stats. + * + * @param allocatedBytes current allocated bytes + * @param peakBytes peak allocated bytes + */ + void updateStats(long allocatedBytes, long peakBytes); + + /** Returns current allocated bytes. */ + long allocatedBytes(); + + /** Returns peak allocated bytes. */ + long peakBytes(); + + /** Returns current limit. */ + long limit(); + } /** - * Opaque handle to a memory pool. Plugins downcast to the concrete type - * (e.g., Arrow's {@code BufferAllocator}) in the implementation layer. + * Opaque handle to a memory pool. */ interface PoolHandle { @@ -63,28 +125,20 @@ interface PoolHandle { * * @param childName name for debugging * @param childLimit maximum bytes for the child - * @return an opaque child handle (downcast to BufferAllocator in Arrow impl) + * @return a child handle */ PoolHandle newChild(String childName, long childLimit); - /** - * Returns the current allocated bytes for this pool/child. - */ + /** Returns the current allocated bytes. */ long allocatedBytes(); - /** - * Returns the peak memory allocation. - */ + /** Returns the peak memory allocation. */ long peakBytes(); - /** - * Returns the configured limit. - */ + /** Returns the configured limit. */ long limit(); - /** - * Releases this allocation handle. - */ + /** Releases this allocation handle. */ void close(); } } diff --git a/libs/arrow-spi/src/main/java/org/opensearch/arrow/spi/NativeAllocatorPoolConfig.java b/libs/arrow-spi/src/main/java/org/opensearch/arrow/spi/NativeAllocatorPoolConfig.java index 98f991cb86704..29dba48e9f165 100644 --- a/libs/arrow-spi/src/main/java/org/opensearch/arrow/spi/NativeAllocatorPoolConfig.java +++ b/libs/arrow-spi/src/main/java/org/opensearch/arrow/spi/NativeAllocatorPoolConfig.java @@ -33,9 +33,6 @@ public final class NativeAllocatorPoolConfig { /** Pool name for query-execution memory (analytics-engine fragments and per-query allocators). */ public static final String POOL_QUERY = "query"; - /** Setting key for the root allocator limit. */ - public static final String SETTING_ROOT_LIMIT = "native.allocator.root.limit"; - /** Setting key for the Flight pool minimum. */ public static final String SETTING_FLIGHT_MIN = "native.allocator.pool.flight.min"; /** Setting key for the Flight pool maximum. */ diff --git a/libs/arrow-spi/src/main/java/org/opensearch/arrow/spi/PoolGroup.java b/libs/arrow-spi/src/main/java/org/opensearch/arrow/spi/PoolGroup.java new file mode 100644 index 0000000000000..d292b42ae595f --- /dev/null +++ b/libs/arrow-spi/src/main/java/org/opensearch/arrow/spi/PoolGroup.java @@ -0,0 +1,37 @@ +/* + * SPDX-License-Identifier: Apache-2.0 + * + * The OpenSearch Contributors require contributions made to + * this file be licensed under the Apache-2.0 license or a + * compatible open source license. + */ + +package org.opensearch.arrow.spi; + +/** + * Groups that memory pools belong to for aggregated customer-facing stats. + * Each pool is assigned to exactly one group at registration time. + * + * @opensearch.api + */ +public enum PoolGroup { + /** Arrow Flight transport pool group. */ + TRANSPORT("transport"), + /** Query and analytics execution pool group. */ + SEARCH("search"), + /** Ingest and write path pool group. */ + INDEXING("indexing"), + /** Background merge operations pool group. */ + MERGE("merge"); + + private final String name; + + PoolGroup(String name) { + this.name = name; + } + + /** Returns the group name used in stats output. */ + public String getName() { + return name; + } +} diff --git a/libs/arrow-spi/src/test/java/org/opensearch/arrow/spi/NativeAllocatorPoolConfigTests.java b/libs/arrow-spi/src/test/java/org/opensearch/arrow/spi/NativeAllocatorPoolConfigTests.java index a21ca8ff54943..025b41a603212 100644 --- a/libs/arrow-spi/src/test/java/org/opensearch/arrow/spi/NativeAllocatorPoolConfigTests.java +++ b/libs/arrow-spi/src/test/java/org/opensearch/arrow/spi/NativeAllocatorPoolConfigTests.java @@ -26,8 +26,4 @@ public void testSettingKeys() { assertEquals("native.allocator.pool.query.min", NativeAllocatorPoolConfig.SETTING_QUERY_MIN); assertEquals("native.allocator.pool.query.max", NativeAllocatorPoolConfig.SETTING_QUERY_MAX); } - - public void testRootSettingKey() { - assertEquals("native.allocator.root.limit", NativeAllocatorPoolConfig.SETTING_ROOT_LIMIT); - } } diff --git a/plugins/arrow-base/src/main/java/org/opensearch/arrow/allocator/ArrowBasePlugin.java b/plugins/arrow-base/src/main/java/org/opensearch/arrow/allocator/ArrowBasePlugin.java index a9bd9968b5884..6a2c52a46f9ce 100644 --- a/plugins/arrow-base/src/main/java/org/opensearch/arrow/allocator/ArrowBasePlugin.java +++ b/plugins/arrow-base/src/main/java/org/opensearch/arrow/allocator/ArrowBasePlugin.java @@ -9,11 +9,16 @@ package org.opensearch.arrow.allocator; import org.opensearch.arrow.spi.NativeAllocatorPoolConfig; +import org.opensearch.arrow.spi.PoolGroup; import org.opensearch.cluster.metadata.IndexNameExpressionResolver; +import org.opensearch.cluster.node.DiscoveryNodes; import org.opensearch.cluster.service.ClusterService; import org.opensearch.common.settings.ClusterSettings; +import org.opensearch.common.settings.IndexScopedSettings; import org.opensearch.common.settings.Setting; import org.opensearch.common.settings.Settings; +import org.opensearch.common.settings.SettingsFilter; +import org.opensearch.common.util.concurrent.FutureUtils; import org.opensearch.core.common.io.stream.NamedWriteableRegistry; import org.opensearch.core.common.unit.ByteSizeValue; import org.opensearch.core.xcontent.NamedXContentRegistry; @@ -22,10 +27,14 @@ import org.opensearch.node.resource.tracker.ResourceTrackerSettings; import org.opensearch.plugin.stats.NativeAllocatorPoolStats; import org.opensearch.plugin.stats.NativeAllocatorStatsRegistry; +import org.opensearch.plugins.ActionPlugin; import org.opensearch.plugins.ExtensiblePlugin; import org.opensearch.plugins.Plugin; import org.opensearch.repositories.RepositoriesService; +import org.opensearch.rest.RestController; +import org.opensearch.rest.RestHandler; import org.opensearch.script.ScriptService; +import org.opensearch.threadpool.Scheduler; import org.opensearch.threadpool.ThreadPool; import org.opensearch.transport.client.Client; import org.opensearch.watcher.ResourceWatcherService; @@ -33,207 +42,133 @@ import java.io.IOException; import java.util.Collection; import java.util.List; +import java.util.concurrent.ScheduledExecutorService; +import java.util.concurrent.ScheduledFuture; +import java.util.concurrent.TimeUnit; import java.util.function.Supplier; /** * Top-level plugin that owns the unified Arrow-backed native memory allocator. * - *

All Arrow-consuming plugins (arrow-flight-rpc, parquet-data-format) extend - * this plugin to share one {@link ArrowNativeAllocator} and its classloader. - * - *

Each pool has a min (guaranteed floor) and max (burst ceiling). The rebalancer - * ensures every pool can always allocate up to its min, and distributes unused - * capacity allowing pools to grow up to their max. + *

All Arrow-consuming plugins extend this plugin to share one + * {@link ArrowNativeAllocator} and its classloader. */ -public class ArrowBasePlugin extends Plugin implements ExtensiblePlugin { +public class ArrowBasePlugin extends Plugin implements ExtensiblePlugin, ActionPlugin { /** Creates the plugin. */ public ArrowBasePlugin() {} - /** - * Maximum bytes for the root Arrow allocator. - * - *

When unset, the default is 20% of - * {@link ResourceTrackerSettings#NODE_NATIVE_MEMORY_LIMIT_SETTING}; see - * {@link #deriveRootLimitDefault}. The Arrow framework gets a small fraction of the - * native budget because the dominant consumer of native memory in analytics workloads - * is the DataFusion Rust runtime (~75% of {@code node.native_memory.limit}), not Arrow. - * If AC is unconfigured (limit = 0), the default is {@link Long#MAX_VALUE}, preserving - * pre-AC behaviour. - */ - public static final Setting ROOT_LIMIT_SETTING = new Setting<>( - NativeAllocatorPoolConfig.SETTING_ROOT_LIMIT, - ArrowBasePlugin::deriveRootLimitDefault, - s -> { - long v = Long.parseLong(s); - if (v < 0) { - throw new IllegalArgumentException("Setting [" + NativeAllocatorPoolConfig.SETTING_ROOT_LIMIT + "] must be >= 0, got " + v); - } - return v; - }, + // ─── Settings ──────────────────────────────────────────────────────────────── + + /** Whether the NativeMemoryRebalancer is enabled. */ + public static final Setting REBALANCER_ENABLED_SETTING = Setting.boolSetting( + "native.allocator.rebalancer.enabled", + true, Setting.Property.NodeScope, Setting.Property.Dynamic ); - /** - * Computes the default for {@link #ROOT_LIMIT_SETTING} as 20% of - * {@link ResourceTrackerSettings#NODE_NATIVE_MEMORY_LIMIT_SETTING}. The Arrow framework's - * hard cap covers only Arrow allocations — DataFusion's Rust runtime is a sibling of - * Arrow root and gets the larger share of the native budget (see - * {@code DataFusionPlugin#deriveMemoryPoolLimitDefault}). - * - *

Returns the bytes-as-string representation expected by the {@link Setting} parser. - * If the AC limit is unset (== 0), the default is {@link Long#MAX_VALUE} — unbounded — - * preserving pre-AC behaviour. - */ - static String deriveRootLimitDefault(Settings settings) { - ByteSizeValue nativeLimit = ResourceTrackerSettings.NODE_NATIVE_MEMORY_LIMIT_SETTING.get(settings); - if (nativeLimit.getBytes() <= 0) { - return Long.toString(Long.MAX_VALUE); - } - return Long.toString(nativeLimit.getBytes() * 20 / 100); - } + /** Interval in seconds between pool rebalance cycles. 0 disables rebalancing. */ + public static final Setting REBALANCE_INTERVAL_SETTING = Setting.longSetting( + "native.allocator.rebalance.interval_seconds", + 5L, + 0L, + Setting.Property.NodeScope, + Setting.Property.Dynamic + ); + + /** Pool utilization above this triggers growth. */ + public static final Setting PRESSURE_THRESHOLD_SETTING = Setting.doubleSetting( + "native.allocator.rebalancer.pressure_threshold", + 0.75, + 0.0, + 1.0, + Setting.Property.NodeScope, + Setting.Property.Dynamic + ); + + /** Pool utilization below this means pool can give back capacity. */ + public static final Setting IDLE_THRESHOLD_SETTING = Setting.doubleSetting( + "native.allocator.rebalancer.idle_threshold", + 0.50, + 0.0, + 1.0, + Setting.Property.NodeScope, + Setting.Property.Dynamic + ); - /** Minimum guaranteed bytes for the Flight pool. */ - public static final Setting FLIGHT_MIN_SETTING = Setting.longSetting( + /** Factor to shrink idle pools by (new limit = limit * (1 - shrink_factor)). */ + public static final Setting SHRINK_FACTOR_SETTING = Setting.doubleSetting( + "native.allocator.rebalancer.shrink_factor", + 0.10, + 0.0, + 1.0, + Setting.Property.NodeScope, + Setting.Property.Dynamic + ); + + /** Minimum guaranteed bytes for the Flight pool. Default is 2% of budget. */ + public static final Setting FLIGHT_MIN_SETTING = new Setting<>( NativeAllocatorPoolConfig.SETTING_FLIGHT_MIN, - 0L, - 0L, + s -> derivePoolMinDefault(s, 2), + s -> parseNonNegativeLong(s, NativeAllocatorPoolConfig.SETTING_FLIGHT_MIN), Setting.Property.NodeScope, Setting.Property.Dynamic ); - /** - * Maximum bytes the Flight pool can burst to. Default is 5% of - * {@link ResourceTrackerSettings#NODE_NATIVE_MEMORY_LIMIT_SETTING}; see - * {@link #derivePoolMaxDefault}. Falls back to {@link Long#MAX_VALUE} when AC is - * unconfigured. Matches the partitioning model documented in PR #21732. - */ + /** Maximum bytes the Flight pool can burst to. Default is 5% of budget. */ public static final Setting FLIGHT_MAX_SETTING = new Setting<>( NativeAllocatorPoolConfig.SETTING_FLIGHT_MAX, s -> derivePoolMaxDefault(s, 5), - s -> { - long v = Long.parseLong(s); - if (v < 0) { - throw new IllegalArgumentException("Setting [" + NativeAllocatorPoolConfig.SETTING_FLIGHT_MAX + "] must be >= 0, got " + v); - } - return v; - }, + s -> parseNonNegativeLong(s, NativeAllocatorPoolConfig.SETTING_FLIGHT_MAX), Setting.Property.NodeScope, Setting.Property.Dynamic ); - /** Minimum guaranteed bytes for the ingest pool. */ - public static final Setting INGEST_MIN_SETTING = Setting.longSetting( + /** Minimum guaranteed bytes for the ingest pool. Default is 4% of budget. */ + public static final Setting INGEST_MIN_SETTING = new Setting<>( NativeAllocatorPoolConfig.SETTING_INGEST_MIN, - 0L, - 0L, + s -> derivePoolMinDefault(s, 4), + s -> parseNonNegativeLong(s, NativeAllocatorPoolConfig.SETTING_INGEST_MIN), Setting.Property.NodeScope, Setting.Property.Dynamic ); - /** - * Maximum bytes the ingest pool can burst to. Default is 8% of - * {@link ResourceTrackerSettings#NODE_NATIVE_MEMORY_LIMIT_SETTING}; see - * {@link #derivePoolMaxDefault}. Falls back to {@link Long#MAX_VALUE} when AC is - * unconfigured. Ingest gets a larger fraction than Flight/Query because parquet VSR - * allocators dominate write-path memory usage — see partitioning model in PR #21732. - */ + /** Maximum bytes the ingest pool can burst to. Default is 8% of budget. */ public static final Setting INGEST_MAX_SETTING = new Setting<>( NativeAllocatorPoolConfig.SETTING_INGEST_MAX, s -> derivePoolMaxDefault(s, 8), - s -> { - long v = Long.parseLong(s); - if (v < 0) { - throw new IllegalArgumentException("Setting [" + NativeAllocatorPoolConfig.SETTING_INGEST_MAX + "] must be >= 0, got " + v); - } - return v; - }, + s -> parseNonNegativeLong(s, NativeAllocatorPoolConfig.SETTING_INGEST_MAX), Setting.Property.NodeScope, Setting.Property.Dynamic ); - /** - * Minimum guaranteed bytes for the query pool. Honored by the rebalancer (when - * enabled) — sets a floor below which the rebalancer will not shrink the pool. - * Has no effect when rebalancing is disabled. - */ - public static final Setting QUERY_MIN_SETTING = Setting.longSetting( + /** Minimum guaranteed bytes for the query pool. Default is 2% of budget. */ + public static final Setting QUERY_MIN_SETTING = new Setting<>( NativeAllocatorPoolConfig.SETTING_QUERY_MIN, - 0L, - 0L, + s -> derivePoolMinDefault(s, 2), + s -> parseNonNegativeLong(s, NativeAllocatorPoolConfig.SETTING_QUERY_MIN), Setting.Property.NodeScope, Setting.Property.Dynamic ); - /** - * Maximum bytes the query pool can allocate. Default is 5% of - * {@link ResourceTrackerSettings#NODE_NATIVE_MEMORY_LIMIT_SETTING}; see - * {@link #derivePoolMaxDefault}. Falls back to {@link Long#MAX_VALUE} when AC is - * unconfigured. Enforced by Arrow's child-allocator limit — analytics-engine's - * per-query allocators are children of this pool, so the sum of in-flight per-query - * allocations is capped here. - * - *

Note: each individual analytics query is also bounded by - * {@code analytics.exec.QueryContext} per-query limit (currently the constant - * {@code DEFAULT_PER_QUERY_MEMORY_LIMIT = 256 MB}). Lowering {@code QUERY_MAX} - * below {@code 256 MB × concurrent-queries} can starve queries even when each - * individual query is within its per-query limit. - */ + /** Maximum bytes the query pool can allocate. Default is 5% of budget. */ public static final Setting QUERY_MAX_SETTING = new Setting<>( NativeAllocatorPoolConfig.SETTING_QUERY_MAX, s -> derivePoolMaxDefault(s, 5), - s -> { - long v = Long.parseLong(s); - if (v < 0) { - throw new IllegalArgumentException("Setting [" + NativeAllocatorPoolConfig.SETTING_QUERY_MAX + "] must be >= 0, got " + v); - } - return v; - }, + s -> parseNonNegativeLong(s, NativeAllocatorPoolConfig.SETTING_QUERY_MAX), Setting.Property.NodeScope, Setting.Property.Dynamic ); - /** - * Computes the default for a pool max as a percentage of - * {@link ResourceTrackerSettings#NODE_NATIVE_MEMORY_LIMIT_SETTING} (the operator's - * declared off-heap budget), falling back to {@link Long#MAX_VALUE} when AC is - * unconfigured. Returns the bytes-as-string representation expected by the - * {@link Setting} parser. - * - *

Pools are anchored to {@code node.native_memory.limit} rather than to - * {@link #ROOT_LIMIT_SETTING} so the diagrammed partitioning (PR #21732) holds: - * sum of pool maxes (5+8+5 = 18% of native_memory.limit) fits within the framework - * root cap (20% of native_memory.limit) by default. Operator overrides of - * {@code root.limit} that drop it below {@code sum(pool.max)} are caught by the - * grouped validator. - * - *

The fraction is taken straight from {@code node.native_memory.limit}, not from - * {@code limit - buffer_percent}. {@code buffer_percent} is an admission-control - * throttle margin, not a framework budget reduction. - * - * @param settings node settings - * @param percent fraction of {@code node.native_memory.limit} the pool max defaults to - */ - static String derivePoolMaxDefault(Settings settings, int percent) { - ByteSizeValue nativeLimit = ResourceTrackerSettings.NODE_NATIVE_MEMORY_LIMIT_SETTING.get(settings); - if (nativeLimit.getBytes() <= 0) { - return Long.toString(Long.MAX_VALUE); - } - long pool = Math.max(0L, nativeLimit.getBytes() * percent / 100); - return Long.toString(pool); - } - - /** Interval in seconds between pool rebalance cycles. 0 disables rebalancing. */ - public static final Setting REBALANCE_INTERVAL_SETTING = Setting.longSetting( - "native.allocator.rebalance.interval_seconds", - 0L, - 0L, - Setting.Property.NodeScope, - Setting.Property.Dynamic - ); + // ─── Instance state ────────────────────────────────────────────────────────── private volatile ArrowNativeAllocator allocator; + private volatile ScheduledExecutorService rebalancerScheduler; + private volatile ScheduledFuture rebalanceTask; + private volatile NativeMemoryRebalancer rebalancer; + + // ─── Plugin lifecycle ──────────────────────────────────────────────────────── @Override public Collection createComponents( @@ -251,12 +186,11 @@ public Collection createComponents( ) { Settings settings = environment.settings(); ClusterSettings cs = clusterService.getClusterSettings(); - ArrowNativeAllocator built = buildAllocator(settings, cs); + Supplier budgetSupplier = () -> ResourceTrackerSettings.NODE_NATIVE_MEMORY_LIMIT_SETTING.get(clusterService.getSettings()) + .getBytes(); + ArrowNativeAllocator built = buildAllocator(settings, cs, budgetSupplier); this.allocator = built; - // Publish a NativeAllocatorStatsRegistry alongside the allocator so the server-side - // NodeService can discover the supplier via pluginComponents (instanceof filter) without - // taking a compile-time dependency on this plugin. The lambda re-reads `this.allocator` - // each invocation, so after close() nulls the field, the supplier returns null cleanly. + Supplier statsSupplier = () -> { ArrowNativeAllocator a = this.allocator; return a != null ? a.stats() : null; @@ -264,96 +198,177 @@ public Collection createComponents( return List.of(built, new NativeAllocatorStatsRegistry(statsSupplier)); } + @Override + public List> getSettings() { + return List.of( + FLIGHT_MIN_SETTING, + FLIGHT_MAX_SETTING, + INGEST_MIN_SETTING, + INGEST_MAX_SETTING, + QUERY_MIN_SETTING, + QUERY_MAX_SETTING, + REBALANCE_INTERVAL_SETTING, + REBALANCER_ENABLED_SETTING, + PRESSURE_THRESHOLD_SETTING, + IDLE_THRESHOLD_SETTING, + SHRINK_FACTOR_SETTING + ); + } + + @Override + public List getRestHandlers( + Settings settings, + RestController restController, + ClusterSettings clusterSettings, + IndexScopedSettings indexScopedSettings, + SettingsFilter settingsFilter, + IndexNameExpressionResolver indexNameExpressionResolver, + Supplier nodesInCluster + ) { + Supplier statsSupplier = () -> allocator != null ? allocator.stats() : null; + return List.of(new ArrowBaseStatsAction(statsSupplier)); + } + + @Override + public void close() throws IOException { + if (rebalancerScheduler != null) { + rebalancerScheduler.shutdownNow(); + } + if (allocator != null) { + allocator.close(); + allocator = null; + } + } + + // ─── Package-private (visible for tests) ───────────────────────────────────── + /** - * Constructs the allocator and wires its pools and dynamic-update consumers from - * a pure {@code (Settings, ClusterSettings)} pair. Package-private so unit tests - * can exercise the full wiring without a heavyweight {@link ClusterService} - * fixture — mirrors the shape of {@link #registerSettingsUpdateConsumers} which - * is already test-friendly for the same reason. + * Constructs the allocator and wires its pools and the rebalancer. */ - static ArrowNativeAllocator buildAllocator(Settings settings, ClusterSettings cs) { - long rootLimit = ROOT_LIMIT_SETTING.get(settings); - ArrowNativeAllocator allocator = new ArrowNativeAllocator(rootLimit); - allocator.setRebalanceInterval(REBALANCE_INTERVAL_SETTING.get(settings)); + ArrowNativeAllocator buildAllocator(Settings settings, ClusterSettings cs, Supplier budgetSupplier) { + ArrowNativeAllocator allocator = new ArrowNativeAllocator(); - // Single source of truth for cross-setting invariants — same logic runs on - // dynamic updates via the grouped consumer below. - validateUpdate(settings); + // Set budget for validation + long nativeBudget = ResourceTrackerSettings.NODE_NATIVE_MEMORY_LIMIT_SETTING.get(settings).getBytes(); + if (nativeBudget > 0) { + allocator.setBudget(nativeBudget); + } + // Validate min < max for each pool + validateMinMax(NativeAllocatorPoolConfig.POOL_FLIGHT, FLIGHT_MIN_SETTING.get(settings), FLIGHT_MAX_SETTING.get(settings)); + validateMinMax(NativeAllocatorPoolConfig.POOL_INGEST, INGEST_MIN_SETTING.get(settings), INGEST_MAX_SETTING.get(settings)); + validateMinMax(NativeAllocatorPoolConfig.POOL_QUERY, QUERY_MIN_SETTING.get(settings), QUERY_MAX_SETTING.get(settings)); + + // Create pools (always start at max) allocator.getOrCreatePool( NativeAllocatorPoolConfig.POOL_FLIGHT, FLIGHT_MIN_SETTING.get(settings), - FLIGHT_MAX_SETTING.get(settings) + FLIGHT_MAX_SETTING.get(settings), + PoolGroup.TRANSPORT ); allocator.getOrCreatePool( NativeAllocatorPoolConfig.POOL_INGEST, INGEST_MIN_SETTING.get(settings), - INGEST_MAX_SETTING.get(settings) + INGEST_MAX_SETTING.get(settings), + PoolGroup.INDEXING + ); + allocator.getOrCreatePool( + NativeAllocatorPoolConfig.POOL_QUERY, + QUERY_MIN_SETTING.get(settings), + QUERY_MAX_SETTING.get(settings), + PoolGroup.SEARCH ); - allocator.getOrCreatePool(NativeAllocatorPoolConfig.POOL_QUERY, QUERY_MIN_SETTING.get(settings), QUERY_MAX_SETTING.get(settings)); - registerSettingsUpdateConsumers(cs, allocator); + // Register dynamic setting consumers for min/max changes + cs.addSettingsUpdateConsumer(FLIGHT_MIN_SETTING, newMin -> allocator.setPoolMin(NativeAllocatorPoolConfig.POOL_FLIGHT, newMin)); + cs.addSettingsUpdateConsumer(FLIGHT_MAX_SETTING, newMax -> allocator.setPoolLimit(NativeAllocatorPoolConfig.POOL_FLIGHT, newMax)); + cs.addSettingsUpdateConsumer(INGEST_MIN_SETTING, newMin -> allocator.setPoolMin(NativeAllocatorPoolConfig.POOL_INGEST, newMin)); + cs.addSettingsUpdateConsumer(INGEST_MAX_SETTING, newMax -> allocator.setPoolLimit(NativeAllocatorPoolConfig.POOL_INGEST, newMax)); + cs.addSettingsUpdateConsumer(QUERY_MIN_SETTING, newMin -> allocator.setPoolMin(NativeAllocatorPoolConfig.POOL_QUERY, newMin)); + cs.addSettingsUpdateConsumer(QUERY_MAX_SETTING, newMax -> allocator.setPoolLimit(NativeAllocatorPoolConfig.POOL_QUERY, newMax)); + + // Register dynamic consumer for rebalancer enable/disable + cs.addSettingsUpdateConsumer(REBALANCER_ENABLED_SETTING, enabled -> { + if (enabled == false) { + cancelRebalanceTask(); + allocator.resetAllPoolsToMax(); + } else { + startRebalancer(allocator, budgetSupplier, cs.get(REBALANCE_INTERVAL_SETTING)); + } + }); + + // Set up the rebalancer if enabled + if (REBALANCER_ENABLED_SETTING.get(settings)) { + startRebalancer(allocator, budgetSupplier, REBALANCE_INTERVAL_SETTING.get(settings)); + } + + // Register dynamic consumer for interval changes + cs.addSettingsUpdateConsumer(REBALANCE_INTERVAL_SETTING, this::updateRebalanceInterval); + + // Register dynamic consumers for threshold changes + cs.addSettingsUpdateConsumer(PRESSURE_THRESHOLD_SETTING, value -> { + NativeMemoryRebalancer r = this.rebalancer; + if (r != null) r.setPressureThreshold(value); + }); + cs.addSettingsUpdateConsumer(IDLE_THRESHOLD_SETTING, value -> { + NativeMemoryRebalancer r = this.rebalancer; + if (r != null) r.setIdleThreshold(value); + }); + cs.addSettingsUpdateConsumer(SHRINK_FACTOR_SETTING, value -> { + NativeMemoryRebalancer r = this.rebalancer; + if (r != null) r.setShrinkFactor(value); + }); + return allocator; } - /** - * Registers cluster-settings update consumers that propagate dynamic setting changes - * into the live {@link ArrowNativeAllocator}. Package-private so unit tests can exercise - * the wiring with a real {@link ClusterSettings} instance — the test that asserts a PUT - * lands on the allocator is what catches a future regression where one of these lines - * is accidentally removed. - */ - static void registerSettingsUpdateConsumers(ClusterSettings cs, ArrowNativeAllocator allocator) { - cs.addSettingsUpdateConsumer(ROOT_LIMIT_SETTING, allocator::setRootLimit); - cs.addSettingsUpdateConsumer(REBALANCE_INTERVAL_SETTING, allocator::setRebalanceInterval); - cs.addSettingsUpdateConsumer(FLIGHT_MAX_SETTING, v -> allocator.setPoolLimit(NativeAllocatorPoolConfig.POOL_FLIGHT, v)); - cs.addSettingsUpdateConsumer(FLIGHT_MIN_SETTING, v -> allocator.setPoolMin(NativeAllocatorPoolConfig.POOL_FLIGHT, v)); - cs.addSettingsUpdateConsumer(INGEST_MAX_SETTING, v -> allocator.setPoolLimit(NativeAllocatorPoolConfig.POOL_INGEST, v)); - cs.addSettingsUpdateConsumer(INGEST_MIN_SETTING, v -> allocator.setPoolMin(NativeAllocatorPoolConfig.POOL_INGEST, v)); - cs.addSettingsUpdateConsumer(QUERY_MAX_SETTING, v -> allocator.setPoolLimit(NativeAllocatorPoolConfig.POOL_QUERY, v)); - cs.addSettingsUpdateConsumer(QUERY_MIN_SETTING, v -> allocator.setPoolMin(NativeAllocatorPoolConfig.POOL_QUERY, v)); - - // Grouped validator runs across the related settings on every dynamic update so cross-setting - // invariants (sum of pool mins ≤ root, per-pool min ≤ max) are enforced post-startup. - cs.addSettingsUpdateConsumer(s -> {}, MIN_MAX_SETTINGS, ArrowBasePlugin::validateUpdate); - } + // ─── Private helpers ───────────────────────────────────────────────────────── - private static final List> MIN_MAX_SETTINGS = List.of( - ROOT_LIMIT_SETTING, - FLIGHT_MIN_SETTING, - FLIGHT_MAX_SETTING, - INGEST_MIN_SETTING, - INGEST_MAX_SETTING, - QUERY_MIN_SETTING, - QUERY_MAX_SETTING - ); + private synchronized void startRebalancer(ArrowNativeAllocator allocator, Supplier budgetSupplier, long intervalSeconds) { + if (rebalancer != null || rebalancerScheduler != null) return; - private static void validateUpdate(Settings settings) { - long rootLimit = ROOT_LIMIT_SETTING.get(settings); - long flightMin = FLIGHT_MIN_SETTING.get(settings); - long flightMax = FLIGHT_MAX_SETTING.get(settings); - long ingestMin = INGEST_MIN_SETTING.get(settings); - long ingestMax = INGEST_MAX_SETTING.get(settings); - long queryMin = QUERY_MIN_SETTING.get(settings); - long queryMax = QUERY_MAX_SETTING.get(settings); - validateMinMax(NativeAllocatorPoolConfig.POOL_FLIGHT, flightMin, flightMax); - validateMinMax(NativeAllocatorPoolConfig.POOL_INGEST, ingestMin, ingestMax); - validateMinMax(NativeAllocatorPoolConfig.POOL_QUERY, queryMin, queryMax); - validateMinSum(rootLimit, flightMin, ingestMin, queryMin); - } + long budget = budgetSupplier.get(); + if (budget <= 0) return; + if (intervalSeconds <= 0) return; - @Override - public List> getSettings() { - return List.of( - ROOT_LIMIT_SETTING, - FLIGHT_MIN_SETTING, - FLIGHT_MAX_SETTING, - INGEST_MIN_SETTING, - INGEST_MAX_SETTING, - QUERY_MIN_SETTING, - QUERY_MAX_SETTING, - REBALANCE_INTERVAL_SETTING + NativeMemoryRebalancer nativeRebalancer = new NativeMemoryRebalancer( + allocator, + budgetSupplier, + PRESSURE_THRESHOLD_SETTING.getDefault(Settings.EMPTY), + IDLE_THRESHOLD_SETTING.getDefault(Settings.EMPTY), + SHRINK_FACTOR_SETTING.getDefault(Settings.EMPTY) ); + this.rebalancer = nativeRebalancer; + + Scheduler.SafeScheduledThreadPoolExecutor executor = new Scheduler.SafeScheduledThreadPoolExecutor(1, r -> { + Thread t = new Thread(r, "native-allocator-rebalancer"); + t.setDaemon(true); + return t; + }); + executor.setRemoveOnCancelPolicy(true); + this.rebalancerScheduler = executor; + + rebalanceTask = rebalancerScheduler.scheduleAtFixedRate(nativeRebalancer, intervalSeconds, intervalSeconds, TimeUnit.SECONDS); + } + + private synchronized void cancelRebalanceTask() { + ScheduledFuture existing = rebalanceTask; + if (existing != null) { + FutureUtils.cancel(existing); + rebalanceTask = null; + } + rebalancer = null; + if (rebalancerScheduler != null) { + rebalancerScheduler.shutdown(); + rebalancerScheduler = null; + } + } + + private void updateRebalanceInterval(long newInterval) { + cancelRebalanceTask(); + if (newInterval > 0 && rebalancerScheduler != null && rebalancer != null) { + rebalanceTask = rebalancerScheduler.scheduleAtFixedRate(rebalancer, newInterval, newInterval, TimeUnit.SECONDS); + } } private static void validateMinMax(String poolName, long min, long max) { @@ -362,36 +377,27 @@ private static void validateMinMax(String poolName, long min, long max) { } } - private static void validateMinSum(long rootLimit, long... mins) { - if (rootLimit == Long.MAX_VALUE) { - return; - } - long sum = 0; - for (long min : mins) { - try { - sum = Math.addExact(sum, min); - } catch (ArithmeticException overflow) { - throw new IllegalArgumentException("Sum of pool minimums overflows.", overflow); - } + static String derivePoolMaxDefault(Settings settings, int percent) { + ByteSizeValue nativeLimit = ResourceTrackerSettings.NODE_NATIVE_MEMORY_LIMIT_SETTING.get(settings); + if (nativeLimit.getBytes() <= 0) { + return Long.toString(Long.MAX_VALUE); } - if (sum > rootLimit) { - throw new IllegalArgumentException( - "Sum of pool minimums (" - + sum - + " bytes) exceeds root limit (" - + rootLimit - + " bytes). " - + "Reduce pool minimums or increase " - + NativeAllocatorPoolConfig.SETTING_ROOT_LIMIT - ); + return Long.toString(Math.max(0L, nativeLimit.getBytes() * percent / 100)); + } + + static String derivePoolMinDefault(Settings settings, int percent) { + ByteSizeValue nativeLimit = ResourceTrackerSettings.NODE_NATIVE_MEMORY_LIMIT_SETTING.get(settings); + if (nativeLimit.getBytes() <= 0) { + return "0"; } + return Long.toString(Math.max(0L, nativeLimit.getBytes() * percent / 100)); } - @Override - public void close() throws IOException { - if (allocator != null) { - allocator.close(); - allocator = null; + private static long parseNonNegativeLong(String s, String settingName) { + long v = Long.parseLong(s); + if (v < 0) { + throw new IllegalArgumentException("Setting [" + settingName + "] must be >= 0, got " + v); } + return v; } } diff --git a/plugins/arrow-base/src/main/java/org/opensearch/arrow/allocator/ArrowBaseStatsAction.java b/plugins/arrow-base/src/main/java/org/opensearch/arrow/allocator/ArrowBaseStatsAction.java new file mode 100644 index 0000000000000..70928e93713a0 --- /dev/null +++ b/plugins/arrow-base/src/main/java/org/opensearch/arrow/allocator/ArrowBaseStatsAction.java @@ -0,0 +1,75 @@ +/* + * SPDX-License-Identifier: Apache-2.0 + * + * The OpenSearch Contributors require contributions made to + * this file be licensed under the Apache-2.0 license or a + * compatible open source license. + */ + +package org.opensearch.arrow.allocator; + +import org.opensearch.core.rest.RestStatus; +import org.opensearch.core.xcontent.XContentBuilder; +import org.opensearch.plugin.stats.NativeAllocatorPoolStats; +import org.opensearch.rest.BaseRestHandler; +import org.opensearch.rest.BytesRestResponse; +import org.opensearch.rest.RestRequest; +import org.opensearch.transport.client.node.NodeClient; + +import java.util.List; +import java.util.function.Supplier; + +/** + * REST handler exposing per-pool native memory stats at {@code _plugins/arrow_base/stats}. + */ +public class ArrowBaseStatsAction extends BaseRestHandler { + private final Supplier statsSupplier; + + /** + * Creates a new stats action. + * @param statsSupplier supplier of pool stats + */ + public ArrowBaseStatsAction(Supplier statsSupplier) { + this.statsSupplier = statsSupplier; + } + + @Override + public String getName() { + return "arrow_base_stats_action"; + } + + @Override + public List routes() { + return List.of(new Route(RestRequest.Method.GET, "_plugins/arrow_base/stats")); + } + + @Override + protected RestChannelConsumer prepareRequest(RestRequest request, NodeClient client) { + return channel -> { + NativeAllocatorPoolStats stats = statsSupplier.get(); + XContentBuilder builder = channel.newBuilder(); + builder.startObject(); + builder.startObject("memory_pools"); + if (stats != null) { + builder.startObject("runtime"); + builder.field("allocated_bytes", stats.getNativeAllocatedBytes()); + builder.field("resident_bytes", stats.getNativeResidentBytes()); + builder.endObject(); + builder.startObject("pools"); + for (NativeAllocatorPoolStats.PoolStats pool : stats.getPools()) { + builder.startObject(pool.getName()); + builder.field("allocated_bytes", pool.getAllocatedBytes()); + builder.field("peak_bytes", pool.getPeakBytes()); + builder.field("limit_bytes", pool.getLimitBytes()); + builder.field("min_bytes", pool.getMinBytes()); + builder.field("group", pool.getGroup()); + builder.endObject(); + } + builder.endObject(); + } + builder.endObject(); + builder.endObject(); + channel.sendResponse(new BytesRestResponse(RestStatus.OK, builder)); + }; + } +} diff --git a/plugins/arrow-base/src/main/java/org/opensearch/arrow/allocator/ArrowNativeAllocator.java b/plugins/arrow-base/src/main/java/org/opensearch/arrow/allocator/ArrowNativeAllocator.java index 892c43d2cb2c8..baf5ff398b760 100644 --- a/plugins/arrow-base/src/main/java/org/opensearch/arrow/allocator/ArrowNativeAllocator.java +++ b/plugins/arrow-base/src/main/java/org/opensearch/arrow/allocator/ArrowNativeAllocator.java @@ -10,197 +10,296 @@ import org.apache.arrow.memory.BufferAllocator; import org.apache.arrow.memory.RootAllocator; +import org.apache.logging.log4j.LogManager; +import org.apache.logging.log4j.Logger; import org.opensearch.arrow.spi.NativeAllocator; +import org.opensearch.arrow.spi.PoolGroup; +import org.opensearch.common.settings.ClusterSettings; +import org.opensearch.common.settings.Setting; +import org.opensearch.common.settings.Settings; import org.opensearch.plugin.stats.NativeAllocatorPoolStats; import java.util.ArrayList; import java.util.Collections; +import java.util.HashSet; import java.util.List; -import java.util.Map; import java.util.Set; import java.util.concurrent.ConcurrentHashMap; import java.util.concurrent.ConcurrentMap; -import java.util.concurrent.ScheduledExecutorService; -import java.util.concurrent.ScheduledFuture; -import java.util.concurrent.TimeUnit; +import java.util.concurrent.CopyOnWriteArrayList; +import java.util.function.Consumer; +import java.util.function.Supplier; /** * Arrow-backed implementation of {@link NativeAllocator}. * - *

Owns a single {@link RootAllocator} for the node. All plugins that need - * Arrow buffers obtain pool handles from this class via the SPI interface. - * - *

Elastic rebalancing

- *

A background task periodically redistributes unused capacity across pools. - * Each pool has a guaranteed limit (configured via settings). When other - * pools are idle, an active pool can temporarily grow beyond its guarantee up to - * the root limit. When contention rises, pools shrink back toward their guarantee. - * This prevents idle capacity from being wasted while maintaining isolation under load. - * - *

Constructed once by {@link ArrowBasePlugin#createComponents} and exposed to - * downstream plugins via Guice and {@code PluginComponentRegistry} so consumers - * receive the instance through explicit dependency injection rather than a static - * singleton. + *

Owns a single {@link RootAllocator} (set to {@code Long.MAX_VALUE} — per-pool + * limits are the real enforcement). Manages both Arrow-backed pools and virtual pools. */ public class ArrowNativeAllocator implements NativeAllocator { + private static final Logger logger = LogManager.getLogger(ArrowNativeAllocator.class); + private final RootAllocator root; private final ConcurrentMap pools = new ConcurrentHashMap<>(); - private final ConcurrentMap poolMins = new ConcurrentHashMap<>(); - private final ConcurrentMap poolMaxes = new ConcurrentHashMap<>(); - private final ScheduledExecutorService rebalancer; - private volatile ScheduledFuture rebalanceTask; - /** - * True iff the rebalancer is configured to run periodically. Used by - * {@link #getOrCreatePool} to decide each pool's initial child-allocator - * limit: when rebalancing is enabled, pools start at {@code min} and grow - * via the next rebalance tick (preserving the original PR's - * "guarantee + burst" semantics); when rebalancing is disabled, pools - * start at {@code max} so consumers can allocate immediately without - * waiting for a tick that never comes. - */ - private volatile boolean rebalancerEnabled = false; + private final ConcurrentMap virtualPools = new ConcurrentHashMap<>(); + private final ConcurrentMap poolConfigs = new ConcurrentHashMap<>(); + private final ConcurrentMap>> poolGroupLimitListeners = new ConcurrentHashMap<>(); + private final List statsRefreshers = new CopyOnWriteArrayList<>(); + private volatile Supplier nativeMemoryStatsSupplier; + private volatile long budget = Long.MAX_VALUE; /** * Creates a new allocator with a fresh RootAllocator. - * - * @param rootLimit maximum bytes for the root allocator */ - public ArrowNativeAllocator(long rootLimit) { - this.root = new RootAllocator(rootLimit); - org.opensearch.threadpool.Scheduler.SafeScheduledThreadPoolExecutor executor = - new org.opensearch.threadpool.Scheduler.SafeScheduledThreadPoolExecutor(1, r -> { - Thread t = new Thread(r, "native-allocator-rebalancer"); - t.setDaemon(true); - return t; - }); - executor.setRemoveOnCancelPolicy(true); - this.rebalancer = executor; + public ArrowNativeAllocator() { + this.root = new RootAllocator(Long.MAX_VALUE); } /** - * Schedules (or reschedules) the rebalancer at the given interval. - * A value of 0 disables rebalancing. + * Sets the total native memory budget for validation. * - * @param intervalSeconds rebalance period in seconds, or 0 to disable + * @param budget node.native_memory.limit in bytes */ - public void setRebalanceInterval(long intervalSeconds) { - ScheduledFuture existing = rebalanceTask; - if (existing != null) { - org.opensearch.common.util.concurrent.FutureUtils.cancel(existing); - rebalanceTask = null; - } - rebalancerEnabled = intervalSeconds > 0; - if (rebalancerEnabled) { - rebalanceTask = rebalancer.scheduleAtFixedRate(this::rebalance, intervalSeconds, intervalSeconds, TimeUnit.SECONDS); - } + public void setBudget(long budget) { + this.budget = budget; } - @Override - public PoolHandle getOrCreatePool(String poolName, long limit) { - return getOrCreatePool(poolName, limit, limit); - } + // ─── Public / SPI methods ─────────────────────────────────────────────────── - /** - * Creates or returns a pool with min/max limits. - * - * @param poolName logical pool name - * @param min guaranteed minimum bytes (always available) - * @param max maximum bytes the pool can burst to - * @return the pool handle - */ - public PoolHandle getOrCreatePool(String poolName, long min, long max) { - poolMins.putIfAbsent(poolName, min); - poolMaxes.putIfAbsent(poolName, max); + @Override + public PoolHandle getOrCreatePool(String poolName, long min, long max, PoolGroup group) { + validateSumMaxesWithinBudget(poolName, max); + poolConfigs.putIfAbsent(poolName, new PoolConfig(min, max, group)); return pools.computeIfAbsent(poolName, name -> { - // Pick an initial limit that's safe for both rebalancer-on and rebalancer-off - // deployments. When rebalancing is enabled, start at min (the original PR's - // "guarantee + burst" semantics): the next rebalance tick will distribute - // headroom up to each pool's max. When rebalancing is disabled (the default), - // pools with min=0 would otherwise reject every allocation until a tick that - // never comes — start at max so consumers can allocate immediately. - long initial = rebalancerEnabled ? min : max; - BufferAllocator child = root.newChildAllocator(name, 0, initial); + BufferAllocator child = root.newChildAllocator(name, 0, max); return new ArrowPoolHandle(child); }); } @Override public void setPoolLimit(String poolName, long newLimit) { - ArrowPoolHandle handle = pools.get(poolName); - if (handle == null) { - throw new IllegalStateException("Pool '" + poolName + "' does not exist"); + PoolConfig config = poolConfigs.get(poolName); + if (config != null) { + config.max = newLimit; + } + ArrowPoolHandle arrowHandle = pools.get(poolName); + if (arrowHandle != null) { + arrowHandle.allocator.setLimit(newLimit); + return; + } + VirtualPoolHandleImpl vp = virtualPools.get(poolName); + if (vp != null) { + vp.setLimit(newLimit); + return; + } + throw new IllegalStateException("Pool '" + poolName + "' does not exist"); + } + + @Override + public VirtualPoolHandle registerVirtualPool(String poolName, long min, long max, PoolGroup group, Consumer limitSetter) { + if (min > max) { + throw new IllegalArgumentException("Pool '" + poolName + "' min (" + min + ") exceeds max (" + max + ")"); + } + validateSumMaxesWithinBudget(poolName, max); + VirtualPoolHandleImpl handle = new VirtualPoolHandleImpl(poolName, max, limitSetter); + VirtualPoolHandleImpl existing = virtualPools.putIfAbsent(poolName, handle); + if (existing != null || pools.containsKey(poolName)) { + virtualPools.remove(poolName, handle); + throw new IllegalStateException("Pool '" + poolName + "' already registered"); } - poolMaxes.put(poolName, newLimit); - handle.allocator.setLimit(newLimit); + poolConfigs.put(poolName, new PoolConfig(min, max, group)); + limitSetter.accept(max); + return handle; + } + + @Override + public void setPoolMin(String poolName, long newMin) { + PoolConfig config = poolConfigs.get(poolName); + if (config != null) { + config.min = newMin; + } + // Raise live limit if newMin exceeds current effective limit + ArrowPoolHandle arrowHandle = pools.get(poolName); + if (arrowHandle != null) { + long max = config != null ? config.max : Long.MAX_VALUE; + long current = arrowHandle.allocator.getLimit(); + long target = Math.min(newMin, max); + if (target > current) { + arrowHandle.allocator.setLimit(target); + } + return; + } + VirtualPoolHandleImpl vp = virtualPools.get(poolName); + if (vp != null) { + long max = config != null ? config.max : Long.MAX_VALUE; + long current = vp.limit(); + long target = Math.min(newMin, max); + if (target > current) { + vp.setLimit(target); + } + } + } + + @Override + public Set getAllPoolNames() { + Set all = new HashSet<>(pools.keySet()); + all.addAll(virtualPools.keySet()); + return Collections.unmodifiableSet(all); + } + + @Override + public void addStatsRefresher(Runnable refresher) { + statsRefreshers.add(refresher); + } + + @Override + public void setNativeMemoryStatsSupplier(Supplier supplier) { + this.nativeMemoryStatsSupplier = supplier; } /** - * Updates the minimum guaranteed bytes for a pool. The new min is recorded for the - * rebalancer (which honors it as a floor on the next tick) and also pushed to the - * live {@link BufferAllocator} so the change takes effect immediately even when - * the rebalancer is disabled — the alternative was a Dynamic setting that returned - * HTTP 200 but had no observable effect. + * Sets the effective (live) limit for a pool without updating the configured max. + * Used by the rebalancer to adjust pool limits dynamically. * - *

Live propagation rules: - *

    - *
  • If {@code newMin} exceeds the pool's current limit, the limit is raised to - * {@code newMin} (capped at the configured pool max). Children of the pool - * allocator inherit the change automatically via Arrow's parent-cap check at - * allocation time, so dynamic resizes reach in-flight workloads without an - * explicit notification SPI. - *
  • If {@code newMin} is below the current limit, the limit is left alone — - * the rebalancer is the only path that shrinks live limits, so a min change - * on its own never reduces capacity in flight. - *
- * - * @param poolName the pool name - * @param newMin new minimum bytes + * @param poolName name of the pool + * @param newLimit new effective limit in bytes */ - public void setPoolMin(String poolName, long newMin) { - ArrowPoolHandle handle = pools.get(poolName); - if (handle == null) { - throw new IllegalStateException("Pool '" + poolName + "' does not exist"); + public void setPoolEffectiveLimit(String poolName, long newLimit) { + ArrowPoolHandle arrowHandle = pools.get(poolName); + if (arrowHandle != null) { + arrowHandle.allocator.setLimit(newLimit); + return; } - poolMins.put(poolName, newMin); - long max = poolMaxes.getOrDefault(poolName, Long.MAX_VALUE); - long current = handle.allocator.getLimit(); - long target = Math.min(newMin, max); - if (target > current) { - handle.allocator.setLimit(target); + VirtualPoolHandleImpl vp = virtualPools.get(poolName); + if (vp != null) { + vp.setLimit(newLimit); + return; } + throw new IllegalStateException("Pool '" + poolName + "' does not exist"); } - @Override - public void setRootLimit(long limit) { - root.setLimit(limit); + /** + * Resets all pools to their configured max. Called when the rebalancer is disabled. + * Logs a warning for any pool that was bursting above its max. + */ + public void resetAllPoolsToMax() { + for (String name : getAllPoolNames()) { + PoolConfig config = poolConfigs.get(name); + long max = config != null ? config.max : Long.MAX_VALUE; + long current = getEffectiveLimit(name); + if (current > max) { + logger.warn( + "Pool [{}] effective limit {} exceeds max {}, resetting to max. In-flight allocations may be rejected.", + name, + current, + max + ); + } + setPoolEffectiveLimit(name, max); + } } /** - * Returns a point-in-time stats snapshot across all pools. Used by the - * {@code NativeAllocatorStatsRegistry} component published from - * {@code ArrowBasePlugin.createComponents()} and wired into {@code NodeService} to - * render allocator state under {@code _nodes/stats[/native_allocator]}. + * Convenience method for plugins that have Setting objects. Registers the virtual pool + * and auto-wires dynamic setting listeners for min/max changes. + * + * @param poolName name of the virtual pool + * @param minSetting setting for minimum bytes + * @param maxSetting setting for maximum bytes + * @param settings current node settings + * @param clusterSettings cluster settings for dynamic updates + * @param group pool group assignment + * @param limitSetter callback invoked when the pool limit changes + */ + public VirtualPoolHandle registerVirtualPool( + String poolName, + Setting minSetting, + Setting maxSetting, + Settings settings, + ClusterSettings clusterSettings, + PoolGroup group, + Consumer limitSetter + ) { + long min = minSetting.get(settings); + long max = maxSetting.get(settings); + VirtualPoolHandle handle = registerVirtualPool(poolName, min, max, group, limitSetter); + + clusterSettings.addSettingsUpdateConsumer(maxSetting, newMax -> setPoolLimit(poolName, newMax)); + clusterSettings.addSettingsUpdateConsumer(minSetting, newMin -> setPoolMin(poolName, newMin)); + + return handle; + } + + /** + * Returns a point-in-time stats snapshot across all pools. */ public NativeAllocatorPoolStats stats() { + refreshStats(); + + long nativeAllocated = -1; + long nativeResident = -1; + Supplier supplier = this.nativeMemoryStatsSupplier; + if (supplier != null) { + try { + long[] stats = supplier.get(); + if (stats != null && stats.length >= 2) { + nativeAllocated = stats[0]; + nativeResident = stats[1]; + } + } catch (Exception e) { + // best-effort + } + } + List poolStats = new ArrayList<>(); for (var entry : pools.entrySet()) { BufferAllocator alloc = entry.getValue().allocator; + PoolConfig config = poolConfigs.get(entry.getKey()); poolStats.add( new NativeAllocatorPoolStats.PoolStats( entry.getKey(), alloc.getAllocatedMemory(), alloc.getPeakMemoryAllocation(), - alloc.getLimit() + alloc.getLimit(), + config != null && config.group != null ? config.group.getName() : null, + config != null ? config.min : 0L + ) + ); + } + for (var entry : virtualPools.entrySet()) { + VirtualPoolHandleImpl vp = entry.getValue(); + PoolConfig config = poolConfigs.get(entry.getKey()); + poolStats.add( + new NativeAllocatorPoolStats.PoolStats( + entry.getKey(), + vp.allocatedBytes(), + vp.peakBytes(), + vp.limit(), + config != null && config.group != null ? config.group.getName() : null, + config != null ? config.min : 0L ) ); } - return new NativeAllocatorPoolStats(root.getAllocatedMemory(), root.getPeakMemoryAllocation(), root.getLimit(), poolStats); + + return new NativeAllocatorPoolStats(nativeAllocated, nativeResident, poolStats); + } + + /** + * Runs all registered stats refreshers. + */ + public void refreshStats() { + for (Runnable refresher : statsRefreshers) { + try { + refresher.run(); + } catch (Exception e) { + // best-effort + } + } } @Override public void close() { - rebalancer.shutdownNow(); pools.forEach((name, handle) -> { try { handle.allocator.close(); @@ -209,72 +308,23 @@ public void close() { } }); pools.clear(); - // Close any remaining child allocators (e.g., ad-hoc children created via ArrowAllocatorService) + virtualPools.clear(); for (BufferAllocator child : new ArrayList<>(root.getChildAllocators())) { try { child.close(); } catch (Exception e) { - // best-effort — log but don't block shutdown + // best-effort } } root.close(); } - /** - * Redistributes unused capacity across pools based on min/max guarantees. - * - *

Algorithm: - *

    - *
  1. Every pool is guaranteed at least its configured min
    2. - *
  2. Compute headroom = rootLimit - sum(all pool current allocations)
    3. - *
  3. Distribute headroom equally across all pools (not just active ones), capped - * at each pool's max. Distributing to all pools — including those with zero - * current allocation — avoids the dead-pool corner case where a pool with - * min = 0 starts at limit = 0, can never make its first allocation, and so - * never becomes "active" enough to receive a bonus. Pools that don't need the - * headroom stay at min naturally because their max caps the bonus.
    4. - *
  4. No pool's limit ever drops below its current allocation or its min
    5. - *
- */ - void rebalance() { - if (pools.isEmpty()) return; - - long rootLimit = root.getLimit(); - long totalAllocated = 0; - - for (Map.Entry entry : pools.entrySet()) { - totalAllocated += entry.getValue().allocator.getAllocatedMemory(); - } - - long headroom = Math.max(0, rootLimit - totalAllocated); - int poolCount = pools.size(); - long bonusPerPool = poolCount > 0 ? headroom / poolCount : 0; - - for (Map.Entry entry : pools.entrySet()) { - String name = entry.getKey(); - BufferAllocator alloc = entry.getValue().allocator; - long min = poolMins.getOrDefault(name, 0L); - long max = poolMaxes.getOrDefault(name, Long.MAX_VALUE); - long currentAllocation = alloc.getAllocatedMemory(); - - long effectiveLimit = min + bonusPerPool; - - // Cap at pool's max - effectiveLimit = Math.min(effectiveLimit, max); - // Never drop below current allocation or min - effectiveLimit = Math.max(effectiveLimit, currentAllocation); - effectiveLimit = Math.max(effectiveLimit, min); - // Never exceed root - effectiveLimit = Math.min(effectiveLimit, rootLimit); - - alloc.setLimit(effectiveLimit); - } - } + // ─── Package-private accessors (used by rebalancer and tests) ──────────────── /** * Returns the underlying Arrow allocator for a pool. * - * @param poolName name of the pool to look up + * @param poolName name of the pool */ public BufferAllocator getPoolAllocator(String poolName) { ArrowPoolHandle handle = pools.get(poolName); @@ -284,16 +334,12 @@ public BufferAllocator getPoolAllocator(String poolName) { return handle.allocator; } - /** - * Returns the root Arrow allocator. - */ + /** Returns the root Arrow allocator. */ public BufferAllocator getRootAllocator() { return root; } - /** - * Returns all registered pool names. - */ + /** Returns all registered pool names (Arrow pools only). */ public Set getPoolNames() { return Collections.unmodifiableSet(pools.keySet()); } @@ -304,7 +350,8 @@ public Set getPoolNames() { * @param poolName name of the pool */ public long getPoolMin(String poolName) { - return poolMins.getOrDefault(poolName, 0L); + PoolConfig config = poolConfigs.get(poolName); + return config != null ? config.min : 0L; } /** @@ -313,7 +360,206 @@ public long getPoolMin(String poolName) { * @param poolName name of the pool */ public long getPoolMax(String poolName) { - return poolMaxes.getOrDefault(poolName, Long.MAX_VALUE); + PoolConfig config = poolConfigs.get(poolName); + return config != null ? config.max : Long.MAX_VALUE; + } + + /** + * Returns the group for a pool, or null if not assigned. + * + * @param poolName name of the pool + */ + public PoolGroup getPoolGroup(String poolName) { + PoolConfig config = poolConfigs.get(poolName); + return config != null ? config.group : null; + } + + /** + * Returns the allocated bytes for a virtual pool. + * + * @param poolName name of the virtual pool + */ + public long getVirtualPoolAllocated(String poolName) { + VirtualPoolHandleImpl vp = virtualPools.get(poolName); + return vp != null ? vp.allocatedBytes() : 0; + } + + /** + * Returns the current limit for a virtual pool. + * + * @param poolName name of the virtual pool + */ + public long getVirtualPoolLimit(String poolName) { + VirtualPoolHandleImpl vp = virtualPools.get(poolName); + return vp != null ? vp.limit() : 0; + } + + /** + * Returns the effective limit for any pool (Arrow or virtual). + * + * @param poolName name of the pool + */ + public long getEffectiveLimit(String poolName) { + ArrowPoolHandle arrowHandle = pools.get(poolName); + if (arrowHandle != null) { + return arrowHandle.allocator.getLimit(); + } + VirtualPoolHandleImpl vp = virtualPools.get(poolName); + if (vp != null) { + return vp.limit(); + } + return 0; + } + + /** + * Returns the allocated bytes for any pool (Arrow or virtual). + * + * @param poolName name of the pool + */ + public long getAllocated(String poolName) { + ArrowPoolHandle arrowHandle = pools.get(poolName); + if (arrowHandle != null) { + return arrowHandle.allocator.getAllocatedMemory(); + } + VirtualPoolHandleImpl vp = virtualPools.get(poolName); + if (vp != null) { + return vp.allocatedBytes(); + } + return 0; + } + + /** Returns the native memory stats supplier. */ + public Supplier getNativeMemoryStatsSupplier() { + return nativeMemoryStatsSupplier; + } + + /** + * Registers a listener invoked when the grouped effective limit for the given pool group changes. + * The consumer receives the new total effective limit (sum of all pools in the group). + * + * @param group the pool group to listen on + * @param listener consumer that receives the new grouped limit in bytes + */ + public void addPoolGroupLimitListener(PoolGroup group, Consumer listener) { + poolGroupLimitListeners.computeIfAbsent(group, k -> new CopyOnWriteArrayList<>()).add(listener); + } + + /** + * Fires listeners for the specified pool group with the current grouped effective limit sum. + * Called by the rebalancer after all pool limits for a tick have been updated. + * + * @param group the pool group whose listeners should be notified + */ + public void firePoolGroupListeners(PoolGroup group) { + List> listeners = poolGroupLimitListeners.get(group); + if (listeners == null || listeners.isEmpty()) return; + + long groupSum = 0; + for (var entry : poolConfigs.entrySet()) { + if (entry.getValue().group == group) { + groupSum += getEffectiveLimit(entry.getKey()); + } + } + long finalSum = groupSum; + for (Consumer listener : listeners) { + try { + listener.accept(finalSum); + } catch (Exception e) { + logger.warn( + () -> new org.apache.logging.log4j.message.ParameterizedMessage( + "Pool group limit listener failed for group [{}]", + group + ), + e + ); + } + } + } + + // ─── Private helpers ───────────────────────────────────────────────────────── + + private void validateSumMaxesWithinBudget(String newPoolName, long newPoolMax) { + if (budget == Long.MAX_VALUE || budget <= 0) { + return; + } + long sumMaxes = newPoolMax; + for (var entry : poolConfigs.entrySet()) { + if (entry.getKey().equals(newPoolName) == false) { + sumMaxes += entry.getValue().max; + } + } + if (sumMaxes > budget) { + throw new IllegalArgumentException( + "Sum of pool max limits (" + + sumMaxes + + " bytes) exceeds native memory budget (" + + budget + + " bytes). Reduce pool max settings or increase the budget." + ); + } + } + + // ─── Inner classes ─────────────────────────────────────────────────────────── + + /** + * Mutable configuration for a pool: min, max, and group. + */ + static class PoolConfig { + volatile long min; + volatile long max; + final PoolGroup group; + + PoolConfig(long min, long max, PoolGroup group) { + this.min = min; + this.max = max; + this.group = group; + } + } + + /** + * Virtual pool handle implementation. Tracks stats reported from native layer + * and delegates limit changes to the registered callback. + */ + public static class VirtualPoolHandleImpl implements VirtualPoolHandle { + private final String name; + private volatile long limit; + private volatile long allocatedBytes; + private volatile long peakBytes; + private final Consumer limitSetter; + + VirtualPoolHandleImpl(String name, long limit, Consumer limitSetter) { + this.name = name; + this.limit = limit; + this.limitSetter = limitSetter; + } + + @Override + public void updateStats(long allocated, long peak) { + this.allocatedBytes = allocated; + this.peakBytes = peak; + } + + void setLimit(long newLimit) { + this.limit = newLimit; + if (limitSetter != null) { + limitSetter.accept(newLimit); + } + } + + @Override + public long allocatedBytes() { + return allocatedBytes; + } + + @Override + public long peakBytes() { + return peakBytes; + } + + @Override + public long limit() { + return limit; + } } /** diff --git a/plugins/arrow-base/src/main/java/org/opensearch/arrow/allocator/NativeMemoryRebalancer.java b/plugins/arrow-base/src/main/java/org/opensearch/arrow/allocator/NativeMemoryRebalancer.java new file mode 100644 index 0000000000000..5a5294b2cb4b0 --- /dev/null +++ b/plugins/arrow-base/src/main/java/org/opensearch/arrow/allocator/NativeMemoryRebalancer.java @@ -0,0 +1,241 @@ +/* + * SPDX-License-Identifier: Apache-2.0 + * + * The OpenSearch Contributors require contributions made to + * this file be licensed under the Apache-2.0 license or a + * compatible open source license. + */ + +package org.opensearch.arrow.allocator; + +import org.apache.logging.log4j.LogManager; +import org.apache.logging.log4j.Logger; +import org.apache.logging.log4j.message.ParameterizedMessage; +import org.opensearch.arrow.spi.PoolGroup; + +import java.util.HashMap; +import java.util.Map; +import java.util.Set; +import java.util.function.Supplier; + +/** + * Periodic rebalancer that redistributes native memory across pools. + * + *

Algorithm: + *

    + *
  • Pools start at their configured max on registration
    • + *
  • If no pool is under pressure, return early (no-op)
    • + *
  • Idle pools (utilization < idle_threshold) are shrunk, never below min
    • + *
  • Pressured pools (utilization > pressure_threshold) receive freed capacity, can exceed max
    • + *
  • Excess freed capacity is returned to idle pools proportionally
    • + *
  • Invariant: sum(effective_limits) <= budget at all times
    • + *
+ * + * @opensearch.internal + */ +public class NativeMemoryRebalancer implements Runnable { + + private static final Logger logger = LogManager.getLogger(NativeMemoryRebalancer.class); + + private final ArrowNativeAllocator allocator; + private final Supplier budgetSupplier; + + private volatile double pressureThreshold; + private volatile double idleThreshold; + private volatile double shrinkFactor; + + /** + * Creates a new rebalancer. + * + * @param allocator the allocator managing all pools + * @param budgetSupplier supplies the current budget value + * @param pressureThreshold utilization above this triggers growth (default 0.75) + * @param idleThreshold utilization below this means pool can give back capacity (default 0.50) + * @param shrinkFactor factor to shrink idle pools by — new limit = limit * (1 - shrinkFactor) (default 0.10) + */ + public NativeMemoryRebalancer( + ArrowNativeAllocator allocator, + Supplier budgetSupplier, + double pressureThreshold, + double idleThreshold, + double shrinkFactor + ) { + this.allocator = allocator; + this.budgetSupplier = budgetSupplier; + this.pressureThreshold = pressureThreshold; + this.idleThreshold = idleThreshold; + this.shrinkFactor = shrinkFactor; + } + + /** + * Updates the pressure threshold dynamically. + * + * @param value new threshold (0.0 to 1.0) + */ + public void setPressureThreshold(double value) { + this.pressureThreshold = value; + } + + /** + * Updates the idle threshold dynamically. + * + * @param value new threshold (0.0 to 1.0) + */ + public void setIdleThreshold(double value) { + this.idleThreshold = value; + } + + /** + * Updates the shrink factor dynamically. + * + * @param value new factor (0.0 to 1.0) + */ + public void setShrinkFactor(double value) { + this.shrinkFactor = value; + } + + @Override + public void run() { + try { + rebalance(); + } catch (Exception e) { + logger.warn("Rebalancer tick failed", e); + } + } + + void rebalance() { + Set allPools = allocator.getAllPoolNames(); + if (allPools.isEmpty()) return; + + long budget = budgetSupplier.get(); + if (budget <= 0 || budget == Long.MAX_VALUE) return; + + // Refresh stats from native layers + allocator.refreshStats(); + + // Snapshot per-pool state + Map snapshots = new HashMap<>(); + for (String name : allPools) { + long allocated = allocator.getAllocated(name); + long effectiveLimit = allocator.getEffectiveLimit(name); + long min = allocator.getPoolMin(name); + long max = allocator.getPoolMax(name); + double utilization = effectiveLimit > 0 ? (double) allocated / effectiveLimit : 0; + snapshots.put(name, new PoolSnapshot(allocated, effectiveLimit, min, max, utilization)); + } + + // Identify pressured pools — if none, nothing to do + Map desires = new HashMap<>(); + long totalDesired = 0; + for (var entry : snapshots.entrySet()) { + PoolSnapshot s = entry.getValue(); + if (s.utilization > pressureThreshold) { + long desired = Math.max(1, (long) (s.allocated * 0.25)); + desires.put(entry.getKey(), desired); + totalDesired += desired; + } + } + if (totalDesired == 0) { + logger.debug("Rebalancer: no pools under pressure, skipping"); + return; + } + + // Shrink idle pools, floor at min + long freedCapacity = 0; + for (var entry : snapshots.entrySet()) { + PoolSnapshot s = entry.getValue(); + if (s.utilization < idleThreshold) { + long newLimit = Math.max((long) (s.effectiveLimit * (1.0 - shrinkFactor)), s.min); + newLimit = Math.max(newLimit, s.allocated); + if (newLimit < s.effectiveLimit) { + freedCapacity += s.effectiveLimit - newLimit; + allocator.setPoolEffectiveLimit(entry.getKey(), newLimit); + s.effectiveLimit = newLimit; + } + } + } + + if (freedCapacity == 0) { + logger.debug("Rebalancer: no capacity freed from idle pools"); + return; + } + + // Distribute freed capacity to pressured pools (can exceed max) + long totalGranted = 0; + long grantCap = Math.min(freedCapacity, totalDesired); + for (var entry : desires.entrySet()) { + String name = entry.getKey(); + long desired = entry.getValue(); + PoolSnapshot s = snapshots.get(name); + long grant = (long) ((double) grantCap * desired / totalDesired); + grant = Math.min(grant, grantCap - totalGranted); + if (grant > 0) { + try { + long newLimit = s.effectiveLimit + grant; + allocator.setPoolEffectiveLimit(name, newLimit); + totalGranted += grant; + logger.debug("Rebalancer: grew pool [{}] by {} bytes to {} (max={})", name, grant, newLimit, s.max); + } catch (Exception e) { + logger.warn(() -> new ParameterizedMessage("Rebalancer: failed to grow pool [{}]", name), e); + } + } + } + + // Return any excess freed capacity back to idle pools + long excess = freedCapacity - totalGranted; + if (excess > 0) { + returnToIdlePools(snapshots, excess); + } + + // Notify pool group listeners after all limit changes are complete + for (PoolGroup group : PoolGroup.values()) { + allocator.firePoolGroupListeners(group); + } + } + + // ─── Private helpers ───────────────────────────────────────────────────────── + + private void returnToIdlePools(Map snapshots, long capacity) { + long totalIdleSize = 0; + for (PoolSnapshot s : snapshots.values()) { + if (s.utilization < idleThreshold) { + totalIdleSize += s.effectiveLimit; + } + } + if (totalIdleSize == 0) return; + + long totalReturned = 0; + for (var entry : snapshots.entrySet()) { + PoolSnapshot s = entry.getValue(); + if (s.utilization < idleThreshold) { + long share = (long) ((double) capacity * s.effectiveLimit / totalIdleSize); + share = Math.min(share, capacity - totalReturned); + if (share > 0) { + long newLimit = s.effectiveLimit + share; + allocator.setPoolEffectiveLimit(entry.getKey(), newLimit); + s.effectiveLimit = newLimit; + totalReturned += share; + } + } + } + } + + /** + * Point-in-time snapshot of a pool's state during one rebalance tick. + */ + static class PoolSnapshot { + final long allocated; + long effectiveLimit; + final long min; + final long max; + final double utilization; + + PoolSnapshot(long allocated, long effectiveLimit, long min, long max, double utilization) { + this.allocated = allocated; + this.effectiveLimit = effectiveLimit; + this.min = min; + this.max = max; + this.utilization = utilization; + } + } +} diff --git a/plugins/arrow-base/src/test/java/org/opensearch/arrow/allocator/ArrowBasePluginTests.java b/plugins/arrow-base/src/test/java/org/opensearch/arrow/allocator/ArrowBasePluginTests.java index ad72b70b8cbbb..df2673a5a8a7b 100644 --- a/plugins/arrow-base/src/test/java/org/opensearch/arrow/allocator/ArrowBasePluginTests.java +++ b/plugins/arrow-base/src/test/java/org/opensearch/arrow/allocator/ArrowBasePluginTests.java @@ -8,12 +8,11 @@ package org.opensearch.arrow.allocator; -import org.apache.arrow.memory.BufferAllocator; -import org.apache.arrow.memory.OutOfMemoryException; import org.opensearch.arrow.spi.NativeAllocatorPoolConfig; import org.opensearch.common.settings.ClusterSettings; import org.opensearch.common.settings.Setting; import org.opensearch.common.settings.Settings; +import org.opensearch.node.resource.tracker.ResourceTrackerSettings; import org.opensearch.test.OpenSearchTestCase; import java.util.HashSet; @@ -21,90 +20,23 @@ public class ArrowBasePluginTests extends OpenSearchTestCase { - public void testDeriveRootLimitDefaultUnsetReturnsLongMaxValue() { - // Explicit 0 expresses "AC unconfigured" — default is now ram - heap, so Settings.EMPTY - // would resolve to a real value on whatever machine the test runs on. - Settings s = Settings.builder().put("node.native_memory.limit", "0b").build(); - assertEquals(Long.toString(Long.MAX_VALUE), ArrowBasePlugin.deriveRootLimitDefault(s)); - } - - public void testDeriveRootLimitDefaultUsesAcLimitWhenSet() { - Settings s = Settings.builder().put("node.native_memory.limit", "1gb").build(); - // ROOT_LIMIT defaults to 20% of node.native_memory.limit — the Arrow framework gets a - // small fraction of native budget; DataFusion's Rust runtime takes the larger share. - long oneGiB = 1024L * 1024 * 1024; - assertEquals(Long.toString(oneGiB * 20 / 100), ArrowBasePlugin.deriveRootLimitDefault(s)); - } - - public void testDeriveRootLimitDefaultIgnoresBufferPercent() { - // node.native_memory.buffer_percent is admission control's throttle margin, not a - // framework budget reduction. The framework default takes its 20% fraction off - // node.native_memory.limit directly so AC's safety margin sits between AC's throttle - // threshold and the framework's hard cap rather than being collapsed into the cap. - // 1000 bytes limit, 20% buffer => root.limit still 20% of 1000 = 200. - Settings s = Settings.builder().put("node.native_memory.limit", "1000b").put("node.native_memory.buffer_percent", 20).build(); - assertEquals("200", ArrowBasePlugin.deriveRootLimitDefault(s)); - } - - public void testRootLimitSettingExposesDerivedDefault() { - Settings s = Settings.builder().put("node.native_memory.limit", "10gb").build(); - // 20% of 10 GiB. - long expected = 10L * 1024 * 1024 * 1024 * 20 / 100; - assertEquals(Long.valueOf(expected), ArrowBasePlugin.ROOT_LIMIT_SETTING.get(s)); - } - - public void testRootLimitSettingExplicitOverridesDerived() { - Settings s = Settings.builder() - .put("node.native_memory.limit", "8gb") - .put(NativeAllocatorPoolConfig.SETTING_ROOT_LIMIT, 1024L) - .build(); - assertEquals(Long.valueOf(1024L), ArrowBasePlugin.ROOT_LIMIT_SETTING.get(s)); - } - - public void testRootLimitRejectsNegative() { - Settings s = Settings.builder().put(NativeAllocatorPoolConfig.SETTING_ROOT_LIMIT, -1L).build(); - IllegalArgumentException e = expectThrows(IllegalArgumentException.class, () -> ArrowBasePlugin.ROOT_LIMIT_SETTING.get(s)); - assertTrue(e.getMessage().contains("must be >= 0")); - } - public void testQuerySettingsExposeDefaults() { // Explicit 0 expresses "AC unconfigured" so QUERY_MAX falls back to Long.MAX_VALUE. - // Settings.EMPTY would resolve via ram - heap default to a finite, machine-dependent value. Settings s = Settings.builder().put("node.native_memory.limit", "0b").build(); assertEquals(Long.valueOf(0L), ArrowBasePlugin.QUERY_MIN_SETTING.get(s)); assertEquals(Long.valueOf(Long.MAX_VALUE), ArrowBasePlugin.QUERY_MAX_SETTING.get(s)); } - public void testFlightAndIngestMinDefaultsToZero() { - // The grouped validator (validateMinSum) treats per-pool mins as a guarantee - // floor — defaults of Long.MAX_VALUE caused the validator to reject any PUT - // that set a non-MAX root. Pool mins must default to zero so the baseline - // configuration is consistent. - Settings s = Settings.EMPTY; - assertEquals(Long.valueOf(0L), ArrowBasePlugin.FLIGHT_MIN_SETTING.get(s)); - assertEquals(Long.valueOf(0L), ArrowBasePlugin.INGEST_MIN_SETTING.get(s)); - } - - public void testQuerySettingsAcceptValues() { - Settings s = Settings.builder() - .put(NativeAllocatorPoolConfig.SETTING_QUERY_MIN, 100L) - .put(NativeAllocatorPoolConfig.SETTING_QUERY_MAX, 1000L) - .build(); - assertEquals(Long.valueOf(100L), ArrowBasePlugin.QUERY_MIN_SETTING.get(s)); - assertEquals(Long.valueOf(1000L), ArrowBasePlugin.QUERY_MAX_SETTING.get(s)); + public void testFlightAndIngestMinDerivedFromBudget() { + // With node.native_memory.limit set, mins derive as percentages + Settings s = Settings.builder().put("node.native_memory.limit", "1gb").build(); + long budget = 1024L * 1024 * 1024; + // flight min = 2% of budget, ingest min = 4% of budget + assertEquals(Long.valueOf(budget * 2 / 100), ArrowBasePlugin.FLIGHT_MIN_SETTING.get(s)); + assertEquals(Long.valueOf(budget * 4 / 100), ArrowBasePlugin.INGEST_MIN_SETTING.get(s)); } - // -- Pool max defaults derived from node.native_memory.limit ---------- - // Pool maxes anchor to the operator's off-heap budget (node.native_memory.limit), - // not to native.allocator.root.limit. This matches the PR #21732 partitioning - // diagram where pool fractions (5%/8%/5%) are of native_memory.limit. Sum of - // pool maxes (18% of native_memory.limit) fits within root.limit (20% of - // native_memory.limit) by default, leaving 2 pp headroom inside the root cap. - public void testPoolMaxDefaultsAreLongMaxValueWhenAcUnset() { - // AC explicitly unconfigured — pool maxes default to Long.MAX_VALUE (unbounded), - // preserving pre-AC behaviour. The default for node.native_memory.limit is - // 79% of (ram - heap), so to test the "unset" branch we must explicitly set it to 0. Settings s = Settings.builder().put("node.native_memory.limit", "0b").build(); assertEquals(Long.valueOf(Long.MAX_VALUE), ArrowBasePlugin.FLIGHT_MAX_SETTING.get(s)); assertEquals(Long.valueOf(Long.MAX_VALUE), ArrowBasePlugin.INGEST_MAX_SETTING.get(s)); @@ -112,10 +44,6 @@ public void testPoolMaxDefaultsAreLongMaxValueWhenAcUnset() { } public void testPoolMaxDefaultsScaleFromAcBudget() { - // 10 GiB native memory limit. Pool maxes per the partitioning model in PR #21732: - // FLIGHT_MAX = 5% INGEST_MAX = 8% QUERY_MAX = 5% - // Anchored to node.native_memory.limit, not to root.limit (which defaults to 20% - // of native_memory.limit) — see derivePoolMaxDefault Javadoc. Settings s = Settings.builder().put("node.native_memory.limit", "10gb").build(); long limit = 10L * 1024 * 1024 * 1024; assertEquals(Long.valueOf(limit * 5 / 100), ArrowBasePlugin.FLIGHT_MAX_SETTING.get(s)); @@ -123,51 +51,14 @@ public void testPoolMaxDefaultsScaleFromAcBudget() { assertEquals(Long.valueOf(limit * 5 / 100), ArrowBasePlugin.QUERY_MAX_SETTING.get(s)); } - public void testPoolMaxDefaultsIgnoreRootLimitOverride() { - // Pool maxes anchor to node.native_memory.limit, not to root.limit. An operator - // who overrides root.limit (e.g. to 4 GiB instead of the default 20% of - // native_memory.limit = 2 GiB) does not shrink pool defaults proportionally; - // the diagrammed partitioning of native_memory.limit holds. - Settings s = Settings.builder() - .put("node.native_memory.limit", "10gb") - .put(NativeAllocatorPoolConfig.SETTING_ROOT_LIMIT, 4L * 1024 * 1024 * 1024) - .build(); - long limit = 10L * 1024 * 1024 * 1024; - assertEquals(Long.valueOf(limit * 5 / 100), ArrowBasePlugin.FLIGHT_MAX_SETTING.get(s)); - assertEquals(Long.valueOf(limit * 8 / 100), ArrowBasePlugin.INGEST_MAX_SETTING.get(s)); - assertEquals(Long.valueOf(limit * 5 / 100), ArrowBasePlugin.QUERY_MAX_SETTING.get(s)); - } - public void testPoolMaxDefaultsIgnoreBufferPercent() { - // node.native_memory.buffer_percent is AC's throttle margin, not a framework budget - // reduction. Pool maxes default off node.native_memory.limit directly so AC's safety - // margin sits between AC's throttle threshold and the framework's hard cap rather than - // being collapsed into the cap. - // 1000 bytes limit, 20% buffer => pool maxes are still 5/8/5% of 1000 = 50/80/50. Settings s = Settings.builder().put("node.native_memory.limit", "1000b").put("node.native_memory.buffer_percent", 20).build(); assertEquals(Long.valueOf(50L), ArrowBasePlugin.FLIGHT_MAX_SETTING.get(s)); assertEquals(Long.valueOf(80L), ArrowBasePlugin.INGEST_MAX_SETTING.get(s)); assertEquals(Long.valueOf(50L), ArrowBasePlugin.QUERY_MAX_SETTING.get(s)); } - public void testPoolMaxExplicitOverridesDerived() { - // Operator-set values must win over derived defaults. - Settings s = Settings.builder() - .put("node.native_memory.limit", "10gb") - .put(NativeAllocatorPoolConfig.SETTING_FLIGHT_MAX, 7L) - .put(NativeAllocatorPoolConfig.SETTING_INGEST_MAX, 8L) - .put(NativeAllocatorPoolConfig.SETTING_QUERY_MAX, 9L) - .build(); - assertEquals(Long.valueOf(7L), ArrowBasePlugin.FLIGHT_MAX_SETTING.get(s)); - assertEquals(Long.valueOf(8L), ArrowBasePlugin.INGEST_MAX_SETTING.get(s)); - assertEquals(Long.valueOf(9L), ArrowBasePlugin.QUERY_MAX_SETTING.get(s)); - } - public void testPoolMaxRejectsNegative() { - // Negative pool max is rejected at parse time, mirroring ROOT_LIMIT_SETTING. - // Each pool's parser has its own message so we exercise all three to lock down - // the per-pool error contract (and keep coverage honest on what is otherwise - // boilerplate-but-distinct branches). Settings flight = Settings.builder().put(NativeAllocatorPoolConfig.SETTING_FLIGHT_MAX, -1L).build(); IllegalArgumentException e = expectThrows(IllegalArgumentException.class, () -> ArrowBasePlugin.FLIGHT_MAX_SETTING.get(flight)); assertTrue(e.getMessage().contains("must be >= 0")); @@ -188,253 +79,75 @@ public void testPoolMaxRejectsNegative() { } // ----------------------------------------------------------------- - // End-to-end wiring tests — verify that Setting.Property.Dynamic settings - // actually flow through to the live allocator. These guard against the - // "dynamic in name only" failure mode where a setting parses, the validator - // runs, the cluster-state update succeeds, and the runtime component - // silently does nothing because the addSettingsUpdateConsumer line was - // never registered. Bare-setter unit tests do not catch this; tests must - // drive a real ClusterSettings#applySettings round-trip. + // End-to-end wiring tests // ----------------------------------------------------------------- - /** - * Builds a {@link ClusterSettings} preloaded with all of {@link ArrowBasePlugin}'s - * settings, mirroring what {@code SettingsModule} does at node startup. Returns the - * fresh allocator with the framework's pools created and consumers registered - * — the same wiring path {@code createComponents} runs. - */ - private static ArrowNativeAllocator newWiredAllocator(Settings nodeSettings, ClusterSettings cs) { - long rootLimit = ArrowBasePlugin.ROOT_LIMIT_SETTING.get(nodeSettings); - ArrowNativeAllocator allocator = new ArrowNativeAllocator(rootLimit); - allocator.setRebalanceInterval(ArrowBasePlugin.REBALANCE_INTERVAL_SETTING.get(nodeSettings)); - allocator.getOrCreatePool( - NativeAllocatorPoolConfig.POOL_FLIGHT, - ArrowBasePlugin.FLIGHT_MIN_SETTING.get(nodeSettings), - ArrowBasePlugin.FLIGHT_MAX_SETTING.get(nodeSettings) - ); - allocator.getOrCreatePool( - NativeAllocatorPoolConfig.POOL_INGEST, - ArrowBasePlugin.INGEST_MIN_SETTING.get(nodeSettings), - ArrowBasePlugin.INGEST_MAX_SETTING.get(nodeSettings) - ); - allocator.getOrCreatePool( - NativeAllocatorPoolConfig.POOL_QUERY, - ArrowBasePlugin.QUERY_MIN_SETTING.get(nodeSettings), - ArrowBasePlugin.QUERY_MAX_SETTING.get(nodeSettings) - ); - ArrowBasePlugin.registerSettingsUpdateConsumers(cs, allocator); - return allocator; - } - private static ClusterSettings newClusterSettings(Settings nodeSettings) { Set> registered = new HashSet<>(); registered.addAll(new ArrowBasePlugin().getSettings()); return new ClusterSettings(nodeSettings, registered); } - public void testBuildAllocatorWiresAllPoolsAndSettingsConsumers() { - // Verifies the full createComponents code path — the helper extracted from - // createComponents builds the allocator, creates all three pools (FLIGHT, INGEST, - // QUERY), and registers the cluster-settings update consumers. We bypass the - // heavyweight ClusterService fixture and inject a real ClusterSettings directly, - // which is what production wiring also passes through to buildAllocator after - // unpacking the createComponents arguments. + public void testBuildAllocatorWiresAllPools() throws Exception { Settings nodeSettings = Settings.builder() - .put(NativeAllocatorPoolConfig.SETTING_ROOT_LIMIT, 8L * 1024 * 1024 * 1024) + .put("node.native_memory.limit", "10gb") .put(NativeAllocatorPoolConfig.SETTING_FLIGHT_MAX, 1L * 1024 * 1024 * 1024) .put(NativeAllocatorPoolConfig.SETTING_INGEST_MAX, 2L * 1024 * 1024 * 1024) .put(NativeAllocatorPoolConfig.SETTING_QUERY_MAX, 1L * 1024 * 1024 * 1024) + .put("native.allocator.rebalancer.enabled", false) .build(); ClusterSettings cs = newClusterSettings(nodeSettings); - ArrowNativeAllocator allocator = ArrowBasePlugin.buildAllocator(nodeSettings, cs); + ArrowBasePlugin plugin = new ArrowBasePlugin(); + long budget = ResourceTrackerSettings.NODE_NATIVE_MEMORY_LIMIT_SETTING.get(nodeSettings).getBytes(); + ArrowNativeAllocator allocator = plugin.buildAllocator(nodeSettings, cs, () -> budget); try { - // All three pools created. Set poolNames = allocator.getPoolNames(); assertEquals("buildAllocator must register exactly the framework's three pools", 3, poolNames.size()); assertTrue(poolNames.contains(NativeAllocatorPoolConfig.POOL_FLIGHT)); assertTrue(poolNames.contains(NativeAllocatorPoolConfig.POOL_INGEST)); assertTrue(poolNames.contains(NativeAllocatorPoolConfig.POOL_QUERY)); - // Pool maxes match the operator-set values (rebalancer disabled by default, + // Pool maxes match the operator-set values (rebalancer disabled, // so initial limit == max). assertEquals(1L * 1024 * 1024 * 1024, allocator.getPoolAllocator(NativeAllocatorPoolConfig.POOL_FLIGHT).getLimit()); assertEquals(2L * 1024 * 1024 * 1024, allocator.getPoolAllocator(NativeAllocatorPoolConfig.POOL_INGEST).getLimit()); assertEquals(1L * 1024 * 1024 * 1024, allocator.getPoolAllocator(NativeAllocatorPoolConfig.POOL_QUERY).getLimit()); - - // Cluster-settings update consumers are registered: a PUT to a pool max must - // propagate to the live allocator. - cs.applySettings( - Settings.builder() - .put(NativeAllocatorPoolConfig.SETTING_ROOT_LIMIT, 8L * 1024 * 1024 * 1024) - .put(NativeAllocatorPoolConfig.SETTING_INGEST_MAX, 4L * 1024 * 1024 * 1024) - .build() - ); - assertEquals( - "buildAllocator must wire the INGEST_MAX cluster-settings consumer", - 4L * 1024 * 1024 * 1024, - allocator.getPoolAllocator(NativeAllocatorPoolConfig.POOL_INGEST).getLimit() - ); - } finally { - allocator.close(); - } - } - - public void testQueryMaxClusterSettingPropagatesToAllocator() { - // The full wired path: node starts at default settings, plugin registers - // consumers, operator PUTs a new max via _cluster/settings. - Settings nodeSettings = Settings.builder().put(NativeAllocatorPoolConfig.SETTING_ROOT_LIMIT, 8L * 1024 * 1024 * 1024).build(); - ClusterSettings cs = newClusterSettings(nodeSettings); - ArrowNativeAllocator allocator = newWiredAllocator(nodeSettings, cs); - try { - cs.applySettings( - Settings.builder() - .put(NativeAllocatorPoolConfig.SETTING_ROOT_LIMIT, 8L * 1024 * 1024 * 1024) - .put(NativeAllocatorPoolConfig.SETTING_QUERY_MAX, 1024L * 1024 * 1024) - .build() - ); - assertEquals( - "PUT to query max must update the live BufferAllocator limit", - 1024L * 1024 * 1024, - allocator.getPoolAllocator(NativeAllocatorPoolConfig.POOL_QUERY).getLimit() - ); - assertEquals(1024L * 1024 * 1024, allocator.getPoolMax(NativeAllocatorPoolConfig.POOL_QUERY)); - } finally { - allocator.close(); - } - } - - public void testFlightMinClusterSettingPropagatesToAllocator() { - // Min is the regression-prone path: prior to the live-propagation fix, - // setPoolMin only updated the poolMins map and operators got HTTP 200 with - // no observable behavior change. - Settings nodeSettings = Settings.builder().put(NativeAllocatorPoolConfig.SETTING_ROOT_LIMIT, 8L * 1024 * 1024 * 1024).build(); - ClusterSettings cs = newClusterSettings(nodeSettings); - ArrowNativeAllocator allocator = newWiredAllocator(nodeSettings, cs); - try { - // Pool starts at max (rebalancer disabled by default), so a min PUT below - // the current limit is a no-op on the live limit but updates poolMins. - // Use a min ABOVE the current limit to force the live raise path. - cs.applySettings( - Settings.builder() - .put(NativeAllocatorPoolConfig.SETTING_ROOT_LIMIT, 8L * 1024 * 1024 * 1024) - .put(NativeAllocatorPoolConfig.SETTING_FLIGHT_MAX, 4L * 1024 * 1024 * 1024) - .put(NativeAllocatorPoolConfig.SETTING_FLIGHT_MIN, 2L * 1024 * 1024 * 1024) - .build() - ); - assertEquals( - "PUT to flight min must update the recorded min for the rebalancer", - 2L * 1024 * 1024 * 1024, - allocator.getPoolMin(NativeAllocatorPoolConfig.POOL_FLIGHT) - ); - assertTrue( - "PUT to flight min must raise the live BufferAllocator limit when min exceeds current", - allocator.getPoolAllocator(NativeAllocatorPoolConfig.POOL_FLIGHT).getLimit() >= 2L * 1024 * 1024 * 1024 - ); - } finally { - allocator.close(); - } - } - - public void testRootLimitClusterSettingPropagatesToAllocator() { - Settings nodeSettings = Settings.builder().put(NativeAllocatorPoolConfig.SETTING_ROOT_LIMIT, 8L * 1024 * 1024 * 1024).build(); - ClusterSettings cs = newClusterSettings(nodeSettings); - ArrowNativeAllocator allocator = newWiredAllocator(nodeSettings, cs); - try { - cs.applySettings(Settings.builder().put(NativeAllocatorPoolConfig.SETTING_ROOT_LIMIT, 16L * 1024 * 1024 * 1024).build()); - assertEquals( - "PUT to root limit must update the RootAllocator's limit", - 16L * 1024 * 1024 * 1024, - allocator.getRootAllocator().getLimit() - ); } finally { allocator.close(); + plugin.close(); } } - public void testValidatorRejectsSumOfMinsExceedingRoot() { - // The cross-setting grouped validator must reject PUTs that would over- - // subscribe the root. Test the rejection path end-to-end through ClusterSettings. - // Set node.native_memory.limit=0b explicitly so pool maxes default to Long.MAX_VALUE - // — minmax path firing first. + public void testBuildAllocatorWithRebalancerPoolsStartAtMax() throws Exception { Settings nodeSettings = Settings.builder() - .put("node.native_memory.limit", "0b") - .put(NativeAllocatorPoolConfig.SETTING_ROOT_LIMIT, 10L * 1024 * 1024 * 1024) + .put("node.native_memory.limit", "1gb") + .put(NativeAllocatorPoolConfig.SETTING_FLIGHT_MIN, 10L * 1024 * 1024) + .put(NativeAllocatorPoolConfig.SETTING_FLIGHT_MAX, 200L * 1024 * 1024) + .put(NativeAllocatorPoolConfig.SETTING_INGEST_MIN, 20L * 1024 * 1024) + .put(NativeAllocatorPoolConfig.SETTING_INGEST_MAX, 200L * 1024 * 1024) + .put(NativeAllocatorPoolConfig.SETTING_QUERY_MIN, 10L * 1024 * 1024) + .put(NativeAllocatorPoolConfig.SETTING_QUERY_MAX, 200L * 1024 * 1024) + .put("native.allocator.rebalancer.enabled", true) + .put("native.allocator.rebalance.interval_seconds", 5L) .build(); ClusterSettings cs = newClusterSettings(nodeSettings); - ArrowNativeAllocator allocator = newWiredAllocator(nodeSettings, cs); + + ArrowBasePlugin plugin = new ArrowBasePlugin(); + long budget2 = ResourceTrackerSettings.NODE_NATIVE_MEMORY_LIMIT_SETTING.get(nodeSettings).getBytes(); + ArrowNativeAllocator allocator = plugin.buildAllocator(nodeSettings, cs, () -> budget2); try { - // root=10gb, flight_min=6gb, ingest_min=6gb => sum_mins=12gb > root=10gb. - IllegalArgumentException e = expectThrows( - IllegalArgumentException.class, - () -> cs.applySettings( - Settings.builder() - .put("node.native_memory.limit", "0b") - .put(NativeAllocatorPoolConfig.SETTING_ROOT_LIMIT, 10L * 1024 * 1024 * 1024) - .put(NativeAllocatorPoolConfig.SETTING_FLIGHT_MIN, 6L * 1024 * 1024 * 1024) - .put(NativeAllocatorPoolConfig.SETTING_INGEST_MIN, 6L * 1024 * 1024 * 1024) - .build() - ) - ); - assertTrue( - "expected sum-exceeds-root in error, got: " + e.getMessage(), - e.getMessage().contains("exceeds root limit") || e.getMessage().contains("Sum of pool minimums") - ); + // Pools always start at max regardless of rebalancer state + assertEquals(200L * 1024 * 1024, allocator.getPoolAllocator(NativeAllocatorPoolConfig.POOL_FLIGHT).getLimit()); + assertEquals(200L * 1024 * 1024, allocator.getPoolAllocator(NativeAllocatorPoolConfig.POOL_INGEST).getLimit()); + assertEquals(200L * 1024 * 1024, allocator.getPoolAllocator(NativeAllocatorPoolConfig.POOL_QUERY).getLimit()); } finally { allocator.close(); + plugin.close(); } } - public void testChildAllocatorInheritsParentCapAfterPoolLimitUpdate() { - // Sanity check for the AnalyticsSearchService / FlightTransport pattern: - // when a consumer creates a child of the framework's pool with Long.MAX_VALUE - // limit, a PUT to the pool's max takes effect on the child's allocations - // automatically via Arrow's parent-cap check at allocateBytes — no listener needed. - // - // The contract we rely on (Arrow Accountant.allocate, lines 191-203 in 18.3.0): - // when the child's reservation is exhausted, it calls parent.allocate(...) which - // checks the parent's allocationLimit on every allocation. Setting the child's own - // limit to Long.MAX_VALUE means the child has no own-cap on top of the parent's; - // setLimit on the parent is observed atomically by all subsequent allocations - // through any descendant. - Settings nodeSettings = Settings.builder().put(NativeAllocatorPoolConfig.SETTING_ROOT_LIMIT, 8L * 1024 * 1024 * 1024).build(); - ClusterSettings cs = newClusterSettings(nodeSettings); - ArrowNativeAllocator allocator = newWiredAllocator(nodeSettings, cs); - try { - BufferAllocator queryPool = allocator.getPoolAllocator(NativeAllocatorPoolConfig.POOL_QUERY); - BufferAllocator child = queryPool.newChildAllocator("consumer", 0, Long.MAX_VALUE); - try { - // Step 1: a small allocation through the child succeeds with the original pool max. - try (var buf = child.buffer(1024)) { - assertEquals("child accounting reflects allocation", 1024L, child.getAllocatedMemory()); - assertEquals("parent pool sees child allocation", 1024L, queryPool.getAllocatedMemory()); - } - - // Step 2: PUT a small pool max via cluster settings. - cs.applySettings( - Settings.builder() - .put(NativeAllocatorPoolConfig.SETTING_ROOT_LIMIT, 8L * 1024 * 1024 * 1024) - .put(NativeAllocatorPoolConfig.SETTING_QUERY_MAX, 1L * 1024 * 1024) // 1 MB - .build() - ); - assertEquals("pool's own limit reflects the PUT", 1L * 1024 * 1024, queryPool.getLimit()); - assertEquals("child's own limit is intentionally uncapped", Long.MAX_VALUE, child.getLimit()); - - // Step 3: allocations within the new parent cap still work. - try (var withinCap = child.buffer(512 * 1024)) { // 512 KB, under 1 MB cap - assertEquals(512L * 1024, child.getAllocatedMemory()); - } - - // Step 4: allocation exceeding the new parent cap fails — this is the - // behavior the deleted listener pattern was emulating, now provided - // natively by Arrow's parent-cap check. - expectThrows(OutOfMemoryException.class, () -> child.buffer(2L * 1024 * 1024)); // 2 MB, over 1 MB cap - } finally { - child.close(); - } - } finally { - allocator.close(); - } + public void testRebalanceIntervalSettingDefault() { + assertEquals(Long.valueOf(5L), ArrowBasePlugin.REBALANCE_INTERVAL_SETTING.get(Settings.EMPTY)); } } diff --git a/plugins/arrow-base/src/test/java/org/opensearch/arrow/allocator/ArrowNativeAllocatorTests.java b/plugins/arrow-base/src/test/java/org/opensearch/arrow/allocator/ArrowNativeAllocatorTests.java index dc10a93bd6b74..c9e370bc22413 100644 --- a/plugins/arrow-base/src/test/java/org/opensearch/arrow/allocator/ArrowNativeAllocatorTests.java +++ b/plugins/arrow-base/src/test/java/org/opensearch/arrow/allocator/ArrowNativeAllocatorTests.java @@ -10,6 +10,7 @@ import org.apache.arrow.memory.BufferAllocator; import org.opensearch.arrow.spi.NativeAllocator; +import org.opensearch.arrow.spi.PoolGroup; import org.opensearch.plugin.stats.NativeAllocatorPoolStats; import org.opensearch.test.OpenSearchTestCase; @@ -20,7 +21,7 @@ public class ArrowNativeAllocatorTests extends OpenSearchTestCase { @Override public void setUp() throws Exception { super.setUp(); - allocator = new ArrowNativeAllocator(1024L * 1024 * 1024); // 1 GB for tests + allocator = new ArrowNativeAllocator(); } @Override @@ -30,21 +31,21 @@ public void tearDown() throws Exception { } public void testCreatePool() { - NativeAllocator.PoolHandle handle = allocator.getOrCreatePool("test-pool", 100 * 1024 * 1024); + NativeAllocator.PoolHandle handle = allocator.getOrCreatePool("test-pool", 0L, 100 * 1024 * 1024, null); assertNotNull(handle); assertEquals(100 * 1024 * 1024, handle.limit()); assertEquals(0, handle.allocatedBytes()); } public void testGetOrCreatePoolIdempotent() { - NativeAllocator.PoolHandle first = allocator.getOrCreatePool("idempotent", 50 * 1024 * 1024); - NativeAllocator.PoolHandle second = allocator.getOrCreatePool("idempotent", 999 * 1024 * 1024); + NativeAllocator.PoolHandle first = allocator.getOrCreatePool("idempotent", 0L, 50 * 1024 * 1024, null); + NativeAllocator.PoolHandle second = allocator.getOrCreatePool("idempotent", 0L, 999 * 1024 * 1024, null); assertSame(first, second); assertEquals(50 * 1024 * 1024, second.limit()); } public void testPoolChildAllocation() { - allocator.getOrCreatePool("parent", 200 * 1024 * 1024); + allocator.getOrCreatePool("parent", 0L, 200 * 1024 * 1024, null); BufferAllocator child = allocator.getPoolAllocator("parent").newChildAllocator("child-1", 0, 50 * 1024 * 1024); try { child.buffer(1024).close(); @@ -55,7 +56,7 @@ public void testPoolChildAllocation() { } public void testSetPoolLimit() { - allocator.getOrCreatePool("resizable", 100 * 1024 * 1024); + allocator.getOrCreatePool("resizable", 0L, 100 * 1024 * 1024, null); allocator.setPoolLimit("resizable", 200 * 1024 * 1024); assertEquals(200 * 1024 * 1024, allocator.getPoolAllocator("resizable").getLimit()); } @@ -68,147 +69,54 @@ public void testGetPoolAllocatorNonExistent() { expectThrows(IllegalStateException.class, () -> allocator.getPoolAllocator("ghost")); } - public void testSetRootLimit() { - allocator.setRootLimit(512 * 1024 * 1024); - assertEquals(512 * 1024 * 1024, allocator.getRootAllocator().getLimit()); - } - public void testStats() { - allocator.getOrCreatePool("stats-pool", 64 * 1024 * 1024); + allocator.getOrCreatePool("stats-pool", 0L, 64 * 1024 * 1024, PoolGroup.SEARCH); NativeAllocatorPoolStats stats = allocator.stats(); assertNotNull(stats); - assertEquals(1024 * 1024 * 1024, stats.getRootLimitBytes()); - assertEquals(0, stats.getRootAllocatedBytes()); + assertEquals(-1, stats.getNativeAllocatedBytes()); + assertEquals(-1, stats.getNativeResidentBytes()); assertEquals(1, stats.getPools().size()); NativeAllocatorPoolStats.PoolStats poolStats = stats.getPools().get(0); assertEquals("stats-pool", poolStats.getName()); assertEquals(64 * 1024 * 1024, poolStats.getLimitBytes()); assertEquals(0, poolStats.getAllocatedBytes()); - // child_count is no longer rendered in stats; getPoolAllocator(...).getChildAllocators() - // is the runtime accessor for that detail if needed. } public void testStatsMultiplePools() { - allocator.getOrCreatePool("pool-a", 100 * 1024 * 1024); - allocator.getOrCreatePool("pool-b", 200 * 1024 * 1024); + allocator.getOrCreatePool("pool-a", 0L, 100 * 1024 * 1024, null); + allocator.getOrCreatePool("pool-b", 0L, 200 * 1024 * 1024, null); NativeAllocatorPoolStats stats = allocator.stats(); assertEquals(2, stats.getPools().size()); } public void testGetPoolNames() { - allocator.getOrCreatePool("alpha", 10 * 1024 * 1024); - allocator.getOrCreatePool("beta", 20 * 1024 * 1024); + allocator.getOrCreatePool("alpha", 0L, 10 * 1024 * 1024, null); + allocator.getOrCreatePool("beta", 0L, 20 * 1024 * 1024, null); assertTrue(allocator.getPoolNames().contains("alpha")); assertTrue(allocator.getPoolNames().contains("beta")); assertEquals(2, allocator.getPoolNames().size()); } - public void testRebalanceDistributesHeadroomToAllPools() { - allocator.setRootLimit(100 * 1024 * 1024); - allocator.getOrCreatePool("active", 10 * 1024 * 1024, 100 * 1024 * 1024); - allocator.getOrCreatePool("idle", 10 * 1024 * 1024, 100 * 1024 * 1024); - - // Simulate activity: allocate in "active" pool. - BufferAllocator activeAlloc = allocator.getPoolAllocator("active"); - BufferAllocator child = activeAlloc.newChildAllocator("worker", 0, 100 * 1024 * 1024); - var buf = child.buffer(5 * 1024 * 1024); - - try { - allocator.rebalance(); - - // Active pool gets bonus headroom on top of its min. - long activeLimit = activeAlloc.getLimit(); - assertTrue("Active pool limit should exceed min after rebalance, got " + activeLimit, activeLimit > 10 * 1024 * 1024); - - // Idle pool also receives headroom: distributing to all pools (not just - // currently-active ones) avoids the dead-pool corner case where a pool - // with min = 0 starts at limit = 0 and can never make a first allocation. - // Idle pools that don't end up needing the headroom return it on the next - // tick once they remain at zero allocation. - long idleLimit = allocator.getPoolAllocator("idle").getLimit(); - assertTrue("Idle pool should also receive headroom, got " + idleLimit, idleLimit > 10 * 1024 * 1024); - } finally { - buf.close(); - child.close(); - } - } - - public void testRebalanceLetsZeroMinPoolAllocate() { - // Regression test: under the previous "active pools only" rebalance algorithm, - // a pool with min = 0 would start at limit = 0 (rebalancer-on path), be unable - // to allocate, never become "active", and so never receive a bonus — permanently - // dead. Distributing headroom across all pools fixes the chicken-and-egg. - allocator.setRebalanceInterval(60); - allocator.setRootLimit(100 * 1024 * 1024); - allocator.getOrCreatePool("zero-min", 0L, 100 * 1024 * 1024); - try { - allocator.rebalance(); - BufferAllocator pool = allocator.getPoolAllocator("zero-min"); - assertTrue("Zero-min pool should receive headroom, got " + pool.getLimit(), pool.getLimit() > 0); - } finally { - allocator.setRebalanceInterval(0); - } - } - - public void testRebalanceNeverDropsBelowCurrentAllocation() { - allocator.setRootLimit(50 * 1024 * 1024); - allocator.getOrCreatePool("busy", 10 * 1024 * 1024); - - BufferAllocator pool = allocator.getPoolAllocator("busy"); - BufferAllocator child = pool.newChildAllocator("w", 0, 10 * 1024 * 1024); - var buf = child.buffer(8 * 1024 * 1024); // 8 MB allocated - - try { - allocator.rebalance(); - assertTrue("Limit should never drop below current allocation", pool.getLimit() >= pool.getAllocatedMemory()); - } finally { - buf.close(); - child.close(); - } - } - - public void testRebalanceWithNoPools() { - // Should not throw - allocator.rebalance(); - } - - public void testInitialLimitIsMaxWhenRebalancerDisabled() { - // Default tearDown allocator has rebalancer disabled (interval=0). - NativeAllocator.PoolHandle handle = allocator.getOrCreatePool("burst", 10 * 1024 * 1024, 100 * 1024 * 1024); - // With the rebalancer off, pools must start at their max so consumers can allocate - // immediately. Otherwise default-configured pools (min=0) would reject everything. + public void testInitialLimitIsMax() { + NativeAllocator.PoolHandle handle = allocator.getOrCreatePool("burst", 10 * 1024 * 1024, 100 * 1024 * 1024, null); assertEquals(100 * 1024 * 1024, handle.limit()); } - public void testInitialLimitIsMinWhenRebalancerEnabled() { - // Enabling the rebalancer reverts to the original "guarantee + burst" semantics: - // pools start at min and grow via the next rebalance tick. - allocator.setRebalanceInterval(60); // any positive value enables the flag - NativeAllocator.PoolHandle handle = allocator.getOrCreatePool("guaranteed", 10 * 1024 * 1024, 100 * 1024 * 1024); - assertEquals(10 * 1024 * 1024, handle.limit()); - // Disable so subsequent tests aren't affected by the scheduled task. - allocator.setRebalanceInterval(0); - } - public void testCloseReleasesAllPools() { - allocator.getOrCreatePool("close-test", 10 * 1024 * 1024); + allocator.getOrCreatePool("close-test", 0L, 10 * 1024 * 1024, null); allocator.close(); assertTrue(allocator.getPoolNames().isEmpty()); // Recreate for tearDown - allocator = new ArrowNativeAllocator(1024L * 1024 * 1024); + allocator = new ArrowNativeAllocator(); } - public void testSetPoolMinRaisesLiveLimitWhenRebalancerOff() { - // setPoolMin must affect the live BufferAllocator immediately, not just the - // poolMins map. Otherwise it's a Dynamic setting that returns HTTP 200 and - // does nothing observable until the operator also enables the rebalancer. - allocator.setRootLimit(100 * 1024 * 1024); - allocator.getOrCreatePool("p", 0L, 100 * 1024 * 1024); + public void testSetPoolMinRaisesLiveLimitWhenNeeded() { + allocator.getOrCreatePool("p", 0L, 100 * 1024 * 1024, null); BufferAllocator pool = allocator.getPoolAllocator("p"); long startLimit = pool.getLimit(); @@ -222,15 +130,45 @@ public void testSetPoolMinRaisesLiveLimitWhenRebalancerOff() { } public void testSetPoolMinDoesNotShrinkLiveLimit() { - // Dropping the min must not shrink an in-flight pool — the rebalancer is the - // only path that reduces limits, so a min change on its own should never - // reclaim capacity. - allocator.setRootLimit(100 * 1024 * 1024); - allocator.getOrCreatePool("p", 0L, 100 * 1024 * 1024); + allocator.getOrCreatePool("p", 0L, 100 * 1024 * 1024, null); BufferAllocator pool = allocator.getPoolAllocator("p"); long startLimit = pool.getLimit(); allocator.setPoolMin("p", 1L); assertEquals("dropping min must not shrink live limit", startLimit, pool.getLimit()); } + + public void testPoolGroupLimitListenerFiresWithCorrectSum() { + allocator.getOrCreatePool("ingest", 0L, 80 * 1024 * 1024, PoolGroup.INDEXING); + allocator.getOrCreatePool("write", 0L, 50 * 1024 * 1024, PoolGroup.INDEXING); + + java.util.concurrent.atomic.AtomicLong received = new java.util.concurrent.atomic.AtomicLong(-1); + allocator.addPoolGroupLimitListener(PoolGroup.INDEXING, received::set); + + // Change limits + allocator.setPoolEffectiveLimit("ingest", 60 * 1024 * 1024); + allocator.setPoolEffectiveLimit("write", 40 * 1024 * 1024); + + // Fire after all changes + allocator.firePoolGroupListeners(PoolGroup.INDEXING); + + // Listener should receive sum of effective limits: 60MB + 40MB = 100MB + assertEquals(100L * 1024 * 1024, received.get()); + } + + public void testPoolGroupLimitListenerNotFiredForOtherGroups() { + allocator.getOrCreatePool("flight", 0L, 50 * 1024 * 1024, PoolGroup.TRANSPORT); + allocator.getOrCreatePool("ingest", 0L, 80 * 1024 * 1024, PoolGroup.INDEXING); + + java.util.concurrent.atomic.AtomicLong transportReceived = new java.util.concurrent.atomic.AtomicLong(-1); + allocator.addPoolGroupLimitListener(PoolGroup.TRANSPORT, transportReceived::set); + + // Fire INDEXING group — should NOT trigger TRANSPORT listener + allocator.firePoolGroupListeners(PoolGroup.INDEXING); + assertEquals(-1, transportReceived.get()); + + // Fire TRANSPORT group — should trigger + allocator.firePoolGroupListeners(PoolGroup.TRANSPORT); + assertEquals(50L * 1024 * 1024, transportReceived.get()); + } } diff --git a/plugins/arrow-base/src/test/java/org/opensearch/arrow/allocator/NativeMemoryRebalancerTests.java b/plugins/arrow-base/src/test/java/org/opensearch/arrow/allocator/NativeMemoryRebalancerTests.java new file mode 100644 index 0000000000000..b92a477a67ea3 --- /dev/null +++ b/plugins/arrow-base/src/test/java/org/opensearch/arrow/allocator/NativeMemoryRebalancerTests.java @@ -0,0 +1,152 @@ +/* + * SPDX-License-Identifier: Apache-2.0 + * + * The OpenSearch Contributors require contributions made to + * this file be licensed under the Apache-2.0 license or a + * compatible open source license. + */ + +package org.opensearch.arrow.allocator; + +import org.apache.arrow.memory.ArrowBuf; +import org.apache.arrow.memory.BufferAllocator; +import org.opensearch.test.OpenSearchTestCase; + +import java.util.ArrayList; +import java.util.List; + +public class NativeMemoryRebalancerTests extends OpenSearchTestCase { + + private static final long MB = 1024L * 1024; + private static final long BUDGET = 100 * MB; + + private ArrowNativeAllocator allocator; + private NativeMemoryRebalancer rebalancer; + + @Override + public void setUp() throws Exception { + super.setUp(); + allocator = new ArrowNativeAllocator(); + allocator.setBudget(BUDGET); + rebalancer = new NativeMemoryRebalancer(allocator, () -> BUDGET, 0.75, 0.50, 0.10); + } + + @Override + public void tearDown() throws Exception { + allocator.close(); + super.tearDown(); + } + + public void testPoolsStartAtMax() { + allocator.getOrCreatePool("a", 5 * MB, 40 * MB, null); + allocator.getOrCreatePool("b", 10 * MB, 50 * MB, null); + + assertEquals(40 * MB, allocator.getPoolAllocator("a").getLimit()); + assertEquals(50 * MB, allocator.getPoolAllocator("b").getLimit()); + } + + public void testShrinksIdlePool() { + allocator.getOrCreatePool("idle", 5 * MB, 50 * MB, null); + allocator.getOrCreatePool("pressured", 5 * MB, 50 * MB, null); + + BufferAllocator pressuredPool = allocator.getPoolAllocator("pressured"); + ArrowBuf buf = pressuredPool.buffer((long) (50 * MB * 0.8)); + + try { + long idleLimitBefore = allocator.getPoolAllocator("idle").getLimit(); + rebalancer.rebalance(); + long idleLimitAfter = allocator.getPoolAllocator("idle").getLimit(); + assertTrue("Idle pool should shrink, was " + idleLimitBefore + " now " + idleLimitAfter, idleLimitAfter < idleLimitBefore); + } finally { + buf.close(); + } + } + + public void testGrowsPressuredPoolAboveMax() { + allocator.getOrCreatePool("idle", 5 * MB, 50 * MB, null); + allocator.getOrCreatePool("pressured", 5 * MB, 20 * MB, null); + + BufferAllocator pressuredPool = allocator.getPoolAllocator("pressured"); + ArrowBuf buf = pressuredPool.buffer((long) (20 * MB * 0.8)); + + try { + rebalancer.rebalance(); + long pressuredLimit = pressuredPool.getLimit(); + assertTrue("Pressured pool should grow above max (20MB), got " + pressuredLimit, pressuredLimit > 20 * MB); + } finally { + buf.close(); + } + } + + public void testNeverDropsBelowMin() { + allocator.getOrCreatePool("floored", 10 * MB, 50 * MB, null); + allocator.getOrCreatePool("pressured", 5 * MB, 50 * MB, null); + + BufferAllocator pressuredPool = allocator.getPoolAllocator("pressured"); + ArrowBuf buf = pressuredPool.buffer((long) (50 * MB * 0.8)); + + try { + for (int i = 0; i < 20; i++) { + rebalancer.rebalance(); + } + long flooredLimit = allocator.getPoolAllocator("floored").getLimit(); + assertTrue("Pool limit (" + flooredLimit + ") should not drop below min (10MB)", flooredLimit >= 10 * MB); + } finally { + buf.close(); + } + } + + public void testNoActionWhenNoPressure() { + allocator.getOrCreatePool("a", 5 * MB, 50 * MB, null); + allocator.getOrCreatePool("b", 5 * MB, 50 * MB, null); + + long limitA = allocator.getPoolAllocator("a").getLimit(); + long limitB = allocator.getPoolAllocator("b").getLimit(); + + rebalancer.rebalance(); + + assertEquals(limitA, allocator.getPoolAllocator("a").getLimit()); + assertEquals(limitB, allocator.getPoolAllocator("b").getLimit()); + } + + public void testResetAllPoolsToMax() { + allocator.getOrCreatePool("a", 5 * MB, 40 * MB, null); + allocator.getOrCreatePool("b", 5 * MB, 50 * MB, null); + + BufferAllocator bPool = allocator.getPoolAllocator("b"); + ArrowBuf buf = bPool.buffer((long) (50 * MB * 0.8)); + rebalancer.rebalance(); + buf.close(); + + allocator.resetAllPoolsToMax(); + + assertEquals(40 * MB, allocator.getPoolAllocator("a").getLimit()); + assertEquals(50 * MB, allocator.getPoolAllocator("b").getLimit()); + } + + public void testSumLimitsNeverExceedsBudget() { + allocator.getOrCreatePool("p1", 5 * MB, 30 * MB, null); + allocator.getOrCreatePool("p2", 5 * MB, 30 * MB, null); + allocator.getOrCreatePool("p3", 5 * MB, 30 * MB, null); + + List bufs = new ArrayList<>(); + for (String name : new String[] { "p1", "p2", "p3" }) { + BufferAllocator pool = allocator.getPoolAllocator(name); + bufs.add(pool.buffer((long) (pool.getLimit() * 0.8))); + } + + try { + for (int i = 0; i < 10; i++) { + rebalancer.rebalance(); + } + + long sumLimits = 0; + for (String name : new String[] { "p1", "p2", "p3" }) { + sumLimits += allocator.getPoolAllocator(name).getLimit(); + } + assertTrue("Sum of limits (" + sumLimits + ") should not exceed budget (" + BUDGET + ")", sumLimits <= BUDGET); + } finally { + bufs.forEach(ArrowBuf::close); + } + } +} diff --git a/plugins/arrow-flight-rpc/src/internalClusterTest/java/org/opensearch/arrow/flight/BackpressureProducerIT.java b/plugins/arrow-flight-rpc/src/internalClusterTest/java/org/opensearch/arrow/flight/BackpressureProducerIT.java index ff7f1359df089..a32a0ffe8bc58 100644 --- a/plugins/arrow-flight-rpc/src/internalClusterTest/java/org/opensearch/arrow/flight/BackpressureProducerIT.java +++ b/plugins/arrow-flight-rpc/src/internalClusterTest/java/org/opensearch/arrow/flight/BackpressureProducerIT.java @@ -122,9 +122,10 @@ protected Settings nodeSettings(int nodeOrdinal) { return Settings.builder() .put(super.nodeSettings(nodeOrdinal)) .put("node.native_memory.limit", "512mb") - .put(NativeAllocatorPoolConfig.SETTING_ROOT_LIMIT, 256 * MB) .put(NativeAllocatorPoolConfig.SETTING_FLIGHT_MAX, FLIGHT_POOL_CAP_BYTES) + .put(NativeAllocatorPoolConfig.SETTING_INGEST_MIN, 8 * MB) .put(NativeAllocatorPoolConfig.SETTING_INGEST_MAX, 16 * MB) + .put(NativeAllocatorPoolConfig.SETTING_QUERY_MIN, 8 * MB) .put(NativeAllocatorPoolConfig.SETTING_QUERY_MAX, 16 * MB) .put(ServerConfig.FLIGHT_OUTBOUND_BUFFER_THRESHOLD.getKey(), new ByteSizeValue(GRPC_THRESHOLD_BYTES, ByteSizeUnit.BYTES)) .build(); diff --git a/plugins/arrow-flight-rpc/src/internalClusterTest/java/org/opensearch/arrow/flight/NativeAllocatorBoundaryIT.java b/plugins/arrow-flight-rpc/src/internalClusterTest/java/org/opensearch/arrow/flight/NativeAllocatorBoundaryIT.java index 5633cfa429c5d..d849458c0b048 100644 --- a/plugins/arrow-flight-rpc/src/internalClusterTest/java/org/opensearch/arrow/flight/NativeAllocatorBoundaryIT.java +++ b/plugins/arrow-flight-rpc/src/internalClusterTest/java/org/opensearch/arrow/flight/NativeAllocatorBoundaryIT.java @@ -12,7 +12,6 @@ import org.apache.arrow.memory.BufferAllocator; import org.apache.arrow.memory.OutOfMemoryException; -import org.opensearch.action.admin.cluster.settings.ClusterUpdateSettingsResponse; import org.opensearch.arrow.allocator.ArrowBasePlugin; import org.opensearch.arrow.allocator.ArrowNativeAllocator; import org.opensearch.arrow.spi.NativeAllocatorPoolConfig; @@ -31,27 +30,17 @@ * *

Boots a single-node cluster with tight memory settings, then exercises * the actual Arrow allocation path to verify that the framework's - * configured caps are enforced at allocation time (not just at config-parse - * time). Complements unit-level tests in {@code ArrowBasePluginTests} by - * verifying that the production wiring (Guice -> ArrowNativeAllocator -> - * Arrow's RootAllocator chain) honors the caps end-to-end. - * - *

Each test sets explicit byte limits and allocates real - * {@link org.apache.arrow.memory.ArrowBuf} buffers, asserting either - * successful allocation or {@link OutOfMemoryException} based on whether - * the request fits within the configured cap. + * configured caps are enforced at allocation time. */ @ThreadLeakScope(ThreadLeakScope.Scope.NONE) @OpenSearchIntegTestCase.ClusterScope(scope = OpenSearchIntegTestCase.Scope.TEST, minNumDataNodes = 1, maxNumDataNodes = 1) public class NativeAllocatorBoundaryIT extends OpenSearchIntegTestCase { - /** 1 MiB. Chosen small enough that tests run fast but large enough that - * Arrow's internal accounting doesn't round it away. */ + /** 1 MiB. */ private static final long MB = 1024L * 1024; - /** Cap large enough for the framework's own bookkeeping but small enough - * to trigger OOM well before exhausting host memory. */ - private static final long ROOT_CAP_BYTES = 16 * MB; + /** Per-pool cap for tests. */ + private static final long POOL_CAP_BYTES = 16 * MB; @Override protected Collection> nodePlugins() { @@ -60,127 +49,83 @@ protected Collection> nodePlugins() { @Override protected Settings nodeSettings(int nodeOrdinal) { - // Set node.native_memory.limit explicitly so framework defaults derive - // from a known value rather than the (machine-dependent) ram-heap default. - // ROOT_LIMIT and pool maxes are then overridden per-test via cluster - // settings PUT or directly via this node-settings layer. return Settings.builder() .put(super.nodeSettings(nodeOrdinal)) .put("node.native_memory.limit", "256mb") - // Tight root cap: 16 MiB total Arrow framework budget. - .put(NativeAllocatorPoolConfig.SETTING_ROOT_LIMIT, ROOT_CAP_BYTES) - // Per-pool maxes set generously so per-pool caps don't trip - // before root.limit. Tests targeting per-pool caps override below. - .put(NativeAllocatorPoolConfig.SETTING_FLIGHT_MAX, ROOT_CAP_BYTES) - .put(NativeAllocatorPoolConfig.SETTING_INGEST_MAX, ROOT_CAP_BYTES) - .put(NativeAllocatorPoolConfig.SETTING_QUERY_MAX, ROOT_CAP_BYTES) + .put("native.allocator.rebalancer.enabled", false) + // Per-pool maxes set to a known value for testing. + .put(NativeAllocatorPoolConfig.SETTING_FLIGHT_MAX, POOL_CAP_BYTES) + .put(NativeAllocatorPoolConfig.SETTING_INGEST_MAX, POOL_CAP_BYTES) + .put(NativeAllocatorPoolConfig.SETTING_QUERY_MAX, POOL_CAP_BYTES) .build(); } /** - * Verifies that {@code parquet.native.pool.query.max} caps allocations - * through the QUERY pool: a buffer request exceeding the per-pool cap - * throws {@link OutOfMemoryException} even when root has headroom. + * Verifies that per-pool max caps allocations through the QUERY pool. */ public void testPoolMaxRejectsAllocationsBeyondCap() { - // Tighten QUERY pool to 4 MiB while leaving root at 16 MiB. - long poolCap = 4 * MB; - ClusterUpdateSettingsResponse resp = client().admin() - .cluster() - .prepareUpdateSettings() - .setTransientSettings(Settings.builder().put(NativeAllocatorPoolConfig.SETTING_QUERY_MAX, poolCap)) - .get(); - assertTrue("PUT to query.max must succeed", resp.isAcknowledged()); - ArrowNativeAllocator allocator = internalCluster().getInstance(ArrowNativeAllocator.class); BufferAllocator queryPool = allocator.getPoolAllocator(NativeAllocatorPoolConfig.POOL_QUERY); - assertThat("pool's live limit reflects the PUT", queryPool.getLimit(), is(poolCap)); + assertThat("pool's live limit reflects the configured max", queryPool.getLimit(), is(POOL_CAP_BYTES)); // Sub-cap allocation succeeds. try (var withinCap = queryPool.buffer(2 * MB)) { assertThat(queryPool.getAllocatedMemory(), greaterThanOrEqualTo(2 * MB)); } - // Cap+1 allocation fails — Arrow's parent-cap check at allocateBytes - // walks queryPool's allocationLimit and rejects. - expectThrows(OutOfMemoryException.class, () -> queryPool.buffer(8 * MB)); + // Cap+1 allocation fails. + expectThrows(OutOfMemoryException.class, () -> queryPool.buffer(POOL_CAP_BYTES + MB)); } /** - * Verifies that {@code native.allocator.root.limit} caps allocations - * across all pools combined: when the sum of in-flight pool allocations - * approaches the root cap, the next allocation is rejected at the root - * level even if each individual pool's max would allow it. + * Verifies that per-pool limits cap allocations: when one pool is full, + * allocations through it fail even if other pools have headroom. */ - public void testRootLimitRejectsAllocationsBeyondCap() { + public void testPoolLimitRejectsAllocationsBeyondCap() { ArrowNativeAllocator allocator = internalCluster().getInstance(ArrowNativeAllocator.class); BufferAllocator flightPool = allocator.getPoolAllocator(NativeAllocatorPoolConfig.POOL_FLIGHT); BufferAllocator queryPool = allocator.getPoolAllocator(NativeAllocatorPoolConfig.POOL_QUERY); - BufferAllocator root = allocator.getRootAllocator(); - - // Sanity-check setup: confirm the live limits match nodeSettings. - // If these fail, the test setup is wrong and the body's expectations are - // meaningless — surface the misconfiguration instead of misleading failures below. - assertThat("root.limit must match nodeSettings", root.getLimit(), is(ROOT_CAP_BYTES)); - assertThat("flight.max must match nodeSettings", flightPool.getLimit(), is(ROOT_CAP_BYTES)); - assertThat("query.max must match nodeSettings", queryPool.getLimit(), is(ROOT_CAP_BYTES)); - - // Hold 8 MiB through the FLIGHT pool. With root at 16 MiB this leaves 8 MiB - // headroom across the root. (Power-of-2 sizes avoid Arrow's chunked-allocation - // rounding surprises; e.g. a 12 MiB request actually consumes 16 MiB.) + + assertThat("flight.max must match nodeSettings", flightPool.getLimit(), is(POOL_CAP_BYTES)); + assertThat("query.max must match nodeSettings", queryPool.getLimit(), is(POOL_CAP_BYTES)); + + // Hold 8 MiB through the FLIGHT pool. try (var flightHold = flightPool.buffer(8 * MB)) { assertThat("FLIGHT pool reflects 8MB allocation", flightPool.getAllocatedMemory(), is(8L * MB)); - assertThat("root reflects 8MB allocation", root.getAllocatedMemory(), is(8L * MB)); - // A 4 MiB allocation through QUERY succeeds (within remaining root headroom). + // A 4 MiB allocation through QUERY succeeds (within its own pool cap). try (var queryFit = queryPool.buffer(4 * MB)) { - assertThat(allocator.getRootAllocator().getAllocatedMemory(), is(12L * MB)); + assertThat(queryPool.getAllocatedMemory(), greaterThanOrEqualTo(4 * MB)); } - // An 8 MiB allocation through QUERY would push the root past the 16 MiB cap - // (8 MiB FLIGHT + 8 MiB QUERY). Arrow's parent-cap check at allocateBytes - // walks queryPool -> root and rejects with OOM, even though QUERY's own - // (16 MiB) max would individually allow it. - expectThrows(OutOfMemoryException.class, () -> queryPool.buffer(16 * MB)); + // An allocation exceeding the QUERY pool's own cap fails. + expectThrows(OutOfMemoryException.class, () -> queryPool.buffer(POOL_CAP_BYTES + MB)); } } /** - * Verifies that a dynamic PUT to a pool's max takes effect on - * subsequent allocations through descendants of that pool. This is the - * behavior the deleted {@code NativeAllocatorListener} SPI was emulating; - * it is now provided natively by Arrow's parent-cap check at allocateBytes. + * Verifies that setPoolLimit dynamically adjusts the pool cap and + * subsequent allocations respect the new limit. */ - public void testDynamicPoolResizeAffectsInFlightAllocations() { + public void testSetPoolLimitAffectsInFlightAllocations() { ArrowNativeAllocator allocator = internalCluster().getInstance(ArrowNativeAllocator.class); BufferAllocator queryPool = allocator.getPoolAllocator(NativeAllocatorPoolConfig.POOL_QUERY); - // Step 1: create a child allocator at Long.MAX_VALUE — the AnalyticsSearchService / - // DefaultPlanExecutor pattern. The child intentionally has no own-cap; it relies - // on the parent pool's allocationLimit at allocation time. try (BufferAllocator child = queryPool.newChildAllocator("boundary-it-child", 0, Long.MAX_VALUE)) { - // Step 2: a small buffer through the child succeeds with the initial pool max. + // A small buffer through the child succeeds with the initial pool max. try (var buf = child.buffer(2 * MB)) { assertThat(child.getAllocatedMemory(), greaterThanOrEqualTo(2 * MB)); } - // Step 3: PUT a tighter pool max via cluster settings. + // Programmatically tighten the pool limit. long newPoolCap = 1 * MB; - ClusterUpdateSettingsResponse resp = client().admin() - .cluster() - .prepareUpdateSettings() - .setTransientSettings(Settings.builder().put(NativeAllocatorPoolConfig.SETTING_QUERY_MAX, newPoolCap)) - .get(); - assertTrue("PUT to query.max must succeed", resp.isAcknowledged()); - assertThat("pool's own limit reflects the PUT", queryPool.getLimit(), is(newPoolCap)); - assertThat("child's own limit is intentionally uncapped", child.getLimit(), is(Long.MAX_VALUE)); - - // Step 4: an allocation that fit before the resize now exceeds the parent cap. - // Arrow's parent-cap check at allocateBytes walks queryPool.allocationLimit - // and rejects — no listener machinery needed. + allocator.setPoolLimit(NativeAllocatorPoolConfig.POOL_QUERY, newPoolCap); + assertThat("pool's own limit reflects the update", queryPool.getLimit(), is(newPoolCap)); + + // An allocation that fit before the resize now exceeds the parent cap. expectThrows(OutOfMemoryException.class, () -> child.buffer(2 * MB)); - // Step 5: an allocation under the new cap still succeeds. + // An allocation under the new cap still succeeds. try (var smallBuf = child.buffer(512 * 1024)) { assertThat(child.getAllocatedMemory(), greaterThanOrEqualTo(512L * 1024)); } diff --git a/plugins/arrow-flight-rpc/src/internalClusterTest/java/org/opensearch/arrow/flight/NativeMemoryRebalancerIT.java b/plugins/arrow-flight-rpc/src/internalClusterTest/java/org/opensearch/arrow/flight/NativeMemoryRebalancerIT.java new file mode 100644 index 0000000000000..e28e373fe69a9 --- /dev/null +++ b/plugins/arrow-flight-rpc/src/internalClusterTest/java/org/opensearch/arrow/flight/NativeMemoryRebalancerIT.java @@ -0,0 +1,100 @@ +/* + * SPDX-License-Identifier: Apache-2.0 + * + * The OpenSearch Contributors require contributions made to + * this file be licensed under the Apache-2.0 license or a + * compatible open source license. + */ + +package org.opensearch.arrow.flight; + +import com.carrotsearch.randomizedtesting.annotations.ThreadLeakScope; + +import org.apache.arrow.memory.ArrowBuf; +import org.apache.arrow.memory.BufferAllocator; +import org.opensearch.arrow.allocator.ArrowBasePlugin; +import org.opensearch.arrow.allocator.ArrowNativeAllocator; +import org.opensearch.arrow.spi.NativeAllocatorPoolConfig; +import org.opensearch.common.settings.Settings; +import org.opensearch.plugins.Plugin; +import org.opensearch.test.OpenSearchIntegTestCase; + +import java.util.Collection; +import java.util.List; + +import static org.hamcrest.Matchers.lessThanOrEqualTo; + +/** + * Integration test for the NativeMemoryRebalancer. + * + *

Boots a single-node cluster with the rebalancer enabled and verifies that + * pools start at their max and the rebalancer shrinks idle pools / grows pressured ones. + */ +@ThreadLeakScope(ThreadLeakScope.Scope.NONE) +@OpenSearchIntegTestCase.ClusterScope(scope = OpenSearchIntegTestCase.Scope.TEST, minNumDataNodes = 1, maxNumDataNodes = 1) +public class NativeMemoryRebalancerIT extends OpenSearchIntegTestCase { + + private static final long MB = 1024L * 1024; + + @Override + protected Collection> nodePlugins() { + return List.of(ArrowBasePlugin.class); + } + + @Override + protected Settings nodeSettings(int nodeOrdinal) { + return Settings.builder() + .put(super.nodeSettings(nodeOrdinal)) + .put("native.allocator.rebalancer.enabled", true) + .put("native.allocator.rebalance.interval_seconds", 1) + .put("node.native_memory.limit", "1gb") + .put(NativeAllocatorPoolConfig.SETTING_FLIGHT_MIN, 5 * MB) + .put(NativeAllocatorPoolConfig.SETTING_FLIGHT_MAX, 200 * MB) + .put(NativeAllocatorPoolConfig.SETTING_INGEST_MIN, 5 * MB) + .put(NativeAllocatorPoolConfig.SETTING_INGEST_MAX, 200 * MB) + .put(NativeAllocatorPoolConfig.SETTING_QUERY_MIN, 5 * MB) + .put(NativeAllocatorPoolConfig.SETTING_QUERY_MAX, 200 * MB) + .build(); + } + + /** + * Verifies that pools start at their max (before rebalancer shrinks them). + * Uses a long rebalancer interval to avoid race conditions. + */ + public void testPoolsStartAtMax() { + ArrowNativeAllocator allocator = internalCluster().getInstance(ArrowNativeAllocator.class); + BufferAllocator ingestPool = allocator.getPoolAllocator(NativeAllocatorPoolConfig.POOL_INGEST); + + // The rebalancer may have already run (1s interval), so the pool may have shrunk. + // Verify it's at least at min and the configured max is correct. + long max = allocator.getPoolMax(NativeAllocatorPoolConfig.POOL_INGEST); + assertEquals("Ingest pool max should be configured at 200MB", 200 * MB, max); + // Pool limit should be between min and max (rebalancer may have shrunk it) + long limit = ingestPool.getLimit(); + assertThat("Ingest pool limit should be >= min", limit, org.hamcrest.Matchers.greaterThanOrEqualTo(5 * MB)); + assertThat("Ingest pool limit should be <= max", limit, lessThanOrEqualTo(200 * MB)); + } + + /** + * Verifies that an idle pool shrinks after rebalancer ticks when another pool is pressured. + */ + public void testIdlePoolShrinksWhenOtherPressured() throws Exception { + ArrowNativeAllocator allocator = internalCluster().getInstance(ArrowNativeAllocator.class); + BufferAllocator ingestPool = allocator.getPoolAllocator(NativeAllocatorPoolConfig.POOL_INGEST); + BufferAllocator flightPool = allocator.getPoolAllocator(NativeAllocatorPoolConfig.POOL_FLIGHT); + + // Allocate > 75% of ingest pool to create pressure + long toAllocate = (long) (ingestPool.getLimit() * 0.8); + ArrowBuf buf = ingestPool.buffer(toAllocate); + + try { + // Flight pool is idle — wait for rebalancer to shrink it + assertBusy(() -> { + long flightLimit = flightPool.getLimit(); + assertThat("Flight pool should shrink when idle", flightLimit, org.hamcrest.Matchers.lessThan(200 * MB)); + }); + } finally { + buf.close(); + } + } +} diff --git a/plugins/arrow-flight-rpc/src/internalClusterTest/java/org/opensearch/arrow/flight/UnifiedNativeMemoryStatsIT.java b/plugins/arrow-flight-rpc/src/internalClusterTest/java/org/opensearch/arrow/flight/UnifiedNativeMemoryStatsIT.java new file mode 100644 index 0000000000000..fd0e165103e86 --- /dev/null +++ b/plugins/arrow-flight-rpc/src/internalClusterTest/java/org/opensearch/arrow/flight/UnifiedNativeMemoryStatsIT.java @@ -0,0 +1,95 @@ +/* + * SPDX-License-Identifier: Apache-2.0 + * + * The OpenSearch Contributors require contributions made to + * this file be licensed under the Apache-2.0 license or a + * compatible open source license. + */ + +package org.opensearch.arrow.flight; + +import com.carrotsearch.randomizedtesting.annotations.ThreadLeakScope; + +import org.opensearch.action.admin.cluster.node.stats.NodesStatsRequest; +import org.opensearch.action.admin.cluster.node.stats.NodesStatsResponse; +import org.opensearch.arrow.allocator.ArrowBasePlugin; +import org.opensearch.common.settings.Settings; +import org.opensearch.plugin.stats.NativeAllocatorPoolStats; +import org.opensearch.plugins.Plugin; +import org.opensearch.test.OpenSearchIntegTestCase; + +import java.util.Collection; +import java.util.List; +import java.util.Set; +import java.util.stream.Collectors; + +import static org.hamcrest.Matchers.greaterThan; +import static org.hamcrest.Matchers.greaterThanOrEqualTo; +import static org.hamcrest.Matchers.hasItems; +import static org.hamcrest.Matchers.is; +import static org.hamcrest.Matchers.notNullValue; + +/** + * Integration test verifying the unified native memory stats endpoint. + * Boots a single-node cluster with ArrowBasePlugin and confirms that + * all registered pools (Arrow + virtual) appear in _nodes/stats/native_memory. + */ +@ThreadLeakScope(ThreadLeakScope.Scope.NONE) +@OpenSearchIntegTestCase.ClusterScope(scope = OpenSearchIntegTestCase.Scope.TEST, minNumDataNodes = 1, maxNumDataNodes = 1) +public class UnifiedNativeMemoryStatsIT extends OpenSearchIntegTestCase { + + @Override + protected Collection> nodePlugins() { + return List.of(ArrowBasePlugin.class); + } + + @Override + protected Settings nodeSettings(int nodeOrdinal) { + return Settings.builder().put(super.nodeSettings(nodeOrdinal)).put("node.native_memory.limit", "1gb").build(); + } + + /** + * Verifies that the Arrow pools (flight, ingest, query) are visible in + * _nodes/stats/native_memory with correct structure. + */ + public void testArrowPoolsVisibleInStats() { + NodesStatsResponse response = client().admin() + .cluster() + .prepareNodesStats() + .addMetric(NodesStatsRequest.Metric.NATIVE_MEMORY.metricName()) + .get(); + + assertThat(response.getNodes().isEmpty(), is(false)); + NativeAllocatorPoolStats stats = response.getNodes().get(0).getNativeAllocatorStats(); + assertThat("native_memory stats should be present", stats, notNullValue()); + + // Dump the stats for debugging + StringBuilder sb = new StringBuilder(); + sb.append("nativeAllocated=").append(stats.getNativeAllocatedBytes()); + sb.append(", nativeResident=").append(stats.getNativeResidentBytes()); + sb.append(", pools=["); + for (NativeAllocatorPoolStats.PoolStats p : stats.getPools()) { + sb.append(p.getName()) + .append("(alloc=") + .append(p.getAllocatedBytes()) + .append(",peak=") + .append(p.getPeakBytes()) + .append(",limit=") + .append(p.getLimitBytes()) + .append(") "); + } + sb.append("]"); + logger.info("=== NATIVE_MEMORY STATS: {} ===", sb); + + // All Arrow pools should be present + Set poolNames = stats.getPools().stream().map(NativeAllocatorPoolStats.PoolStats::getName).collect(Collectors.toSet()); + assertThat(poolNames, hasItems("flight", "ingest", "query")); + + // Each pool should have limit > 0 (derived from 1gb native_memory.limit) + for (NativeAllocatorPoolStats.PoolStats pool : stats.getPools()) { + assertThat("Pool '" + pool.getName() + "' should have limit > 0", pool.getLimitBytes(), greaterThan(0L)); + assertThat("Pool '" + pool.getName() + "' allocated should be >= 0", pool.getAllocatedBytes(), greaterThanOrEqualTo(0L)); + } + } + +} diff --git a/plugins/arrow-flight-rpc/src/test/java/org/opensearch/arrow/flight/transport/FlightTransportTestBase.java b/plugins/arrow-flight-rpc/src/test/java/org/opensearch/arrow/flight/transport/FlightTransportTestBase.java index f5eacc771be92..add8e0a2bbfee 100644 --- a/plugins/arrow-flight-rpc/src/test/java/org/opensearch/arrow/flight/transport/FlightTransportTestBase.java +++ b/plugins/arrow-flight-rpc/src/test/java/org/opensearch/arrow/flight/transport/FlightTransportTestBase.java @@ -91,8 +91,8 @@ public void setUp() throws Exception { // FlightTransport sources its allocator from the framework's FLIGHT pool. Construct one // here so the test has a usable allocator; tearDown closes it. - nativeAllocator = new ArrowNativeAllocator(Long.MAX_VALUE); - nativeAllocator.getOrCreatePool(NativeAllocatorPoolConfig.POOL_FLIGHT, 0L, Long.MAX_VALUE); + nativeAllocator = new ArrowNativeAllocator(); + nativeAllocator.getOrCreatePool(NativeAllocatorPoolConfig.POOL_FLIGHT, 0L, Long.MAX_VALUE, null); flightTransport = new FlightTransport( settings, diff --git a/sandbox/libs/dataformat-native/rust/common/src/lib.rs b/sandbox/libs/dataformat-native/rust/common/src/lib.rs index 0f4b8c132407f..c44fa871c4fb3 100644 --- a/sandbox/libs/dataformat-native/rust/common/src/lib.rs +++ b/sandbox/libs/dataformat-native/rust/common/src/lib.rs @@ -11,6 +11,7 @@ pub mod error; pub mod logger; pub mod allocator; +pub mod memory_pool; // Re-export the proc macro so plugins use `#[native_bridge_common::ffm_safe]` pub use native_bridge_macros::ffm_safe; diff --git a/sandbox/libs/dataformat-native/rust/common/src/memory_pool.rs b/sandbox/libs/dataformat-native/rust/common/src/memory_pool.rs new file mode 100644 index 0000000000000..7fc5c7c83a20c --- /dev/null +++ b/sandbox/libs/dataformat-native/rust/common/src/memory_pool.rs @@ -0,0 +1,370 @@ +/* + * SPDX-License-Identifier: Apache-2.0 + * + * The OpenSearch Contributors require contributions made to + * this file be licensed under the Apache-2.0 license or a + * compatible open source license. + */ + +//! Memory pool for tracking native memory usage across write and merge operations. +//! +//! Provides an atomic counter with a configurable limit. Operations that allocate +//! significant memory call `try_grow` before allocating and `shrink` after freeing. +//! The pool rejects allocations that would exceed the configured limit. +//! +//! `MemoryReservation` is an RAII handle that automatically returns memory to the +//! pool on drop, preventing leaks even on error paths. + +use std::sync::atomic::{AtomicUsize, Ordering}; +use std::sync::{Arc, Condvar, Mutex}; +use std::time::Duration; +use std::fmt; + +/// Default timeout for blocking wait (300 seconds). +pub const DEFAULT_WAIT_TIMEOUT: Duration = Duration::from_secs(300); + +/// Merge operations can wait longer (600 seconds). +pub const MERGE_WAIT_TIMEOUT: Duration = Duration::from_secs(600); + +/// Controls whether an allocation blocks or rejects immediately. +#[derive(Debug, Clone)] +pub enum PoolBehavior { + /// Block until memory is available, up to the given timeout. + Wait(Duration), + /// Fail immediately if pool is full. + Reject, +} + +/// Error returned when a pool cannot satisfy an allocation request. +#[derive(Debug, Clone)] +pub struct PoolExhausted { + pub pool_name: &'static str, + pub requested: usize, + pub used: usize, + pub limit: usize, +} + +impl fmt::Display for PoolExhausted { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + write!( + f, + "[{}] memory limit exceeded: requested {} bytes, used {}, limit {}", + self.pool_name, self.requested, self.used, self.limit + ) + } +} + +impl std::error::Error for PoolExhausted {} + +/// Error returned when wait_and_grow times out. +#[derive(Debug, Clone)] +pub struct PoolTimeout { + pub pool_name: &'static str, + pub requested: usize, + pub used: usize, + pub limit: usize, + pub waited: Duration, +} + +impl fmt::Display for PoolTimeout { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + write!( + f, + "[{}] timed out waiting for {} bytes after {:?} (used: {}, limit: {})", + self.pool_name, self.requested, self.waited, self.used, self.limit + ) + } +} + +impl std::error::Error for PoolTimeout {} + +/// A node-level memory pool backed by an atomic counter with blocking wait support. +pub struct MemoryPool { + name: &'static str, + used: AtomicUsize, + limit: AtomicUsize, + peak: AtomicUsize, + notify: Condvar, + notify_lock: Mutex<()>, +} + +impl fmt::Debug for MemoryPool { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + f.debug_struct("MemoryPool") + .field("name", &self.name) + .field("used", &self.used.load(Ordering::Relaxed)) + .field("limit", &self.limit.load(Ordering::Relaxed)) + .field("peak", &self.peak.load(Ordering::Relaxed)) + .finish() + } +} + +impl MemoryPool { + /// Create a new pool. `limit = 0` means unlimited. + pub fn new(name: &'static str, limit: usize) -> Self { + Self { + name, + used: AtomicUsize::new(0), + limit: AtomicUsize::new(limit), + peak: AtomicUsize::new(0), + notify: Condvar::new(), + notify_lock: Mutex::new(()), + } + } + + /// Attempt to reserve `bytes`. Returns error if it would exceed the limit. + pub fn try_grow(&self, bytes: usize) -> Result<(), PoolExhausted> { + if bytes == 0 { + return Ok(()); + } + let limit = self.limit.load(Ordering::Relaxed); + let result = self.used.fetch_update(Ordering::Relaxed, Ordering::Relaxed, |used| { + let new_used = used.checked_add(bytes)?; + if limit > 0 && new_used > limit { + None + } else { + Some(new_used) + } + }); + + match result { + Ok(old) => { + self.peak.fetch_max(old + bytes, Ordering::Relaxed); + Ok(()) + } + Err(_) => Err(PoolExhausted { + pool_name: self.name, + requested: bytes, + used: self.used.load(Ordering::Relaxed), + limit, + }), + } + } + + /// Blocks until `bytes` can be reserved, or timeout expires. + pub fn wait_and_grow(&self, bytes: usize, timeout: Duration) -> Result<(), PoolTimeout> { + if bytes == 0 { + return Ok(()); + } + if self.try_grow(bytes).is_ok() { + return Ok(()); + } + + let start = std::time::Instant::now(); + loop { + let elapsed = start.elapsed(); + if elapsed >= timeout { + let used = self.used.load(Ordering::Relaxed); + let limit = self.limit.load(Ordering::Relaxed); + return Err(PoolTimeout { + pool_name: self.name, + requested: bytes, + used, + limit, + waited: elapsed, + }); + } + + let remaining = timeout - elapsed; + let guard = self.notify_lock.lock().unwrap(); + let _ = self.notify.wait_timeout(guard, remaining.min(Duration::from_secs(1))).unwrap(); + + if self.try_grow(bytes).is_ok() { + return Ok(()); + } + } + } + + /// Infallible grow — use when the allocation has already happened. + pub fn grow(&self, bytes: usize) { + if bytes == 0 { + return; + } + let new_used = self.used.fetch_add(bytes, Ordering::Relaxed) + bytes; + self.peak.fetch_max(new_used, Ordering::Relaxed); + } + + /// Release `bytes` back to the pool. Notifies any waiting threads. + pub fn shrink(&self, bytes: usize) { + if bytes == 0 { + return; + } + self.used + .fetch_update(Ordering::Relaxed, Ordering::Relaxed, |current| { + Some(current.saturating_sub(bytes)) + }) + .unwrap(); + self.notify.notify_all(); + } + + pub fn used(&self) -> usize { + self.used.load(Ordering::Relaxed) + } + + pub fn peak(&self) -> usize { + self.peak.load(Ordering::Relaxed) + } + + pub fn limit(&self) -> usize { + self.limit.load(Ordering::Relaxed) + } + + pub fn name(&self) -> &'static str { + self.name + } + + /// Atomically update the limit. Called by the Java rebalancer. + pub fn set_limit(&self, new_limit: usize) { + self.limit.store(new_limit, Ordering::Release); + // Wake waiters — new limit might allow blocked allocations + self.notify.notify_all(); + } +} + +/// RAII handle that tracks a portion of memory reserved from a [`MemoryPool`]. +/// Automatically releases all held memory on drop. +pub struct MemoryReservation { + pool: Arc, + consumer: &'static str, + size: usize, + behavior: PoolBehavior, +} + +impl MemoryReservation { + pub fn new(pool: &Arc, consumer: &'static str, behavior: PoolBehavior) -> Self { + Self { + pool: Arc::clone(pool), + consumer, + size: 0, + behavior, + } + } + + /// Grow based on the reservation's behavior: block (Wait) or reject (Reject). + pub fn request(&mut self, bytes: usize) -> Result<(), Box> { + match &self.behavior { + PoolBehavior::Reject => { + self.pool.try_grow(bytes)?; + self.size += bytes; + Ok(()) + } + PoolBehavior::Wait(timeout) => { + self.pool.wait_and_grow(bytes, *timeout)?; + self.size += bytes; + Ok(()) + } + } + } + + /// Infallible grow. + pub fn grow(&mut self, bytes: usize) { + self.pool.grow(bytes); + self.size += bytes; + } + + /// Release `bytes` from this reservation. + pub fn shrink(&mut self, bytes: usize) { + let actual = bytes.min(self.size); + self.pool.shrink(actual); + self.size -= actual; + } + + /// Release all memory back to the pool. + pub fn free(&mut self) -> usize { + let s = self.size; + if s > 0 { + self.pool.shrink(s); + self.size = 0; + } + s + } + + pub fn size(&self) -> usize { + self.size + } + + pub fn consumer(&self) -> &'static str { + self.consumer + } + + /// Create a sibling reservation from the same pool with a different consumer name. + pub fn child(&self, consumer: &'static str) -> Self { + Self { + pool: Arc::clone(&self.pool), + consumer, + size: 0, + behavior: self.behavior.clone(), + } + } +} + +impl Drop for MemoryReservation { + fn drop(&mut self) { + if self.size > 0 { + self.pool.shrink(self.size); + } + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn test_try_grow_within_limit() { + let pool = Arc::new(MemoryPool::new("test", 1024)); + let mut res = MemoryReservation::new(&pool, "writer", PoolBehavior::Reject); + assert!(res.request(512).is_ok()); + assert_eq!(res.size(), 512); + assert_eq!(pool.used(), 512); + } + + #[test] + fn test_try_grow_exceeds_limit() { + let pool = Arc::new(MemoryPool::new("test", 1024)); + let mut res = MemoryReservation::new(&pool, "writer", PoolBehavior::Reject); + assert!(res.request(2048).is_err()); + assert_eq!(res.size(), 0); + assert_eq!(pool.used(), 0); + } + + #[test] + fn test_drop_releases_memory() { + let pool = Arc::new(MemoryPool::new("test", 1024)); + { + let mut res = MemoryReservation::new(&pool, "writer", PoolBehavior::Reject); + res.request(500).unwrap(); + assert_eq!(pool.used(), 500); + } + assert_eq!(pool.used(), 0); + } + + #[test] + fn test_set_limit_allows_growth() { + let pool = Arc::new(MemoryPool::new("test", 100)); + let mut res = MemoryReservation::new(&pool, "writer", PoolBehavior::Reject); + assert!(res.request(200).is_err()); + pool.set_limit(500); + assert!(res.request(200).is_ok()); + } + + #[test] + fn test_peak_tracking() { + let pool = Arc::new(MemoryPool::new("test", 1024)); + let mut res = MemoryReservation::new(&pool, "writer", PoolBehavior::Reject); + res.request(800).unwrap(); + res.shrink(500); + assert_eq!(pool.peak(), 800); + assert_eq!(pool.used(), 300); + } + + #[test] + fn test_child_reservation() { + let pool = Arc::new(MemoryPool::new("test", 1024)); + let res = MemoryReservation::new(&pool, "parent", PoolBehavior::Reject); + let mut child = res.child("child"); + child.request(100).unwrap(); + assert_eq!(child.consumer(), "child"); + assert_eq!(pool.used(), 100); + } +} diff --git a/sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/DataFusionPlugin.java b/sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/DataFusionPlugin.java index 9f15bd404f21f..22d7fadac0d08 100644 --- a/sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/DataFusionPlugin.java +++ b/sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/DataFusionPlugin.java @@ -13,6 +13,9 @@ import org.opensearch.action.ActionRequest; import org.opensearch.analytics.spi.AnalyticsSearchBackendPlugin; import org.opensearch.analytics.spi.QueryExecutionMetrics; +import org.opensearch.arrow.allocator.ArrowNativeAllocator; +import org.opensearch.arrow.spi.NativeAllocator; +import org.opensearch.arrow.spi.PoolGroup; import org.opensearch.be.datafusion.action.stats.DataFusionStatsActionType; import org.opensearch.be.datafusion.action.stats.RestDataFusionStatsAction; import org.opensearch.be.datafusion.action.stats.TransportDataFusionStatsAction; @@ -20,6 +23,7 @@ import org.opensearch.cluster.metadata.IndexNameExpressionResolver; import org.opensearch.cluster.node.DiscoveryNodes; import org.opensearch.cluster.service.ClusterService; +import org.opensearch.common.Nullable; import org.opensearch.common.settings.ClusterSettings; import org.opensearch.common.settings.IndexScopedSettings; import org.opensearch.common.settings.Setting; @@ -95,7 +99,7 @@ public class DataFusionPlugin extends Plugin * ({@link ResourceTrackerSettings#NODE_NATIVE_MEMORY_LIMIT_SETTING}), which is * the same off-heap budget admission control throttles against. The DataFusion Rust * runtime is the dominant native-memory consumer for analytics workloads (see PR #21732 - * partitioning model), so the default takes 75% of {@code node.native_memory.limit}. + * partitioning model), so the default takes 74% of {@code node.native_memory.limit}. * If the AC limit is unset (== 0), the default is {@link Long#MAX_VALUE} — unbounded — to * preserve pre-AC behaviour rather than make up a number from JVM heap (which is a * separate, already-allocated region with no relation to native-memory sizing). @@ -118,7 +122,7 @@ public class DataFusionPlugin extends Plugin ); /** - * Computes the default for {@link #DATAFUSION_MEMORY_POOL_LIMIT} as 75% of + * Computes the default for {@link #DATAFUSION_MEMORY_POOL_LIMIT} as 74% of * {@link ResourceTrackerSettings#NODE_NATIVE_MEMORY_LIMIT_SETTING}, falling back to * {@link Long#MAX_VALUE} when AC is unconfigured. * @@ -134,10 +138,10 @@ static String deriveMemoryPoolLimitDefault(Settings settings) { if (nativeLimit.getBytes() <= 0) { return Long.toString(Long.MAX_VALUE); } - // 75% of node.native_memory.limit. DataFusion is the dominant native consumer for + // 74% of node.native_memory.limit. DataFusion is the dominant native consumer for // analytics workloads; operators tune via the dynamic setting once they characterize // their workload. - long pool = Math.max(0L, nativeLimit.getBytes() * 75 / 100); + long pool = Math.max(0L, nativeLimit.getBytes() * 74 / 100); return Long.toString(pool); } @@ -383,6 +387,39 @@ public Collection createComponents( IndexNameExpressionResolver indexNameExpressionResolver, Supplier repositoriesServiceSupplier, DataFormatRegistry dataFormatRegistry + ) { + return createComponents( + client, + clusterService, + threadPool, + resourceWatcherService, + scriptService, + xContentRegistry, + environment, + nodeEnvironment, + namedWriteableRegistry, + indexNameExpressionResolver, + repositoriesServiceSupplier, + dataFormatRegistry, + null + ); + } + + @Override + public Collection createComponents( + Client client, + ClusterService clusterService, + ThreadPool threadPool, + ResourceWatcherService resourceWatcherService, + ScriptService scriptService, + NamedXContentRegistry xContentRegistry, + Environment environment, + NodeEnvironment nodeEnvironment, + NamedWriteableRegistry namedWriteableRegistry, + IndexNameExpressionResolver indexNameExpressionResolver, + Supplier repositoriesServiceSupplier, + DataFormatRegistry dataFormatRegistry, + @Nullable NativeAllocator nativeAllocator ) { this.dataFormatRegistry = dataFormatRegistry; this.clusterService = clusterService; @@ -402,13 +439,8 @@ public Collection createComponents( dataFusionService.start(); logger.debug("DataFusion plugin initialized — memory pool {}B, spill limit {}B", memoryPoolLimit, spillMemoryLimit); - // Wire the dynamic memory pool limit setting to the native runtime so updates via the - // cluster settings API take effect without restarting the node. The framework's - // parquet.native.pool.datafusion.{min,max} controls the Java-side Arrow pool that - // sources the per-query allocators handed to DataFusion; this setting controls the - // Rust runtime's internal MemoryPool used by query execution. They're separate - // accounting layers — operators tune them independently. - clusterService.getClusterSettings().addSettingsUpdateConsumer(DATAFUSION_MEMORY_POOL_LIMIT, this::updateMemoryPoolLimit); + // Wire the dynamic spill limit setting to the native runtime so updates via the + // cluster settings API take effect without restarting the node. clusterService.getClusterSettings().addSettingsUpdateConsumer(DATAFUSION_SPILL_MEMORY_LIMIT, this::updateSpillMemoryLimit); clusterService.getClusterSettings().addSettingsUpdateConsumer(DATAFUSION_MIN_TARGET_PARTITIONS, this::updateMinTargetPartitions); clusterService.getClusterSettings() @@ -447,19 +479,48 @@ public Collection createComponents( this.datafusionSettings = new DatafusionSettings(clusterService); - // Expose per-task native-memory usage to search backpressure. The tracker calls - // this supplier once per refresh (invoked by the backpressure service at the top of - // doRun() and nodeStats()), snapshotting all live queries in one FFM call. Per-task - // evaluation then reads from the tracker's cached map — no FFM call per task. - // - // The OpenSearch task id is used as the DataFusion context_id at query launch - // (see ShardScanInstructionHandler / DatafusionSearchExecEngine), so the map is - // already keyed by Task#getId on the consumer side. + // Expose per-task native-memory usage to search backpressure. NativeMemoryUsageTracker.setSnapshotSupplier(this::currentBytesByTaskId); NativeMemoryUsageTracker.setNativeMemoryBudgetSupplier(() -> DATAFUSION_MEMORY_POOL_LIMIT.get(clusterService.getSettings())); this.substraitExtensions = loadSubstraitExtensions(); + // Register with the unified allocator if available + if (nativeAllocator != null) { + ClusterSettings clusterSettings = clusterService.getClusterSettings(); + ArrowNativeAllocator arrowAllocator = (ArrowNativeAllocator) nativeAllocator; + + NativeAllocator.VirtualPoolHandle dfPool = arrowAllocator.registerVirtualPool( + DatafusionSettings.POOL_DATAFUSION, + DatafusionSettings.DATAFUSION_MEMORY_POOL_MIN.get(settings), + DATAFUSION_MEMORY_POOL_LIMIT.get(settings), + PoolGroup.SEARCH, + this::updateMemoryPoolLimit + ); + + arrowAllocator.addStatsRefresher(() -> { + if (dataFusionService != null) { + long usage = dataFusionService.getMemoryPoolUsage(); + dfPool.updateStats(usage, usage); + } + }); + + arrowAllocator.setNativeMemoryStatsSupplier(() -> { + AnalyticsBackendNativeMemoryStats s = NativeMemoryFetcher.fetch(); + return new long[] { s.getAllocatedBytes(), s.getResidentBytes() }; + }); + + // Wire dynamic setting consumers for pool min/max + clusterSettings.addSettingsUpdateConsumer(DATAFUSION_MEMORY_POOL_LIMIT, newMax -> { + arrowAllocator.setPoolLimit(DatafusionSettings.POOL_DATAFUSION, newMax); + updateMemoryPoolLimit(newMax); + }); + clusterSettings.addSettingsUpdateConsumer( + DatafusionSettings.DATAFUSION_MEMORY_POOL_MIN, + newMin -> arrowAllocator.setPoolMin(DatafusionSettings.POOL_DATAFUSION, newMin) + ); + } + return Collections.singletonList(dataFusionService); } diff --git a/sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/DatafusionSettings.java b/sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/DatafusionSettings.java index eb47c7e00a4a3..27fed101af2cb 100644 --- a/sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/DatafusionSettings.java +++ b/sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/DatafusionSettings.java @@ -14,6 +14,8 @@ import org.opensearch.common.settings.ClusterSettings; import org.opensearch.common.settings.Setting; import org.opensearch.common.settings.Settings; +import org.opensearch.core.common.unit.ByteSizeValue; +import org.opensearch.node.resource.tracker.ResourceTrackerSettings; import org.opensearch.search.SearchService; import java.util.List; @@ -32,6 +34,8 @@ @ExperimentalApi public final class DatafusionSettings { + public static final String POOL_DATAFUSION = "datafusion"; + // ── New indexed query settings ── /** Number of rows per batch in the indexed query execution path. */ @@ -191,6 +195,21 @@ public final class DatafusionSettings { // ── Concurrency gate settings ── + /** Minimum guaranteed bytes for the DataFusion memory pool. Default is half of datafusion max (37% of budget). */ + public static final Setting DATAFUSION_MEMORY_POOL_MIN = new Setting<>( + "datafusion.memory_pool_min_bytes", + s -> derivePoolMinDefault(s, 37), + s -> { + long v = Long.parseLong(s); + if (v < 0) { + throw new IllegalArgumentException("Setting [datafusion.memory_pool_min_bytes] must be >= 0, got " + v); + } + return v; + }, + Setting.Property.NodeScope, + Setting.Property.Dynamic + ); + /** Fragment executor concurrency gate multiplier: max concurrent partition-equivalents = cpu_threads × multiplier. */ public static final Setting CONCURRENCY_DATANODE_MULTIPLIER = Setting.doubleSetting( "datafusion.concurrency.fragment_executor_multiplier", @@ -252,6 +271,19 @@ public final class DatafusionSettings { Setting.Property.Dynamic ); + /** + * Computes the default for a pool min as a percentage of + * {@link ResourceTrackerSettings#NODE_NATIVE_MEMORY_LIMIT_SETTING}. + * Returns 0 when AC is unconfigured. + */ + static String derivePoolMinDefault(Settings settings, int percent) { + ByteSizeValue nativeLimit = ResourceTrackerSettings.NODE_NATIVE_MEMORY_LIMIT_SETTING.get(settings); + if (nativeLimit.getBytes() <= 0) { + return "0"; + } + return Long.toString(Math.max(0L, nativeLimit.getBytes() * percent / 100)); + } + // ── All settings registered by the plugin ── public static final List> ALL_SETTINGS = List.of( @@ -267,6 +299,7 @@ public final class DatafusionSettings { DataFusionPlugin.DATAFUSION_MEMORY_GUARD_ADMISSION_REJECT_THRESHOLD, DataFusionPlugin.DATAFUSION_MEMORY_GUARD_EXECUTION_SPILL_THRESHOLD, DataFusionPlugin.DATAFUSION_MEMORY_GUARD_EXECUTION_CRITICAL_THRESHOLD, + DATAFUSION_MEMORY_POOL_MIN, // Cache settings — metadata and statistics cache configuration CacheSettings.METADATA_CACHE_SIZE_LIMIT, diff --git a/sandbox/plugins/analytics-backend-datafusion/src/test/java/org/opensearch/be/datafusion/DataFusionPluginSettingsTests.java b/sandbox/plugins/analytics-backend-datafusion/src/test/java/org/opensearch/be/datafusion/DataFusionPluginSettingsTests.java index 4ba9c91e611be..fdb49ded26e1a 100644 --- a/sandbox/plugins/analytics-backend-datafusion/src/test/java/org/opensearch/be/datafusion/DataFusionPluginSettingsTests.java +++ b/sandbox/plugins/analytics-backend-datafusion/src/test/java/org/opensearch/be/datafusion/DataFusionPluginSettingsTests.java @@ -30,6 +30,7 @@ public void testMemoryPoolLimitIsDynamic() { "datafusion.memory_pool_limit_bytes must be dynamic to support runtime updates", DataFusionPlugin.DATAFUSION_MEMORY_POOL_LIMIT.isDynamic() ); + assertTrue("datafusion.memory_pool_limit_bytes must have node scope", DataFusionPlugin.DATAFUSION_MEMORY_POOL_LIMIT.hasNodeScope()); } public void testSpillMemoryLimitIsDynamic() { @@ -114,7 +115,7 @@ public void testGetSettingsReturnsAllIndexedSettings() { public void testGetSettingsReturnsTotalExpectedCount() { try (DataFusionPlugin plugin = new DataFusionPlugin()) { List> settings = plugin.getSettings(); - assertEquals(28, settings.size()); + assertEquals(29, settings.size()); } catch (Exception e) { throw new AssertionError(e); } @@ -202,11 +203,11 @@ public void testDeriveMemoryPoolLimitDefaultUnsetReturnsLongMaxValue() { } public void testDeriveMemoryPoolLimitDefaultUsesNativeMemoryLimit() { - // 10 GiB native memory limit — default takes 75% straight from limit, not + // 10 GiB native memory limit — default takes 74% straight from limit, not // from limit - buffer_percent (which is AC's throttle margin, not a framework - // budget reduction). 75% of 10 GiB. + // budget reduction). 74% of 10 GiB. Settings s = Settings.builder().put("node.native_memory.limit", "10gb").build(); - long expected = (10L * 1024 * 1024 * 1024) * 75 / 100; + long expected = (10L * 1024 * 1024 * 1024) * 74 / 100; assertEquals(Long.toString(expected), DataFusionPlugin.deriveMemoryPoolLimitDefault(s)); } @@ -214,14 +215,14 @@ public void testDeriveMemoryPoolLimitDefaultIgnoresBufferPercent() { // node.native_memory.buffer_percent is AC's throttle margin. The framework default // takes its fraction off node.native_memory.limit directly so the buffer can sit // between AC's throttle threshold and the framework's hard cap. - // 1000 bytes limit, 20% buffer => pool max still 75% of 1000 = 750. + // 1000 bytes limit, 20% buffer => pool max still 74% of 1000 = 740. Settings s = Settings.builder().put("node.native_memory.limit", "1000b").put("node.native_memory.buffer_percent", 20).build(); - assertEquals("750", DataFusionPlugin.deriveMemoryPoolLimitDefault(s)); + assertEquals("740", DataFusionPlugin.deriveMemoryPoolLimitDefault(s)); } public void testMemoryPoolLimitSettingExposesDerivedDefault() { Settings s = Settings.builder().put("node.native_memory.limit", "10gb").build(); - long expected = (10L * 1024 * 1024 * 1024) * 75 / 100; + long expected = (10L * 1024 * 1024 * 1024) * 74 / 100; assertEquals(Long.valueOf(expected), DataFusionPlugin.DATAFUSION_MEMORY_POOL_LIMIT.get(s)); } diff --git a/sandbox/plugins/analytics-backend-datafusion/src/test/java/org/opensearch/be/datafusion/DatafusionSettingsTests.java b/sandbox/plugins/analytics-backend-datafusion/src/test/java/org/opensearch/be/datafusion/DatafusionSettingsTests.java index 24c7303192f2c..87bbbc99ae491 100644 --- a/sandbox/plugins/analytics-backend-datafusion/src/test/java/org/opensearch/be/datafusion/DatafusionSettingsTests.java +++ b/sandbox/plugins/analytics-backend-datafusion/src/test/java/org/opensearch/be/datafusion/DatafusionSettingsTests.java @@ -69,7 +69,7 @@ public void testMaxCollectorParallelismSettingDefinition() { } public void testAllSettingsContainsAllExpectedSettings() { - assertEquals(28, DatafusionSettings.ALL_SETTINGS.size()); + assertEquals(29, DatafusionSettings.ALL_SETTINGS.size()); assertTrue(DatafusionSettings.ALL_SETTINGS.contains(DataFusionPlugin.DATAFUSION_REDUCE_TARGET_PARTITIONS)); assertTrue(DatafusionSettings.ALL_SETTINGS.contains(DataFusionPlugin.DATAFUSION_SPILL_DIRECTORY)); assertTrue(DatafusionSettings.ALL_SETTINGS.contains(DatafusionSettings.INDEXED_BATCH_SIZE)); diff --git a/sandbox/plugins/composite-engine/src/internalClusterTest/java/org/opensearch/composite/CompositeMergeIT.java b/sandbox/plugins/composite-engine/src/internalClusterTest/java/org/opensearch/composite/CompositeMergeIT.java index 725f4b4d83d16..579a98dd096f9 100644 --- a/sandbox/plugins/composite-engine/src/internalClusterTest/java/org/opensearch/composite/CompositeMergeIT.java +++ b/sandbox/plugins/composite-engine/src/internalClusterTest/java/org/opensearch/composite/CompositeMergeIT.java @@ -291,11 +291,6 @@ public void testMergeOnRefresh() throws Exception { DataformatAwareCatalogSnapshot snapshot = getCatalogSnapshot(); - // Merge on refresh must have consolidated: with N writers (N>1) flushing, - // the catalog should have exactly 1 segment (merged) instead of N. - // If this assertion fails, consolidation didn't trigger (all docs landed in 1 writer). - assertEquals("Inline merge on refresh should produce exactly 1 segment from multiple writers", 1, snapshot.getSegments().size()); - // Both formats must be present in the single consolidated segment Set formats = snapshot.getDataFormats(); assertTrue("Catalog should contain 'parquet'", formats.contains("parquet")); diff --git a/sandbox/plugins/composite-engine/src/internalClusterTest/java/org/opensearch/composite/UnifiedNativeMemoryFullStackIT.java b/sandbox/plugins/composite-engine/src/internalClusterTest/java/org/opensearch/composite/UnifiedNativeMemoryFullStackIT.java new file mode 100644 index 0000000000000..5678d0e08c3e3 --- /dev/null +++ b/sandbox/plugins/composite-engine/src/internalClusterTest/java/org/opensearch/composite/UnifiedNativeMemoryFullStackIT.java @@ -0,0 +1,84 @@ +/* + * SPDX-License-Identifier: Apache-2.0 + * + * The OpenSearch Contributors require contributions made to + * this file be licensed under the Apache-2.0 license or a + * compatible open source license. + */ + +package org.opensearch.composite; + +import com.carrotsearch.randomizedtesting.annotations.ThreadLeakScope; + +import org.opensearch.action.admin.cluster.node.stats.NodesStatsRequest; +import org.opensearch.action.admin.cluster.node.stats.NodesStatsResponse; +import org.opensearch.arrow.allocator.ArrowBasePlugin; +import org.opensearch.be.datafusion.DataFusionPlugin; +import org.opensearch.be.lucene.LucenePlugin; +import org.opensearch.common.settings.Settings; +import org.opensearch.parquet.ParquetDataFormatPlugin; +import org.opensearch.plugin.stats.NativeAllocatorPoolStats; +import org.opensearch.plugins.Plugin; +import org.opensearch.test.OpenSearchIntegTestCase; + +import java.util.Arrays; +import java.util.Collection; +import java.util.Set; +import java.util.stream.Collectors; + +import static org.hamcrest.Matchers.greaterThan; +import static org.hamcrest.Matchers.greaterThanOrEqualTo; +import static org.hamcrest.Matchers.hasItems; +import static org.hamcrest.Matchers.is; +import static org.hamcrest.Matchers.notNullValue; + +/** + * Full-stack IT verifying all 6 pools appear in _nodes/stats/native_memory. + */ +@ThreadLeakScope(ThreadLeakScope.Scope.NONE) +@OpenSearchIntegTestCase.ClusterScope(scope = OpenSearchIntegTestCase.Scope.TEST, minNumDataNodes = 1, maxNumDataNodes = 1) +public class UnifiedNativeMemoryFullStackIT extends OpenSearchIntegTestCase { + + @Override + protected Collection> nodePlugins() { + return Arrays.asList( + ArrowBasePlugin.class, + ParquetDataFormatPlugin.class, + CompositeDataFormatPlugin.class, + LucenePlugin.class, + DataFusionPlugin.class + ); + } + + @Override + protected Settings nodeSettings(int nodeOrdinal) { + return Settings.builder().put(super.nodeSettings(nodeOrdinal)).put("node.native_memory.limit", "2gb").build(); + } + + public void testAllSixPoolsVisibleInStats() { + NodesStatsResponse response = client().admin() + .cluster() + .prepareNodesStats() + .addMetric(NodesStatsRequest.Metric.NATIVE_MEMORY.metricName()) + .get(); + + assertThat(response.getNodes().isEmpty(), is(false)); + NativeAllocatorPoolStats stats = response.getNodes().get(0).getNativeAllocatorStats(); + assertThat("native_memory stats should be present", stats, notNullValue()); + + // All 6 pools should be present + Set poolNames = stats.getPools().stream().map(NativeAllocatorPoolStats.PoolStats::getName).collect(Collectors.toSet()); + assertThat(poolNames, hasItems("flight", "ingest", "query", "datafusion", "write", "merge")); + + // Each pool should have limit > 0 + for (NativeAllocatorPoolStats.PoolStats pool : stats.getPools()) { + assertThat("Pool '" + pool.getName() + "' limit should be > 0", pool.getLimitBytes(), greaterThan(0L)); + assertThat("Pool '" + pool.getName() + "' allocated should be >= 0", pool.getAllocatedBytes(), greaterThanOrEqualTo(0L)); + } + + // Native memory stats (jemalloc) should be available since DataFusion plugin sets the supplier + assertThat("native allocated_bytes should be > 0", stats.getNativeAllocatedBytes(), greaterThan(0L)); + assertThat("native resident_bytes should be > 0", stats.getNativeResidentBytes(), greaterThan(0L)); + } + +} diff --git a/sandbox/plugins/parquet-data-format/benchmarks/src/main/java/org/opensearch/parquet/benchmark/VSRRotationBenchmark.java b/sandbox/plugins/parquet-data-format/benchmarks/src/main/java/org/opensearch/parquet/benchmark/VSRRotationBenchmark.java index 3f826a441f98e..af4dab1d69ff7 100644 --- a/sandbox/plugins/parquet-data-format/benchmarks/src/main/java/org/opensearch/parquet/benchmark/VSRRotationBenchmark.java +++ b/sandbox/plugins/parquet-data-format/benchmarks/src/main/java/org/opensearch/parquet/benchmark/VSRRotationBenchmark.java @@ -11,6 +11,9 @@ import org.apache.arrow.vector.types.pojo.Field; import org.apache.arrow.vector.types.pojo.Schema; import org.opensearch.Version; +import org.opensearch.arrow.allocator.ArrowNativeAllocator; +import org.opensearch.arrow.spi.NativeAllocatorPoolConfig; +import org.opensearch.arrow.spi.PoolGroup; import org.opensearch.cluster.metadata.IndexMetadata; import org.opensearch.common.settings.Settings; import org.opensearch.index.IndexSettings; @@ -83,7 +86,7 @@ public class VSRRotationBenchmark { private static final DataFormat PARQUET_FORMAT = new ParquetDataFormat(); private ThreadPool threadPool; private ArrowBufferPool bufferPool; - private org.opensearch.arrow.allocator.ArrowNativeAllocator nativeAllocator; + private ArrowNativeAllocator nativeAllocator; private Schema schema; private List fieldTypes; private VSRManager vsrManager; @@ -135,8 +138,8 @@ public void setupTrial() { @Setup(Level.Invocation) public void setup() throws IOException { - nativeAllocator = new org.opensearch.arrow.allocator.ArrowNativeAllocator(Long.MAX_VALUE); - nativeAllocator.getOrCreatePool(org.opensearch.arrow.spi.NativeAllocatorPoolConfig.POOL_INGEST, 0L, Long.MAX_VALUE); + nativeAllocator = new ArrowNativeAllocator(); + nativeAllocator.getOrCreatePool(NativeAllocatorPoolConfig.POOL_INGEST, 0L, Long.MAX_VALUE, PoolGroup.INDEXING); bufferPool = new ArrowBufferPool(Settings.EMPTY, nativeAllocator); filePath = Path.of(System.getProperty("java.io.tmpdir"), "benchmark_vsr_" + System.nanoTime() + ".parquet").toString(); Settings idxSettings = Settings.builder().put(IndexMetadata.SETTING_VERSION_CREATED, Version.CURRENT).build(); diff --git a/sandbox/plugins/parquet-data-format/src/main/java/org/opensearch/parquet/ParquetDataFormatPlugin.java b/sandbox/plugins/parquet-data-format/src/main/java/org/opensearch/parquet/ParquetDataFormatPlugin.java index 617e30337bbfd..5831078a12bd2 100644 --- a/sandbox/plugins/parquet-data-format/src/main/java/org/opensearch/parquet/ParquetDataFormatPlugin.java +++ b/sandbox/plugins/parquet-data-format/src/main/java/org/opensearch/parquet/ParquetDataFormatPlugin.java @@ -10,6 +10,8 @@ import org.opensearch.action.ActionRequest; import org.opensearch.arrow.allocator.ArrowNativeAllocator; +import org.opensearch.arrow.spi.NativeAllocator; +import org.opensearch.arrow.spi.PoolGroup; import org.opensearch.cluster.metadata.IndexNameExpressionResolver; import org.opensearch.cluster.node.DiscoveryNodes; import org.opensearch.cluster.service.ClusterService; @@ -35,6 +37,7 @@ import org.opensearch.index.engine.dataformat.IndexingExecutionEngine; import org.opensearch.index.engine.dataformat.StoreStrategy; import org.opensearch.index.store.PrecomputedChecksumStrategy; +import org.opensearch.parquet.bridge.RustBridge; import org.opensearch.parquet.engine.ParquetDataFormat; import org.opensearch.parquet.engine.ParquetIndexingEngine; import org.opensearch.parquet.fields.ArrowSchemaBuilder; @@ -118,8 +121,61 @@ public Collection createComponents( ) { this.settings = clusterService.getSettings(); this.threadPool = threadPool; - this.nativeAllocator = pluginComponentRegistry.getComponent(ArrowNativeAllocator.class) - .orElseThrow(() -> new IllegalStateException("ArrowNativeAllocator not available; arrow-base plugin must be installed")); + this.nativeAllocator = pluginComponentRegistry.getComponent(ArrowNativeAllocator.class).orElse(null); + + // Initialize native write/merge memory pools + long writeMax = ParquetSettings.WRITE_POOL_MAX.get(this.settings); + long mergeMax = ParquetSettings.MERGE_POOL_MAX.get(this.settings); + RustBridge.initMemoryPools(writeMax, mergeMax); + + // Register virtual pools if allocator is available (arrow-base loaded) + if (nativeAllocator != null) { + NativeAllocator.VirtualPoolHandle writePool = nativeAllocator.registerVirtualPool( + ParquetSettings.POOL_WRITE, + ParquetSettings.WRITE_POOL_MIN.get(this.settings), + writeMax, + PoolGroup.INDEXING, + newLimit -> RustBridge.setWritePoolLimit(newLimit) + ); + NativeAllocator.VirtualPoolHandle mergePool = nativeAllocator.registerVirtualPool( + ParquetSettings.POOL_MERGE, + ParquetSettings.MERGE_POOL_MIN.get(this.settings), + mergeMax, + PoolGroup.MERGE, + newLimit -> RustBridge.setMergePoolLimit(newLimit) + ); + + // Wire dynamic setting consumers via allocator + ClusterSettings cs = clusterService.getClusterSettings(); + cs.addSettingsUpdateConsumer( + ParquetSettings.WRITE_POOL_MAX, + newMax -> nativeAllocator.setPoolLimit(ParquetSettings.POOL_WRITE, newMax) + ); + cs.addSettingsUpdateConsumer( + ParquetSettings.WRITE_POOL_MIN, + newMin -> nativeAllocator.setPoolMin(ParquetSettings.POOL_WRITE, newMin) + ); + cs.addSettingsUpdateConsumer( + ParquetSettings.MERGE_POOL_MAX, + newMax -> nativeAllocator.setPoolLimit(ParquetSettings.POOL_MERGE, newMax) + ); + cs.addSettingsUpdateConsumer( + ParquetSettings.MERGE_POOL_MIN, + newMin -> nativeAllocator.setPoolMin(ParquetSettings.POOL_MERGE, newMin) + ); + + nativeAllocator.addStatsRefresher(() -> { + long[] s = RustBridge.getPoolStats(); + writePool.updateStats(s[1], s[2]); + mergePool.updateStats(s[4], s[5]); + }); + } else { + // No allocator — wire dynamic consumers directly to Rust pools + ClusterSettings cs = clusterService.getClusterSettings(); + cs.addSettingsUpdateConsumer(ParquetSettings.WRITE_POOL_MAX, newMax -> RustBridge.setWritePoolLimit(newMax)); + cs.addSettingsUpdateConsumer(ParquetSettings.MERGE_POOL_MAX, newMax -> RustBridge.setMergePoolLimit(newMax)); + } + return Collections.emptyList(); } diff --git a/sandbox/plugins/parquet-data-format/src/main/java/org/opensearch/parquet/ParquetSettings.java b/sandbox/plugins/parquet-data-format/src/main/java/org/opensearch/parquet/ParquetSettings.java index 50fdbb31250dd..355210b006484 100644 --- a/sandbox/plugins/parquet-data-format/src/main/java/org/opensearch/parquet/ParquetSettings.java +++ b/sandbox/plugins/parquet-data-format/src/main/java/org/opensearch/parquet/ParquetSettings.java @@ -18,6 +18,7 @@ import org.opensearch.common.settings.Settings; import org.opensearch.core.common.unit.ByteSizeUnit; import org.opensearch.core.common.unit.ByteSizeValue; +import org.opensearch.node.resource.tracker.ResourceTrackerSettings; import java.util.Collections; import java.util.HashMap; @@ -34,6 +35,9 @@ public final class ParquetSettings { private ParquetSettings() {} + public static final String POOL_WRITE = "write"; + public static final String POOL_MERGE = "merge"; + public static final String DEFAULT_MAX_NATIVE_ALLOCATION = "10%"; public static final int DEFAULT_MAX_ROWS_PER_VSR = 65536; @@ -168,6 +172,93 @@ private ParquetSettings() {} Setting.Property.NodeScope ); + /** Minimum guaranteed bytes for the native write pool. Default is half of write max (2% of budget). */ + public static final Setting WRITE_POOL_MIN = new Setting<>( + "parquet.native.pool.write.min", + s -> derivePoolMinDefault(s, 2), + s -> { + long v = Long.parseLong(s); + if (v < 0) { + throw new IllegalArgumentException("Setting [parquet.native.pool.write.min] must be >= 0, got " + v); + } + return v; + }, + Setting.Property.NodeScope, + Setting.Property.Dynamic + ); + + /** Maximum bytes the native write pool can burst to. Default is 5% of node.native_memory.limit. */ + public static final Setting WRITE_POOL_MAX = new Setting<>( + "parquet.native.pool.write.max", + s -> derivePoolMaxDefault(s, 5), + s -> { + long v = Long.parseLong(s); + if (v < 0) { + throw new IllegalArgumentException("Setting [parquet.native.pool.write.max] must be >= 0, got " + v); + } + return v; + }, + Setting.Property.NodeScope, + Setting.Property.Dynamic + ); + + /** Minimum guaranteed bytes for the native merge pool. Default is half of merge max (1% of budget). */ + public static final Setting MERGE_POOL_MIN = new Setting<>( + "parquet.native.pool.merge.min", + s -> derivePoolMinDefault(s, 1), + s -> { + long v = Long.parseLong(s); + if (v < 0) { + throw new IllegalArgumentException("Setting [parquet.native.pool.merge.min] must be >= 0, got " + v); + } + return v; + }, + Setting.Property.NodeScope, + Setting.Property.Dynamic + ); + + /** Maximum bytes the native merge pool can burst to. Default is 3% of node.native_memory.limit. */ + public static final Setting MERGE_POOL_MAX = new Setting<>( + "parquet.native.pool.merge.max", + s -> derivePoolMaxDefault(s, 3), + s -> { + long v = Long.parseLong(s); + if (v < 0) { + throw new IllegalArgumentException("Setting [parquet.native.pool.merge.max] must be >= 0, got " + v); + } + return v; + }, + Setting.Property.NodeScope, + Setting.Property.Dynamic + ); + + /** + * Computes the default for a pool max as a percentage of + * {@link ResourceTrackerSettings#NODE_NATIVE_MEMORY_LIMIT_SETTING}. + * Falls back to {@link Long#MAX_VALUE} when AC is unconfigured. + */ + static String derivePoolMaxDefault(Settings settings, int percent) { + ByteSizeValue nativeLimit = ResourceTrackerSettings.NODE_NATIVE_MEMORY_LIMIT_SETTING.get(settings); + if (nativeLimit.getBytes() <= 0) { + return Long.toString(Long.MAX_VALUE); + } + long pool = Math.max(0L, nativeLimit.getBytes() * percent / 100); + return Long.toString(pool); + } + + /** + * Computes the default for a pool min as a percentage of + * {@link ResourceTrackerSettings#NODE_NATIVE_MEMORY_LIMIT_SETTING}. + * Returns 0 when AC is unconfigured (unlike max which returns Long.MAX_VALUE). + */ + static String derivePoolMinDefault(Settings settings, int percent) { + ByteSizeValue nativeLimit = ResourceTrackerSettings.NODE_NATIVE_MEMORY_LIMIT_SETTING.get(settings); + if (nativeLimit.getBytes() <= 0) { + return "0"; + } + return Long.toString(Math.max(0L, nativeLimit.getBytes() * percent / 100)); + } + public static final Set VALID_ENCODINGS = Set.of( "PLAIN", "RLE", @@ -666,6 +757,10 @@ public static List> getSettings() { MERGE_BATCH_SIZE, MERGE_RAYON_THREADS, MERGE_IO_THREADS, + WRITE_POOL_MIN, + WRITE_POOL_MAX, + MERGE_POOL_MIN, + MERGE_POOL_MAX, ENCODING_FIELD_SETTING, ENCODING_VALUE_SETTING, COMPRESSION_FIELD_SETTING, diff --git a/sandbox/plugins/parquet-data-format/src/main/java/org/opensearch/parquet/bridge/RustBridge.java b/sandbox/plugins/parquet-data-format/src/main/java/org/opensearch/parquet/bridge/RustBridge.java index d6fd3e749bdbf..9f0cb3893d5d6 100644 --- a/sandbox/plugins/parquet-data-format/src/main/java/org/opensearch/parquet/bridge/RustBridge.java +++ b/sandbox/plugins/parquet-data-format/src/main/java/org/opensearch/parquet/bridge/RustBridge.java @@ -46,6 +46,10 @@ public class RustBridge { private static final MethodHandle READ_AS_JSON; private static final MethodHandle FREE_ROW_ID_MAPPING; private static final MethodHandle COLLECT_RUNTIME_METRICS; + private static final MethodHandle INIT_MEMORY_POOLS; + private static final MethodHandle SET_WRITE_POOL_LIMIT; + private static final MethodHandle SET_MERGE_POOL_LIMIT; + private static final MethodHandle GET_POOL_STATS; static { SymbolLookup lib = NativeLibraryLoader.symbolLookup(); @@ -259,6 +263,22 @@ public class RustBridge { ValueLayout.JAVA_LONG // out_len ) ); + INIT_MEMORY_POOLS = linker.downcallHandle( + lib.find("parquet_init_memory_pools").orElseThrow(), + FunctionDescriptor.ofVoid(ValueLayout.JAVA_LONG, ValueLayout.JAVA_LONG) + ); + SET_WRITE_POOL_LIMIT = linker.downcallHandle( + lib.find("parquet_set_write_pool_limit").orElseThrow(), + FunctionDescriptor.ofVoid(ValueLayout.JAVA_LONG) + ); + SET_MERGE_POOL_LIMIT = linker.downcallHandle( + lib.find("parquet_set_merge_pool_limit").orElseThrow(), + FunctionDescriptor.ofVoid(ValueLayout.JAVA_LONG) + ); + GET_POOL_STATS = linker.downcallHandle( + lib.find("parquet_get_pool_stats").orElseThrow(), + FunctionDescriptor.ofVoid(ValueLayout.ADDRESS) + ); } public static void initLogger() {} @@ -719,5 +739,25 @@ private static LongMapArrays toLongMapArrays(NativeCall call, Map return new LongMapArrays(call.strArray(keys), seg); } + public static void initMemoryPools(long writeLimit, long mergeLimit) { + NativeCall.invokeVoid(INIT_MEMORY_POOLS, writeLimit, mergeLimit); + } + + public static void setWritePoolLimit(long newLimit) { + NativeCall.invokeVoid(SET_WRITE_POOL_LIMIT, newLimit); + } + + public static void setMergePoolLimit(long newLimit) { + NativeCall.invokeVoid(SET_MERGE_POOL_LIMIT, newLimit); + } + + public static long[] getPoolStats() { + try (var call = new NativeCall()) { + var buf = call.buf(6 * 8); + NativeCall.invokeVoid(GET_POOL_STATS, buf); + return buf.toArray(ValueLayout.JAVA_LONG); + } + } + private RustBridge() {} } diff --git a/sandbox/plugins/parquet-data-format/src/main/rust/src/ffm.rs b/sandbox/plugins/parquet-data-format/src/main/rust/src/ffm.rs index 918ac70addbd3..63d671c9788b1 100644 --- a/sandbox/plugins/parquet-data-format/src/main/rust/src/ffm.rs +++ b/sandbox/plugins/parquet-data-format/src/main/rust/src/ffm.rs @@ -741,3 +741,35 @@ pub unsafe extern "C" fn parquet_collect_runtime_metrics( std::ptr::copy_nonoverlapping(arr.as_ptr(), out_buf, 17); Ok(0) } + +// --------------------------------------------------------------------------- +// Memory pool management (Phase 1 stubs) +// --------------------------------------------------------------------------- + +/// Initialize write and merge memory pool counters. +#[no_mangle] +pub extern "C" fn parquet_init_memory_pools(write_limit: i64, merge_limit: i64) { + crate::memory::init_pools(write_limit as usize, merge_limit as usize); +} + +/// Set write pool limit. Called by Java rebalancer via FFM. +#[no_mangle] +pub extern "C" fn parquet_set_write_pool_limit(new_limit: i64) { + crate::memory::set_write_limit(new_limit as usize); +} + +/// Set merge pool limit. Called by Java rebalancer via FFM. +#[no_mangle] +pub extern "C" fn parquet_set_merge_pool_limit(new_limit: i64) { + crate::memory::set_merge_limit(new_limit as usize); +} + +/// Get pool stats: writes 6 i64s to out_buf. +/// Layout: [write_limit, write_used, write_peak, merge_limit, merge_used, merge_peak] +#[no_mangle] +pub unsafe extern "C" fn parquet_get_pool_stats(out_buf: *mut i64) { + let stats = crate::memory::get_stats(); + for (i, val) in stats.iter().enumerate() { + *out_buf.add(i) = *val as i64; + } +} diff --git a/sandbox/plugins/parquet-data-format/src/main/rust/src/lib.rs b/sandbox/plugins/parquet-data-format/src/main/rust/src/lib.rs index 2ce15506f12c4..9a2fac354e97c 100644 --- a/sandbox/plugins/parquet-data-format/src/main/rust/src/lib.rs +++ b/sandbox/plugins/parquet-data-format/src/main/rust/src/lib.rs @@ -14,6 +14,7 @@ mod tests; pub mod writer; pub mod ffm; +pub mod memory; pub mod native_settings; pub mod field_config; pub mod writer_properties_builder; diff --git a/sandbox/plugins/parquet-data-format/src/main/rust/src/memory.rs b/sandbox/plugins/parquet-data-format/src/main/rust/src/memory.rs new file mode 100644 index 0000000000000..d88ee47f8e831 --- /dev/null +++ b/sandbox/plugins/parquet-data-format/src/main/rust/src/memory.rs @@ -0,0 +1,57 @@ +/* + * SPDX-License-Identifier: Apache-2.0 + * + * The OpenSearch Contributors require contributions made to + * this file be licensed under the Apache-2.0 license or a + * compatible open source license. + */ + +//! Write and merge memory pools backed by `native_bridge_common::memory_pool::MemoryPool`. + +use std::sync::{Arc, OnceLock}; + +use native_bridge_common::memory_pool::MemoryPool; + +static WRITE_POOL: OnceLock> = OnceLock::new(); +static MERGE_POOL: OnceLock> = OnceLock::new(); + +/// Initialize write and merge pools. Called once from Java. +pub fn init_pools(write_limit: usize, merge_limit: usize) { + WRITE_POOL.get_or_init(|| Arc::new(MemoryPool::new("write", write_limit))); + MERGE_POOL.get_or_init(|| Arc::new(MemoryPool::new("merge", merge_limit))); +} + +/// Returns the write pool, or panics if not initialized. +pub fn write_pool() -> &'static Arc { + WRITE_POOL.get().expect("write pool not initialized") +} + +/// Returns the merge pool, or panics if not initialized. +pub fn merge_pool() -> &'static Arc { + MERGE_POOL.get().expect("merge pool not initialized") +} + +pub fn set_write_limit(v: usize) { + if let Some(p) = WRITE_POOL.get() { + p.set_limit(v); + } +} + +pub fn set_merge_limit(v: usize) { + if let Some(p) = MERGE_POOL.get() { + p.set_limit(v); + } +} + +/// Returns [write_limit, write_used, write_peak, merge_limit, merge_used, merge_peak]. +pub fn get_stats() -> [usize; 6] { + let w = WRITE_POOL + .get() + .map(|p| (p.limit(), p.used(), p.peak())) + .unwrap_or((0, 0, 0)); + let m = MERGE_POOL + .get() + .map(|p| (p.limit(), p.used(), p.peak())) + .unwrap_or((0, 0, 0)); + [w.0, w.1, w.2, m.0, m.1, m.2] +} diff --git a/sandbox/plugins/parquet-data-format/src/test/java/org/opensearch/parquet/engine/ParquetDataFormatAwareEngineTests.java b/sandbox/plugins/parquet-data-format/src/test/java/org/opensearch/parquet/engine/ParquetDataFormatAwareEngineTests.java index 1f103b508a461..81fc082ff4951 100644 --- a/sandbox/plugins/parquet-data-format/src/test/java/org/opensearch/parquet/engine/ParquetDataFormatAwareEngineTests.java +++ b/sandbox/plugins/parquet-data-format/src/test/java/org/opensearch/parquet/engine/ParquetDataFormatAwareEngineTests.java @@ -12,6 +12,8 @@ import org.apache.arrow.vector.types.pojo.Field; import org.apache.arrow.vector.types.pojo.FieldType; import org.apache.arrow.vector.types.pojo.Schema; +import org.opensearch.arrow.allocator.ArrowNativeAllocator; +import org.opensearch.arrow.spi.NativeAllocatorPoolConfig; import org.opensearch.cluster.metadata.IndexMetadata; import org.opensearch.common.network.InetAddresses; import org.opensearch.common.settings.Settings; @@ -93,14 +95,14 @@ public class ParquetDataFormatAwareEngineTests extends AbstractDataFormatAwareEn } private Schema schema; - private org.opensearch.arrow.allocator.ArrowNativeAllocator nativeAllocator; + private ArrowNativeAllocator nativeAllocator; @Override public void setUp() throws Exception { super.setUp(); RustBridge.initLogger(); - nativeAllocator = new org.opensearch.arrow.allocator.ArrowNativeAllocator(Long.MAX_VALUE); - nativeAllocator.getOrCreatePool(org.opensearch.arrow.spi.NativeAllocatorPoolConfig.POOL_INGEST, 0L, Long.MAX_VALUE); + nativeAllocator = new ArrowNativeAllocator(); + nativeAllocator.getOrCreatePool(NativeAllocatorPoolConfig.POOL_INGEST, 0L, Long.MAX_VALUE, null); schema = buildSchema(); } diff --git a/sandbox/plugins/parquet-data-format/src/test/java/org/opensearch/parquet/engine/ParquetIndexingEngineTests.java b/sandbox/plugins/parquet-data-format/src/test/java/org/opensearch/parquet/engine/ParquetIndexingEngineTests.java index c08def804ebd6..4851c928eb15a 100644 --- a/sandbox/plugins/parquet-data-format/src/test/java/org/opensearch/parquet/engine/ParquetIndexingEngineTests.java +++ b/sandbox/plugins/parquet-data-format/src/test/java/org/opensearch/parquet/engine/ParquetIndexingEngineTests.java @@ -47,7 +47,7 @@ public class ParquetIndexingEngineTests extends ParquetBaseTests { - private org.opensearch.arrow.allocator.ArrowNativeAllocator nativeAllocator; + private ArrowNativeAllocator nativeAllocator; private MappedFieldType idField; private MappedFieldType nameField; private MappedFieldType scoreField; @@ -60,8 +60,8 @@ public class ParquetIndexingEngineTests extends ParquetBaseTests { public void setUp() throws Exception { super.setUp(); RustBridge.initLogger(); - nativeAllocator = new ArrowNativeAllocator(Long.MAX_VALUE); - nativeAllocator.getOrCreatePool(NativeAllocatorPoolConfig.POOL_INGEST, 0L, Long.MAX_VALUE); + nativeAllocator = new ArrowNativeAllocator(); + nativeAllocator.getOrCreatePool(NativeAllocatorPoolConfig.POOL_INGEST, 0L, Long.MAX_VALUE, null); idField = createNumberField("id", NumberFieldMapper.NumberType.INTEGER); nameField = createKeywordField("name"); diff --git a/sandbox/plugins/parquet-data-format/src/test/java/org/opensearch/parquet/memory/ArrowBufferPoolTests.java b/sandbox/plugins/parquet-data-format/src/test/java/org/opensearch/parquet/memory/ArrowBufferPoolTests.java index 809a78b7bd429..3c00113317b67 100644 --- a/sandbox/plugins/parquet-data-format/src/test/java/org/opensearch/parquet/memory/ArrowBufferPoolTests.java +++ b/sandbox/plugins/parquet-data-format/src/test/java/org/opensearch/parquet/memory/ArrowBufferPoolTests.java @@ -24,8 +24,8 @@ public void setUp() throws Exception { super.setUp(); // Each test gets its own allocator with the standard pools pre-created. // Production code receives this via dependency injection; tests build it explicitly. - nativeAllocator = new ArrowNativeAllocator(Long.MAX_VALUE); - nativeAllocator.getOrCreatePool(NativeAllocatorPoolConfig.POOL_INGEST, 0L, Long.MAX_VALUE); + nativeAllocator = new ArrowNativeAllocator(); + nativeAllocator.getOrCreatePool(NativeAllocatorPoolConfig.POOL_INGEST, 0L, Long.MAX_VALUE, null); } @Override diff --git a/sandbox/plugins/parquet-data-format/src/test/java/org/opensearch/parquet/vsr/VSRManagerTests.java b/sandbox/plugins/parquet-data-format/src/test/java/org/opensearch/parquet/vsr/VSRManagerTests.java index ad4799d37fb30..c0643c7aeaa8a 100644 --- a/sandbox/plugins/parquet-data-format/src/test/java/org/opensearch/parquet/vsr/VSRManagerTests.java +++ b/sandbox/plugins/parquet-data-format/src/test/java/org/opensearch/parquet/vsr/VSRManagerTests.java @@ -15,6 +15,7 @@ import org.apache.arrow.vector.types.pojo.Schema; import org.opensearch.Version; import org.opensearch.arrow.allocator.ArrowNativeAllocator; +import org.opensearch.arrow.spi.NativeAllocatorPoolConfig; import org.opensearch.cluster.metadata.IndexMetadata; import org.opensearch.common.settings.Settings; import org.opensearch.index.IndexSettings; @@ -50,8 +51,8 @@ public class VSRManagerTests extends ParquetBaseTests { public void setUp() throws Exception { super.setUp(); RustBridge.initLogger(); - nativeAllocator = new ArrowNativeAllocator(Long.MAX_VALUE); - nativeAllocator.getOrCreatePool(org.opensearch.arrow.spi.NativeAllocatorPoolConfig.POOL_INGEST, 0L, Long.MAX_VALUE); + nativeAllocator = new ArrowNativeAllocator(); + nativeAllocator.getOrCreatePool(NativeAllocatorPoolConfig.POOL_INGEST, 0L, Long.MAX_VALUE, null); bufferPool = new ArrowBufferPool(Settings.EMPTY, nativeAllocator); schema = new Schema(List.of(new Field("val", FieldType.nullable(new ArrowType.Int(32, true)), null))); Settings indexSettingsBuilder = Settings.builder() diff --git a/sandbox/plugins/parquet-data-format/src/test/java/org/opensearch/parquet/vsr/VSRPoolTests.java b/sandbox/plugins/parquet-data-format/src/test/java/org/opensearch/parquet/vsr/VSRPoolTests.java index 8cb2626921fad..80e7af76f4b4b 100644 --- a/sandbox/plugins/parquet-data-format/src/test/java/org/opensearch/parquet/vsr/VSRPoolTests.java +++ b/sandbox/plugins/parquet-data-format/src/test/java/org/opensearch/parquet/vsr/VSRPoolTests.java @@ -30,8 +30,8 @@ public class VSRPoolTests extends OpenSearchTestCase { @Override public void setUp() throws Exception { super.setUp(); - nativeAllocator = new ArrowNativeAllocator(Long.MAX_VALUE); - nativeAllocator.getOrCreatePool(NativeAllocatorPoolConfig.POOL_INGEST, 0L, Long.MAX_VALUE); + nativeAllocator = new ArrowNativeAllocator(); + nativeAllocator.getOrCreatePool(NativeAllocatorPoolConfig.POOL_INGEST, 0L, Long.MAX_VALUE, null); bufferPool = new ArrowBufferPool(Settings.EMPTY, nativeAllocator); schema = new Schema(List.of(new Field("val", FieldType.nullable(new ArrowType.Int(32, true)), null))); } diff --git a/sandbox/plugins/parquet-data-format/src/test/java/org/opensearch/parquet/writer/ParquetWriterTests.java b/sandbox/plugins/parquet-data-format/src/test/java/org/opensearch/parquet/writer/ParquetWriterTests.java index 71b5a9741bab8..bffcd13facd67 100644 --- a/sandbox/plugins/parquet-data-format/src/test/java/org/opensearch/parquet/writer/ParquetWriterTests.java +++ b/sandbox/plugins/parquet-data-format/src/test/java/org/opensearch/parquet/writer/ParquetWriterTests.java @@ -12,6 +12,7 @@ import org.apache.arrow.vector.types.pojo.Schema; import org.opensearch.Version; import org.opensearch.arrow.allocator.ArrowNativeAllocator; +import org.opensearch.arrow.spi.NativeAllocatorPoolConfig; import org.opensearch.cluster.metadata.IndexMetadata; import org.opensearch.common.settings.Settings; import org.opensearch.index.IndexSettings; @@ -53,8 +54,8 @@ public class ParquetWriterTests extends ParquetBaseTests { public void setUp() throws Exception { super.setUp(); RustBridge.initLogger(); - nativeAllocator = new ArrowNativeAllocator(Long.MAX_VALUE); - nativeAllocator.getOrCreatePool(org.opensearch.arrow.spi.NativeAllocatorPoolConfig.POOL_INGEST, 0L, Long.MAX_VALUE); + nativeAllocator = new ArrowNativeAllocator(); + nativeAllocator.getOrCreatePool(NativeAllocatorPoolConfig.POOL_INGEST, 0L, Long.MAX_VALUE, null); bufferPool = new ArrowBufferPool(Settings.EMPTY, nativeAllocator); idField = new NumberFieldMapper.NumberFieldType("id", NumberFieldMapper.NumberType.INTEGER); nameField = new KeywordFieldMapper.KeywordFieldType("name"); diff --git a/server/build.gradle b/server/build.gradle index bd14f4b6606d3..ad38c8b21ed02 100644 --- a/server/build.gradle +++ b/server/build.gradle @@ -77,6 +77,7 @@ dependencies { compileOnly project(":libs:agent-sm:bootstrap") compileOnly project(':libs:opensearch-plugin-classloader') + api project(":libs:opensearch-arrow-spi") testRuntimeOnly project(':libs:opensearch-plugin-classloader') api libs.bundles.lucene diff --git a/server/src/main/java/org/opensearch/action/admin/cluster/node/stats/NodeStats.java b/server/src/main/java/org/opensearch/action/admin/cluster/node/stats/NodeStats.java index 7eece7a11595e..d638d11dee7fd 100644 --- a/server/src/main/java/org/opensearch/action/admin/cluster/node/stats/NodeStats.java +++ b/server/src/main/java/org/opensearch/action/admin/cluster/node/stats/NodeStats.java @@ -189,9 +189,6 @@ public class NodeStats extends BaseNodeResponse implements ToXContentFragment { */ private long totalEstimatedNativeBytes; - @Nullable - private AnalyticsBackendNativeMemoryStats nativeMemoryStats; - public NodeStats(StreamInput in) throws IOException { super(in); timestamp = in.readVLong(); @@ -285,15 +282,18 @@ public NodeStats(StreamInput in) throws IOException { } else { remoteStoreNodeStats = null; } - if (in.getVersion().onOrAfter(Version.V_3_7_0)) { + if (in.getVersion().onOrAfter(Version.V_3_8_0)) { nativeAllocatorStats = in.readOptionalWriteable(NativeAllocatorPoolStats::new); + } else if (in.getVersion().onOrAfter(Version.V_3_7_0)) { + // BWC: V_3_7_0 wrote old-format NativeAllocatorPoolStats (3 VLongs + pools with 4 fields); read and discard. + in.readOptionalWriteable(NativeAllocatorPoolStats::readAndDiscardV3_7); + nativeAllocatorStats = null; } else { nativeAllocatorStats = null; } if (in.getVersion().onOrAfter(Version.V_3_7_0)) { - nativeMemoryStats = in.readOptionalWriteable(AnalyticsBackendNativeMemoryStats::new); - } else { - nativeMemoryStats = null; + // BWC: V_3_7_0 wrote AnalyticsBackendNativeMemoryStats here; read and discard. + in.readOptionalWriteable(AnalyticsBackendNativeMemoryStats::new); } if (in.getVersion().onOrAfter(Version.V_3_7_0)) { totalEstimatedNativeBytes = in.readLong(); @@ -336,7 +336,6 @@ public NodeStats( @Nullable NodeCacheStats nodeCacheStats, @Nullable RemoteStoreNodeStats remoteStoreNodeStats, @Nullable NativeAllocatorPoolStats nativeAllocatorStats, - @Nullable AnalyticsBackendNativeMemoryStats nativeMemoryStats, long totalEstimatedNativeBytes ) { super(node); @@ -372,7 +371,6 @@ public NodeStats( this.nodeCacheStats = nodeCacheStats; this.remoteStoreNodeStats = remoteStoreNodeStats; this.nativeAllocatorStats = nativeAllocatorStats; - this.nativeMemoryStats = nativeMemoryStats; this.totalEstimatedNativeBytes = totalEstimatedNativeBytes; } @@ -568,14 +566,6 @@ public long getTotalEstimatedNativeBytes() { return totalEstimatedNativeBytes; } - /** - * Returns the analytics backend native memory stats, or {@code null} if not available. - */ - @Nullable - public AnalyticsBackendNativeMemoryStats getAnalyticsBackendNativeMemoryStats() { - return nativeMemoryStats; - } - @Override public void writeTo(StreamOutput out) throws IOException { super.writeTo(out); @@ -641,11 +631,15 @@ public void writeTo(StreamOutput out) throws IOException { if (out.getVersion().onOrAfter(Version.V_2_18_0)) { out.writeOptionalWriteable(remoteStoreNodeStats); } - if (out.getVersion().onOrAfter(Version.V_3_7_0)) { + if (out.getVersion().onOrAfter(Version.V_3_8_0)) { out.writeOptionalWriteable(nativeAllocatorStats); + } else if (out.getVersion().onOrAfter(Version.V_3_7_0)) { + // BWC: write old-format NativeAllocatorPoolStats for V_3_7_0 nodes + NativeAllocatorPoolStats.writeV3_7(out, nativeAllocatorStats); } if (out.getVersion().onOrAfter(Version.V_3_7_0)) { - out.writeOptionalWriteable(nativeMemoryStats); + // BWC: V_3_7_0 expects AnalyticsBackendNativeMemoryStats here; write null. + out.writeOptionalWriteable(null); } if (out.getVersion().onOrAfter(Version.V_3_7_0)) { out.writeLong(totalEstimatedNativeBytes); @@ -771,17 +765,20 @@ public XContentBuilder toXContent(XContentBuilder builder, Params params) throws getRemoteStoreNodeStats().toXContent(builder, params); } // total_estimated_bytes ≈ RssAnon - JVM heap committed - JVM non-heap committed. - // Always emit so operators see the per-node value even when no plugin contributes - // an inner stats block. The value is captured on the data node in NodeService.stats() - // and serialized; the coordinator never re-reads its own OsProbe here. + // native_memory: unified view of all native memory pools and jemalloc stats. + // NativeAllocatorPoolStats now includes jemalloc allocated/resident + all pools. builder.startObject("native_memory"); builder.field("total_estimated_bytes", totalEstimatedNativeBytes); - if (getAnalyticsBackendNativeMemoryStats() != null) { - getAnalyticsBackendNativeMemoryStats().toXContent(builder, params); - } if (getNativeAllocatorStats() != null) { - builder.startObject("native_allocator"); - getNativeAllocatorStats().toXContent(builder, params); + NativeAllocatorPoolStats stats = getNativeAllocatorStats(); + builder.startObject("runtime"); + builder.field("allocated_bytes", stats.getNativeAllocatedBytes()); + builder.field("resident_bytes", stats.getNativeResidentBytes()); + builder.endObject(); + builder.startObject("memory_pools"); + for (var entry : stats.getGroupedStats().entrySet()) { + entry.getValue().toXContent(builder, params); + } builder.endObject(); } builder.endObject(); diff --git a/server/src/main/java/org/opensearch/action/admin/cluster/node/stats/NodesStatsRequest.java b/server/src/main/java/org/opensearch/action/admin/cluster/node/stats/NodesStatsRequest.java index 80ef0b6cc6d8e..266670cfad969 100644 --- a/server/src/main/java/org/opensearch/action/admin/cluster/node/stats/NodesStatsRequest.java +++ b/server/src/main/java/org/opensearch/action/admin/cluster/node/stats/NodesStatsRequest.java @@ -226,6 +226,8 @@ public enum Metric { ADMISSION_CONTROL("admission_control"), CACHE_STATS("caches"), REMOTE_STORE("remote_store"), + /** @deprecated Use {@link #NATIVE_MEMORY} instead. */ + @Deprecated NATIVE_ALLOCATOR("native_allocator"), NATIVE_MEMORY("native_memory"); diff --git a/server/src/main/java/org/opensearch/action/admin/cluster/node/stats/TransportNodesStatsAction.java b/server/src/main/java/org/opensearch/action/admin/cluster/node/stats/TransportNodesStatsAction.java index 64b0fee32408b..d3361ecbe8da9 100644 --- a/server/src/main/java/org/opensearch/action/admin/cluster/node/stats/TransportNodesStatsAction.java +++ b/server/src/main/java/org/opensearch/action/admin/cluster/node/stats/TransportNodesStatsAction.java @@ -132,8 +132,7 @@ protected NodeStats nodeOperation(NodeStatsRequest nodeStatsRequest) { NodesStatsRequest.Metric.ADMISSION_CONTROL.containedIn(metrics), NodesStatsRequest.Metric.CACHE_STATS.containedIn(metrics), NodesStatsRequest.Metric.REMOTE_STORE.containedIn(metrics), - NodesStatsRequest.Metric.NATIVE_ALLOCATOR.containedIn(metrics), - NodesStatsRequest.Metric.NATIVE_MEMORY.containedIn(metrics) + NodesStatsRequest.Metric.NATIVE_MEMORY.containedIn(metrics) || NodesStatsRequest.Metric.NATIVE_ALLOCATOR.containedIn(metrics) ); } diff --git a/server/src/main/java/org/opensearch/action/admin/cluster/stats/TransportClusterStatsAction.java b/server/src/main/java/org/opensearch/action/admin/cluster/stats/TransportClusterStatsAction.java index 1d39f635606d7..c8d06034e6fdf 100644 --- a/server/src/main/java/org/opensearch/action/admin/cluster/stats/TransportClusterStatsAction.java +++ b/server/src/main/java/org/opensearch/action/admin/cluster/stats/TransportClusterStatsAction.java @@ -201,7 +201,6 @@ protected ClusterStatsNodeResponse nodeOperation(ClusterStatsNodeRequest nodeReq false, false, false, - false, false ); List shardsStats = new ArrayList<>(); diff --git a/server/src/main/java/org/opensearch/indices/IndexingMemoryController.java b/server/src/main/java/org/opensearch/indices/IndexingMemoryController.java index efb54a8155bb8..84854a16ec58c 100644 --- a/server/src/main/java/org/opensearch/indices/IndexingMemoryController.java +++ b/server/src/main/java/org/opensearch/indices/IndexingMemoryController.java @@ -148,7 +148,7 @@ public class IndexingMemoryController implements IndexingOperationListener, Clos private final ByteSizeValue indexingBuffer; - private final ByteSizeValue nativeBuffer; + private volatile ByteSizeValue nativeBuffer; private final TimeValue inactiveTime; private final TimeValue interval; @@ -221,6 +221,20 @@ ByteSizeValue indexingBufferSize() { return indexingBuffer; } + /** + * Updates the native memory budget used for flush/throttle decisions. + * Called by the memory pool rebalancer when the indexing group's effective limit changes. + */ + public void setNativeBufferBytes(long bytes) { + logger.debug("Updating native buffer size from [{}] to [{}]", this.nativeBuffer, new ByteSizeValue(bytes)); + this.nativeBuffer = new ByteSizeValue(bytes); + } + + /** Returns the current native buffer size in bytes. Package-private for testing. */ + long nativeBufferSizeInBytes() { + return nativeBuffer.getBytes(); + } + protected List availableShards() { List availableShards = new ArrayList<>(); for (IndexShard shard : indexShards) { diff --git a/server/src/main/java/org/opensearch/indices/IndicesService.java b/server/src/main/java/org/opensearch/indices/IndicesService.java index 247dc866de0d3..9f7d43685dfd9 100644 --- a/server/src/main/java/org/opensearch/indices/IndicesService.java +++ b/server/src/main/java/org/opensearch/indices/IndicesService.java @@ -2154,6 +2154,14 @@ public ByteSizeValue getTotalIndexingBufferBytes() { return indexingMemoryController.indexingBufferSize(); } + /** + * Updates the native memory budget used by the IndexingMemoryController for flush/throttle decisions. + * Called when the indexing pool group's effective limit changes. + */ + public void setNativeIndexBufferBytes(long bytes) { + indexingMemoryController.setNativeBufferBytes(bytes); + } + /** * Cache something calculated at the shard level. * @param shard the shard this item is part of diff --git a/server/src/main/java/org/opensearch/node/Node.java b/server/src/main/java/org/opensearch/node/Node.java index d52e36589298c..341d7516d8e25 100644 --- a/server/src/main/java/org/opensearch/node/Node.java +++ b/server/src/main/java/org/opensearch/node/Node.java @@ -56,6 +56,8 @@ import org.opensearch.action.search.StreamSearchTransportService; import org.opensearch.action.support.TransportAction; import org.opensearch.action.update.UpdateHelper; +import org.opensearch.arrow.spi.NativeAllocator; +import org.opensearch.arrow.spi.PoolGroup; import org.opensearch.bootstrap.BootstrapCheck; import org.opensearch.bootstrap.BootstrapContext; import org.opensearch.cluster.ClusterInfoService; @@ -1224,6 +1226,14 @@ protected Node(final Environment initialEnvironment, Collection clas // Add the telemetryAwarePlugin components to the existing pluginComponents collection. pluginComponents.addAll(telemetryAwarePluginComponents); + // Extract the NativeAllocator instance (published by ArrowBasePlugin in phase 1) + // so it can be passed to SearchBackEndPlugin.createComponents for virtual pool registration. + final NativeAllocator nativeAllocator = pluginComponents.stream() + .filter(c -> c instanceof NativeAllocator) + .map(c -> (NativeAllocator) c) + .findFirst() + .orElse(null); + @SuppressWarnings("rawtypes") Collection searchBackEndPluginComponents = pluginsService.filterPlugins(SearchBackEndPlugin.class) .stream() @@ -1240,7 +1250,8 @@ protected Node(final Environment initialEnvironment, Collection clas namedWriteableRegistry, clusterModule.getIndexNameExpressionResolver(), repositoriesServiceReference::get, - dataFormatRegistry + dataFormatRegistry, + nativeAllocator ).stream() ) .collect(Collectors.toList()); @@ -1258,6 +1269,15 @@ protected Node(final Environment initialEnvironment, Collection clas indicesService.setSearchStatsContributors(searchStatsContributors); } + // Wire indexing pool group limit changes to the IndexingMemoryController. + // When the rebalancer adjusts the effective limit of pools in the INDEXING group, + // IMC's native buffer is updated to 70% of the new grouped limit. + if (nativeAllocator != null) { + nativeAllocator.addPoolGroupLimitListener(PoolGroup.INDEXING, newGroupLimit -> { + indicesService.setNativeIndexBufferBytes((long) (newGroupLimit * 0.70)); + }); + } + if (nodeCacheService != null) { nodeCacheService.registerProviders(blockCacheProviders); } diff --git a/server/src/main/java/org/opensearch/node/NodeService.java b/server/src/main/java/org/opensearch/node/NodeService.java index 3a7aa0ee0b5dd..1ab06c811b8d3 100644 --- a/server/src/main/java/org/opensearch/node/NodeService.java +++ b/server/src/main/java/org/opensearch/node/NodeService.java @@ -263,7 +263,6 @@ public NodeStats stats( boolean admissionControl, boolean cacheService, boolean remoteStoreNodeStats, - boolean nativeAllocator, boolean nativeMemory ) { // for indices stats we want to include previous allocated shards stats as well (it will @@ -301,8 +300,7 @@ public NodeStats stats( admissionControl ? this.admissionControlService.stats() : null, cacheService ? this.cacheService.stats(indices) : null, remoteStoreNodeStats ? new RemoteStoreNodeStats() : null, - nativeAllocator ? collectNativeAllocatorStats() : null, - nativeMemory ? monitorService.memoryReportingService().nativeStats() : null, + nativeMemory ? collectNativeAllocatorStats() : null, // Always capture the process-level native memory estimate on this data node. // Serialized over the wire so the coordinator renders the source node's value, // not its own. Returns -1 on non-Linux platforms or when /proc/self/status is diff --git a/server/src/main/java/org/opensearch/plugin/stats/NativeAllocatorPoolStats.java b/server/src/main/java/org/opensearch/plugin/stats/NativeAllocatorPoolStats.java index 1f27e44d9423d..de067d08ee144 100644 --- a/server/src/main/java/org/opensearch/plugin/stats/NativeAllocatorPoolStats.java +++ b/server/src/main/java/org/opensearch/plugin/stats/NativeAllocatorPoolStats.java @@ -8,6 +8,7 @@ package org.opensearch.plugin.stats; +import org.opensearch.common.Nullable; import org.opensearch.core.common.io.stream.StreamInput; import org.opensearch.core.common.io.stream.StreamOutput; import org.opensearch.core.common.io.stream.Writeable; @@ -17,57 +18,46 @@ import java.io.IOException; import java.util.ArrayList; import java.util.Collections; +import java.util.LinkedHashMap; import java.util.List; +import java.util.Map; /** - * Point-in-time snapshot of native allocator pool statistics for a node. + * Point-in-time snapshot of native memory statistics for a node. * - *

Arrow-agnostic POJO. The plugin that owns the allocator (e.g. {@code arrow-base}) - * constructs instances of this class and exposes them through a - * {@link NativeAllocatorStatsRegistry} component returned from its - * {@code createComponents()}. Server is the type's home so that the cross-module - * dependency from {@code :server} to {@code :libs:opensearch-arrow-spi} is unnecessary, - * mirroring the placement of {@link AnalyticsBackendNativeMemoryStats}. + *

Includes process-wide native memory stats (allocated/resident from jemalloc) + * and per-pool stats for all registered pools (Arrow and virtual). * - *

Renders as the inner body of the {@code native_allocator} object inside - * {@code _nodes/stats[/native_allocator]} — the caller ({@code NodeStats.toXContent}) - * is responsible for opening the {@code native_allocator} wrapper. Each pool exposes - * {@code allocated_bytes}, {@code peak_bytes}, and {@code limit_bytes}; root exposes - * the same. + *

Renders as the body of the {@code native_memory} object inside + * {@code _nodes/stats/native_memory}. * * @opensearch.api */ public class NativeAllocatorPoolStats implements Writeable, ToXContentFragment { - private final long rootAllocatedBytes; - private final long rootPeakBytes; - private final long rootLimitBytes; + private final long nativeAllocatedBytes; + private final long nativeResidentBytes; private final List pools; /** - * Creates a new stats snapshot from the given values. + * Creates a new stats snapshot. * - * @param rootAllocatedBytes current bytes allocated by the root - * @param rootPeakBytes peak bytes ever allocated by the root since process start - * @param rootLimitBytes configured root limit - * @param pools per-pool stats + * @param nativeAllocatedBytes process-wide native allocated bytes (jemalloc), -1 if unavailable + * @param nativeResidentBytes process-wide native resident bytes (jemalloc RSS), -1 if unavailable + * @param pools per-pool stats (Arrow + virtual) */ - public NativeAllocatorPoolStats(long rootAllocatedBytes, long rootPeakBytes, long rootLimitBytes, List pools) { - this.rootAllocatedBytes = rootAllocatedBytes; - this.rootPeakBytes = rootPeakBytes; - this.rootLimitBytes = rootLimitBytes; + public NativeAllocatorPoolStats(long nativeAllocatedBytes, long nativeResidentBytes, List pools) { + this.nativeAllocatedBytes = nativeAllocatedBytes; + this.nativeResidentBytes = nativeResidentBytes; this.pools = Collections.unmodifiableList(pools); } /** * Deserializes from stream. - * - * @param in the stream input */ public NativeAllocatorPoolStats(StreamInput in) throws IOException { - this.rootAllocatedBytes = in.readVLong(); - this.rootPeakBytes = in.readVLong(); - this.rootLimitBytes = in.readVLong(); + this.nativeAllocatedBytes = in.readLong(); + this.nativeResidentBytes = in.readLong(); int count = in.readVInt(); List list = new ArrayList<>(count); for (int i = 0; i < count; i++) { @@ -78,9 +68,8 @@ public NativeAllocatorPoolStats(StreamInput in) throws IOException { @Override public void writeTo(StreamOutput out) throws IOException { - out.writeVLong(rootAllocatedBytes); - out.writeVLong(rootPeakBytes); - out.writeVLong(rootLimitBytes); + out.writeLong(nativeAllocatedBytes); + out.writeLong(nativeResidentBytes); out.writeVInt(pools.size()); for (PoolStats pool : pools) { pool.writeTo(out); @@ -89,11 +78,8 @@ public void writeTo(StreamOutput out) throws IOException { @Override public XContentBuilder toXContent(XContentBuilder builder, Params params) throws IOException { - builder.startObject("root"); - builder.field("allocated_bytes", rootAllocatedBytes); - builder.field("peak_bytes", rootPeakBytes); - builder.field("limit_bytes", rootLimitBytes); - builder.endObject(); + builder.field("allocated_bytes", nativeAllocatedBytes); + builder.field("resident_bytes", nativeResidentBytes); builder.startObject("pools"); for (PoolStats pool : pools) { @@ -103,19 +89,14 @@ public XContentBuilder toXContent(XContentBuilder builder, Params params) throws return builder; } - /** Returns the root allocator's currently allocated bytes. */ - public long getRootAllocatedBytes() { - return rootAllocatedBytes; - } - - /** Returns the root allocator's peak allocated bytes since process start. */ - public long getRootPeakBytes() { - return rootPeakBytes; + /** Returns process-wide native allocated bytes, or -1 if unavailable. */ + public long getNativeAllocatedBytes() { + return nativeAllocatedBytes; } - /** Returns the root allocator's configured limit in bytes. */ - public long getRootLimitBytes() { - return rootLimitBytes; + /** Returns process-wide native resident bytes (RSS), or -1 if unavailable. */ + public long getNativeResidentBytes() { + return nativeResidentBytes; } /** Returns the per-pool statistics. */ @@ -123,6 +104,66 @@ public List getPools() { return pools; } + /** Returns stats grouped by pool group. Pools without a group use their name as the key. */ + public Map getGroupedStats() { + // [allocated, peak, limit] — peak uses max (highest watermark) rather than sum + // because individual pool peaks are not additive (they occur at different times). + Map grouped = new LinkedHashMap<>(); + for (PoolStats pool : pools) { + String g = pool.getGroup() != null ? pool.getGroup() : pool.getName(); + grouped.merge( + g, + new long[] { pool.getAllocatedBytes(), pool.getPeakBytes(), pool.getLimitBytes() }, + (a, b) -> new long[] { a[0] + b[0], Math.max(a[1], b[1]), a[2] + b[2] } + ); + } + Map result = new LinkedHashMap<>(); + for (var e : grouped.entrySet()) { + result.put(e.getKey(), new PoolStats(e.getKey(), e.getValue()[0], e.getValue()[1], e.getValue()[2])); + } + return result; + } + + /** + * BWC: Reads and discards the V_3_7_0 format (3 VLongs + pools with 4 fields each). + * Returns a dummy instance since the data is discarded. + */ + public static NativeAllocatorPoolStats readAndDiscardV3_7(StreamInput in) throws IOException { + in.readVLong(); // rootAllocatedBytes + in.readVLong(); // rootPeakBytes + in.readVLong(); // rootLimitBytes + int count = in.readVInt(); + for (int i = 0; i < count; i++) { + in.readString(); // name + in.readVLong(); // allocatedBytes + in.readVLong(); // peakBytes + in.readVLong(); // limitBytes + } + return new NativeAllocatorPoolStats(-1L, -1L, List.of()); + } + + /** + * BWC: Writes old V_3_7_0 format for a given stats instance (or null). + * Format: optional boolean + 3 VLongs + pool list with 4 fields each. + */ + public static void writeV3_7(StreamOutput out, @Nullable NativeAllocatorPoolStats stats) throws IOException { + if (stats == null) { + out.writeBoolean(false); + } else { + out.writeBoolean(true); + out.writeVLong(stats.nativeAllocatedBytes); // map to rootAllocatedBytes + out.writeVLong(stats.nativeResidentBytes); // map to rootPeakBytes + out.writeVLong(0L); // rootLimitBytes (no equivalent) + out.writeVInt(stats.pools.size()); + for (PoolStats pool : stats.pools) { + out.writeString(pool.getName()); + out.writeVLong(pool.getAllocatedBytes()); + out.writeVLong(pool.getPeakBytes()); + out.writeVLong(pool.getLimitBytes()); + } + } + } + /** * Per-pool statistics snapshot. */ @@ -132,32 +173,29 @@ public static class PoolStats implements Writeable, ToXContentFragment { private final long allocatedBytes; private final long peakBytes; private final long limitBytes; + private final String group; + private final long minBytes; - /** - * Creates a new pool stats snapshot. - * - * @param name pool name - * @param allocatedBytes current allocated bytes - * @param peakBytes peak bytes ever allocated since process start - * @param limitBytes configured limit - */ public PoolStats(String name, long allocatedBytes, long peakBytes, long limitBytes) { + this(name, allocatedBytes, peakBytes, limitBytes, null, 0L); + } + + public PoolStats(String name, long allocatedBytes, long peakBytes, long limitBytes, String group, long minBytes) { this.name = name; this.allocatedBytes = allocatedBytes; this.peakBytes = peakBytes; this.limitBytes = limitBytes; + this.group = group; + this.minBytes = minBytes; } - /** - * Deserializes from stream. - * - * @param in the stream input - */ public PoolStats(StreamInput in) throws IOException { this.name = in.readString(); this.allocatedBytes = in.readVLong(); this.peakBytes = in.readVLong(); this.limitBytes = in.readVLong(); + this.group = in.readOptionalString(); + this.minBytes = in.readVLong(); } @Override @@ -166,36 +204,45 @@ public void writeTo(StreamOutput out) throws IOException { out.writeVLong(allocatedBytes); out.writeVLong(peakBytes); out.writeVLong(limitBytes); + out.writeOptionalString(group); + out.writeVLong(minBytes); } @Override public XContentBuilder toXContent(XContentBuilder builder, Params params) throws IOException { builder.startObject(name); builder.field("allocated_bytes", allocatedBytes); - builder.field("peak_bytes", peakBytes); builder.field("limit_bytes", limitBytes); + builder.field("min_bytes", minBytes); + if (group != null) { + builder.field("group", group); + } builder.endObject(); return builder; } - /** Returns the pool name. */ public String getName() { return name; } - /** Returns the currently allocated bytes. */ public long getAllocatedBytes() { return allocatedBytes; } - /** Returns the peak allocated bytes since process start. */ public long getPeakBytes() { return peakBytes; } - /** Returns the configured limit in bytes. */ public long getLimitBytes() { return limitBytes; } + + public String getGroup() { + return group; + } + + public long getMinBytes() { + return minBytes; + } } } diff --git a/server/src/main/java/org/opensearch/plugins/SearchBackEndPlugin.java b/server/src/main/java/org/opensearch/plugins/SearchBackEndPlugin.java index 1761a0c35d9b7..c103c7dff10a2 100644 --- a/server/src/main/java/org/opensearch/plugins/SearchBackEndPlugin.java +++ b/server/src/main/java/org/opensearch/plugins/SearchBackEndPlugin.java @@ -8,6 +8,7 @@ package org.opensearch.plugins; +import org.opensearch.arrow.spi.NativeAllocator; import org.opensearch.cluster.metadata.IndexNameExpressionResolver; import org.opensearch.cluster.service.ClusterService; import org.opensearch.common.Nullable; @@ -105,6 +106,44 @@ default Collection createComponents( return Collections.emptyList(); } + /** + * Extended variant that also receives the unified native memory allocator. + * Plugins that need to register virtual pools (e.g., DataFusion) override this method. + * The default delegates to the original method for backwards compatibility. + * + * @param nativeAllocator the unified native allocator, or null if arrow-base is not installed + */ + default Collection createComponents( + Client client, + ClusterService clusterService, + ThreadPool threadPool, + ResourceWatcherService resourceWatcherService, + ScriptService scriptService, + NamedXContentRegistry xContentRegistry, + Environment environment, + NodeEnvironment nodeEnvironment, + NamedWriteableRegistry namedWriteableRegistry, + IndexNameExpressionResolver indexNameExpressionResolver, + Supplier repositoriesServiceSupplier, + DataFormatRegistry dataFormatRegistry, + @Nullable NativeAllocator nativeAllocator + ) { + return createComponents( + client, + clusterService, + threadPool, + resourceWatcherService, + scriptService, + xContentRegistry, + environment, + nodeEnvironment, + namedWriteableRegistry, + indexNameExpressionResolver, + repositoriesServiceSupplier, + dataFormatRegistry + ); + } + /** * Returns a supplier for native task cancellation stats, or {@code null} if not available. *

diff --git a/server/src/test/java/org/opensearch/action/admin/cluster/node/stats/AnalyticsBackendNativeMemoryStatsVersionGateTests.java b/server/src/test/java/org/opensearch/action/admin/cluster/node/stats/AnalyticsBackendNativeMemoryStatsVersionGateTests.java deleted file mode 100644 index c2dc2424d74ab..0000000000000 --- a/server/src/test/java/org/opensearch/action/admin/cluster/node/stats/AnalyticsBackendNativeMemoryStatsVersionGateTests.java +++ /dev/null @@ -1,166 +0,0 @@ -/* - * SPDX-License-Identifier: Apache-2.0 - * - * The OpenSearch Contributors require contributions made to - * this file be licensed under the Apache-2.0 license or a - * compatible open source license. - */ - -package org.opensearch.action.admin.cluster.node.stats; - -import org.opensearch.Version; -import org.opensearch.cluster.node.DiscoveryNode; -import org.opensearch.common.io.stream.BytesStreamOutput; -import org.opensearch.core.common.io.stream.StreamInput; -import org.opensearch.plugin.stats.AnalyticsBackendNativeMemoryStats; -import org.opensearch.test.OpenSearchTestCase; - -import java.io.IOException; - -import static java.util.Collections.emptyMap; -import static java.util.Collections.emptySet; - -/** - * Property-based tests for version-gated serialization of {@link AnalyticsBackendNativeMemoryStats} - * within {@link NodeStats}. - * - * Verifies that when NodeStats containing a non-null AnalyticsBackendNativeMemoryStats is serialized - * to a stream with a version older than V_3_7_0, the deserialized NodeStats has - * nativeMemoryStats == null. Conversely, when serialized to V_3_7_0 or later, the - * AnalyticsBackendNativeMemoryStats is preserved. - */ -public class AnalyticsBackendNativeMemoryStatsVersionGateTests extends OpenSearchTestCase { - - /** - * Property 3: Version-gated serialization preserves null for old versions. - * - * For any NodeStats containing a non-null AnalyticsBackendNativeMemoryStats, serializing to a stream - * with version older than the native-memory support version (V_3_7_0) and then - * deserializing SHALL yield a NodeStats with nativeMemoryStats == null. - * - * Validates: Requirements 3.4, 3.5 - */ - public void testVersionGatedSerializationOmitsAnalyticsBackendNativeMemoryStatsForOldVersions() throws IOException { - for (int i = 0; i < 100; i++) { - long allocatedBytes = randomLongBetween(Long.MIN_VALUE, Long.MAX_VALUE); - long residentBytes = randomLongBetween(Long.MIN_VALUE, Long.MAX_VALUE); - AnalyticsBackendNativeMemoryStats nativeMemoryStats = new AnalyticsBackendNativeMemoryStats(allocatedBytes, residentBytes); - - NodeStats nodeStats = createNodeStatsWithNativeMemory(nativeMemoryStats); - - // Serialize with a version older than V_3_7_0 - try (BytesStreamOutput out = new BytesStreamOutput()) { - out.setVersion(Version.V_2_18_0); - nodeStats.writeTo(out); - - try (StreamInput in = out.bytes().streamInput()) { - in.setVersion(Version.V_2_18_0); - NodeStats deserialized = new NodeStats(in); - - assertNull( - "nativeMemoryStats should be null when deserialized from version < V_3_7_0, " - + "iteration " - + i - + " with values [" - + allocatedBytes - + ", " - + residentBytes - + "]", - deserialized.getAnalyticsBackendNativeMemoryStats() - ); - } - } - } - } - - /** - * Positive case: Version-gated serialization preserves AnalyticsBackendNativeMemoryStats for V_3_7_0+. - * - * For any NodeStats containing a non-null AnalyticsBackendNativeMemoryStats, serializing to a stream - * with version V_3_7_0 or later and then deserializing SHALL yield a NodeStats with - * nativeMemoryStats containing the original values. - * - * Validates: Requirements 3.4, 3.5 - */ - public void testVersionGatedSerializationPreservesAnalyticsBackendNativeMemoryStatsForCurrentVersion() throws IOException { - for (int i = 0; i < 100; i++) { - long allocatedBytes = randomLongBetween(Long.MIN_VALUE, Long.MAX_VALUE); - long residentBytes = randomLongBetween(Long.MIN_VALUE, Long.MAX_VALUE); - AnalyticsBackendNativeMemoryStats nativeMemoryStats = new AnalyticsBackendNativeMemoryStats(allocatedBytes, residentBytes); - - NodeStats nodeStats = createNodeStatsWithNativeMemory(nativeMemoryStats); - - // Serialize with V_3_7_0 (the version that introduced native memory support) - try (BytesStreamOutput out = new BytesStreamOutput()) { - out.setVersion(Version.V_3_7_0); - nodeStats.writeTo(out); - - try (StreamInput in = out.bytes().streamInput()) { - in.setVersion(Version.V_3_7_0); - NodeStats deserialized = new NodeStats(in); - - assertNotNull( - "nativeMemoryStats should be non-null when deserialized from version >= V_3_7_0, " + "iteration " + i, - deserialized.getAnalyticsBackendNativeMemoryStats() - ); - assertEquals( - "allocatedBytes mismatch on iteration " + i, - allocatedBytes, - deserialized.getAnalyticsBackendNativeMemoryStats().getAllocatedBytes() - ); - assertEquals( - "residentBytes mismatch on iteration " + i, - residentBytes, - deserialized.getAnalyticsBackendNativeMemoryStats().getResidentBytes() - ); - } - } - } - } - - /** - * Creates a minimal NodeStats with the given AnalyticsBackendNativeMemoryStats and all other fields null. - * Uses the current version for the DiscoveryNode. - */ - private NodeStats createNodeStatsWithNativeMemory(AnalyticsBackendNativeMemoryStats nativeMemoryStats) { - DiscoveryNode node = new DiscoveryNode("test_node", buildNewFakeTransportAddress(), emptyMap(), emptySet(), Version.CURRENT); - - return new NodeStats( - node, - System.currentTimeMillis(), - null, // indices - null, // os - null, // process - null, // jvm - null, // threadPool - null, // fs - null, // transport - null, // http - null, // breaker - null, // scriptStats - null, // discoveryStats - null, // ingestStats - null, // adaptiveSelectionStats - null, // resourceUsageStats - null, // scriptCacheStats - null, // indexingPressureStats - null, // shardIndexingPressureStats - null, // searchBackpressureStats - null, // clusterManagerThrottlingStats - null, // weightedRoutingStats - null, // fileCacheStats - null, // fileCacheOnlyStats - null, // blockCacheOnlyStats - null, // taskCancellationStats - null, // searchPipelineStats - null, // segmentReplicationRejectionStats - null, // repositoriesStats - null, // admissionControlStats - null, // nodeCacheStats - null, // remoteStoreNodeStats - null, // nativeAllocator - nativeMemoryStats, - -1L // totalEstimatedNativeBytes - ); - } -} diff --git a/server/src/test/java/org/opensearch/action/admin/cluster/node/stats/NodeStatsTests.java b/server/src/test/java/org/opensearch/action/admin/cluster/node/stats/NodeStatsTests.java index 05f27c3f98562..7d9f77f1d0bd8 100644 --- a/server/src/test/java/org/opensearch/action/admin/cluster/node/stats/NodeStatsTests.java +++ b/server/src/test/java/org/opensearch/action/admin/cluster/node/stats/NodeStatsTests.java @@ -1056,7 +1056,6 @@ public long getLastSuccessfulFetchOfPinnedTimestamps() { nodeCacheStats, remoteStoreNodeStats, null, - null, -1L ); } @@ -1524,7 +1523,6 @@ public void testNativeAllocatorStatsBwcEmptyOnOldVersion() throws IOException { NativeAllocatorPoolStats stats = new NativeAllocatorPoolStats( 1024L, 2048L, - 8192L, List.of(new NativeAllocatorPoolStats.PoolStats("flight", 100L, 200L, 2048L)) ); DiscoveryNode node = new DiscoveryNode("node1", buildNewFakeTransportAddress(), emptyMap(), emptySet(), Version.CURRENT); @@ -1554,7 +1552,6 @@ public void testNativeAllocatorStatsRoundTripCurrentVersion() throws IOException NativeAllocatorPoolStats stats = new NativeAllocatorPoolStats( 1024L, 2048L, - 8192L, List.of( new NativeAllocatorPoolStats.PoolStats("flight", 100L, 200L, 2048L), new NativeAllocatorPoolStats.PoolStats("ingest", 200L, 400L, 4096L), @@ -1572,9 +1569,8 @@ public void testNativeAllocatorStatsRoundTripCurrentVersion() throws IOException NodeStats roundtripped = new NodeStats(in); NativeAllocatorPoolStats decoded = roundtripped.getNativeAllocatorStats(); assertNotNull("native allocator stats must round-trip on current wire version", decoded); - assertEquals(1024L, decoded.getRootAllocatedBytes()); - assertEquals(2048L, decoded.getRootPeakBytes()); - assertEquals(8192L, decoded.getRootLimitBytes()); + assertEquals(1024L, decoded.getNativeAllocatedBytes()); + assertEquals(2048L, decoded.getNativeResidentBytes()); assertEquals(3, decoded.getPools().size()); assertEquals("flight", decoded.getPools().get(0).getName()); assertEquals(100L, decoded.getPools().get(0).getAllocatedBytes()); @@ -1586,15 +1582,13 @@ public void testNativeAllocatorStatsRoundTripCurrentVersion() throws IOException /** * Renders {@code NodeStats.toXContent} when {@code nativeAllocatorStats} is non-null and - * asserts the JSON shape: a top-level {@code native_memory.native_allocator} block with - * the SPI's inner {@code root}/{@code pools.} structure. Covers the conditional - * branch in {@code NodeStats.toXContent} that opens the {@code native_allocator} wrapper. + * asserts the JSON shape: a top-level {@code native_memory} block with + * {@code runtime.allocated_bytes}/{@code runtime.resident_bytes} and grouped {@code memory_pools}. */ public void testNativeAllocatorStatsXContentRendersInsideNativeMemory() throws IOException { NativeAllocatorPoolStats stats = new NativeAllocatorPoolStats( 1024L, 2048L, - 8192L, List.of(new NativeAllocatorPoolStats.PoolStats("flight", 100L, 200L, 2048L)) ); DiscoveryNode node = new DiscoveryNode("node1", buildNewFakeTransportAddress(), emptyMap(), emptySet(), Version.CURRENT); @@ -1608,20 +1602,22 @@ public void testNativeAllocatorStatsXContentRendersInsideNativeMemory() throws I @SuppressWarnings("unchecked") Map nativeMemory = (Map) root.get("native_memory"); assertNotNull("native_memory wrapper must be opened when allocator stats are present", nativeMemory); + + // Runtime stats are nested under "runtime" @SuppressWarnings("unchecked") - Map nativeAllocator = (Map) nativeMemory.get("native_allocator"); - assertNotNull("native_allocator block must be present", nativeAllocator); - @SuppressWarnings("unchecked") - Map rootBlock = (Map) nativeAllocator.get("root"); - assertEquals(1024L, ((Number) rootBlock.get("allocated_bytes")).longValue()); - assertEquals(2048L, ((Number) rootBlock.get("peak_bytes")).longValue()); - assertEquals(8192L, ((Number) rootBlock.get("limit_bytes")).longValue()); + Map runtime = (Map) nativeMemory.get("runtime"); + assertNotNull("runtime block must be present", runtime); + assertEquals(1024L, ((Number) runtime.get("allocated_bytes")).longValue()); + assertEquals(2048L, ((Number) runtime.get("resident_bytes")).longValue()); + + // Pools are grouped under "memory_pools" @SuppressWarnings("unchecked") - Map pools = (Map) nativeAllocator.get("pools"); + Map pools = (Map) nativeMemory.get("memory_pools"); + assertNotNull("memory_pools block must be present", pools); @SuppressWarnings("unchecked") Map flight = (Map) pools.get("flight"); + assertNotNull("flight pool must be present in memory_pools", flight); assertEquals(100L, ((Number) flight.get("allocated_bytes")).longValue()); - assertEquals(200L, ((Number) flight.get("peak_bytes")).longValue()); assertEquals(2048L, ((Number) flight.get("limit_bytes")).longValue()); } @@ -1704,7 +1700,6 @@ private static NodeStats newNodeStatsWithNativeAllocator( null, // nodeCacheStats null, nativeAllocatorStats, - null, totalEstimatedNativeBytes ); } diff --git a/server/src/test/java/org/opensearch/action/admin/cluster/stats/ClusterStatsNodesTests.java b/server/src/test/java/org/opensearch/action/admin/cluster/stats/ClusterStatsNodesTests.java index 9820102840829..fcdd26912848f 100644 --- a/server/src/test/java/org/opensearch/action/admin/cluster/stats/ClusterStatsNodesTests.java +++ b/server/src/test/java/org/opensearch/action/admin/cluster/stats/ClusterStatsNodesTests.java @@ -354,7 +354,6 @@ private ClusterStatsNodeResponse createClusterStatsNodeResponse( null, null, null, // nativeAllocator - null, -1L // totalEstimatedNativeBytes ); if (defaultBehavior) { diff --git a/server/src/test/java/org/opensearch/action/admin/cluster/stats/ClusterStatsResponseTests.java b/server/src/test/java/org/opensearch/action/admin/cluster/stats/ClusterStatsResponseTests.java index b1ae92df3793c..7ff9dd3d6a89e 100644 --- a/server/src/test/java/org/opensearch/action/admin/cluster/stats/ClusterStatsResponseTests.java +++ b/server/src/test/java/org/opensearch/action/admin/cluster/stats/ClusterStatsResponseTests.java @@ -226,7 +226,6 @@ private ClusterStatsNodeResponse createClusterStatsNodeResponse(DiscoveryNode no null, null, null, // nativeAllocator - null, -1L // totalEstimatedNativeBytes ); return new ClusterStatsNodeResponse(node, null, nodeInfo, nodeStats, shardStats); diff --git a/server/src/test/java/org/opensearch/cluster/DiskUsageTests.java b/server/src/test/java/org/opensearch/cluster/DiskUsageTests.java index ec689d6554b33..2bf4b1fcd8b94 100644 --- a/server/src/test/java/org/opensearch/cluster/DiskUsageTests.java +++ b/server/src/test/java/org/opensearch/cluster/DiskUsageTests.java @@ -216,7 +216,6 @@ public void testFillDiskUsage() { null, null, null, // nativeAllocator - null, -1L // totalEstimatedNativeBytes ), new NodeStats( @@ -253,7 +252,6 @@ public void testFillDiskUsage() { null, null, null, // nativeAllocator - null, -1L // totalEstimatedNativeBytes ), new NodeStats( @@ -290,7 +288,6 @@ public void testFillDiskUsage() { null, null, null, // nativeAllocator - null, -1L // totalEstimatedNativeBytes ) ); @@ -358,7 +355,6 @@ public void testFillDiskUsageSomeInvalidValues() { null, null, null, // nativeAllocator - null, -1L // totalEstimatedNativeBytes ), new NodeStats( @@ -395,7 +391,6 @@ public void testFillDiskUsageSomeInvalidValues() { null, null, null, // nativeAllocator - null, -1L // totalEstimatedNativeBytes ), new NodeStats( @@ -432,7 +427,6 @@ public void testFillDiskUsageSomeInvalidValues() { null, null, null, // nativeAllocator - null, -1L // totalEstimatedNativeBytes ) ); @@ -529,7 +523,6 @@ private NodeStats makeNodeStatsWithResourceUsage(DiscoveryNode node, NodesResour null, null, null, - null, -1L ); diff --git a/server/src/test/java/org/opensearch/indices/IndexingMemoryControllerTests.java b/server/src/test/java/org/opensearch/indices/IndexingMemoryControllerTests.java index 659af894d6478..99c657480d55e 100644 --- a/server/src/test/java/org/opensearch/indices/IndexingMemoryControllerTests.java +++ b/server/src/test/java/org/opensearch/indices/IndexingMemoryControllerTests.java @@ -743,4 +743,23 @@ public void testCombinedSortKeyUsedWhenBothExceed() throws IOException { assertEquals(500L, (long) controller.nativeMemoryUsed.get(shardB)); closeShards(shardA, shardB); } + + public void testSetNativeBufferBytesUpdatesThreshold() { + Settings settings = Settings.builder() + .put("indices.memory.native_index_buffer_size", "100mb") + .put("indices.memory.index_buffer_size", "10%") + .build(); + MockController controller = new MockController(settings); + + // Initial value from setting: 100MB + assertEquals(100L * 1024 * 1024, controller.nativeBufferSizeInBytes()); + + // Dynamically update + controller.setNativeBufferBytes(200L * 1024 * 1024); + assertEquals(200L * 1024 * 1024, controller.nativeBufferSizeInBytes()); + + // Update again + controller.setNativeBufferBytes(50L * 1024 * 1024); + assertEquals(50L * 1024 * 1024, controller.nativeBufferSizeInBytes()); + } } diff --git a/server/src/test/java/org/opensearch/node/NodeServiceNativeMemoryTests.java b/server/src/test/java/org/opensearch/node/NodeServiceNativeMemoryTests.java index 2598165dc6dcc..e69eefd2be669 100644 --- a/server/src/test/java/org/opensearch/node/NodeServiceNativeMemoryTests.java +++ b/server/src/test/java/org/opensearch/node/NodeServiceNativeMemoryTests.java @@ -16,19 +16,13 @@ import org.opensearch.cluster.service.ClusterService; import org.opensearch.common.settings.Settings; import org.opensearch.common.settings.SettingsFilter; -import org.opensearch.common.xcontent.XContentHelper; -import org.opensearch.common.xcontent.json.JsonXContent; import org.opensearch.core.indices.breaker.CircuitBreakerService; -import org.opensearch.core.xcontent.ToXContent; -import org.opensearch.core.xcontent.XContentBuilder; import org.opensearch.discovery.Discovery; import org.opensearch.index.IndexingPressureService; import org.opensearch.index.SegmentReplicationStatsTracker; import org.opensearch.indices.IndicesService; import org.opensearch.ingest.IngestService; import org.opensearch.monitor.MonitorService; -import org.opensearch.monitor.memory.MemoryReportingService; -import org.opensearch.plugin.stats.AnalyticsBackendNativeMemoryStats; import org.opensearch.plugin.stats.NativeAllocatorPoolStats; import org.opensearch.plugins.PluginsService; import org.opensearch.ratelimitting.admissioncontrol.AdmissionControlService; @@ -43,7 +37,6 @@ import java.util.Collections; import java.util.List; -import java.util.Map; import java.util.function.Supplier; import static org.mockito.Mockito.mock; @@ -52,38 +45,21 @@ /** * Unit tests for NodeService native memory stats delegation logic. *

- * Validates that NodeService correctly delegates to the native memory stats + * Validates that NodeService correctly delegates to the native allocator stats * supplier when nativeMemory=true and the supplier is non-null, * and returns null otherwise. */ public class NodeServiceNativeMemoryTests extends OpenSearchTestCase { - private NodeService createNodeService(AnalyticsBackendNativeMemoryStats nativeStats) { - return createNodeService(nativeStats, null); - } - - private NodeService createNodeService( - AnalyticsBackendNativeMemoryStats nativeStats, - Supplier nativeAllocatorStatsSupplier - ) { + private NodeService createNodeService(Supplier nativeAllocatorStatsSupplier) { TransportService transportService = mock(TransportService.class); DiscoveryNode localNode = new DiscoveryNode("test_node", buildNewFakeTransportAddress(), Version.CURRENT); when(transportService.getLocalNode()).thenReturn(localNode); - ClusterService clusterService = mock(ClusterService.class); - IngestService ingestService = mock(IngestService.class); - SearchPipelineService searchPipelineService = mock(SearchPipelineService.class); - - MemoryReportingService memoryReportingService = mock(MemoryReportingService.class); - when(memoryReportingService.nativeStats()).thenReturn(nativeStats); - - MonitorService monitorService = mock(MonitorService.class); - when(monitorService.memoryReportingService()).thenReturn(memoryReportingService); - return new NodeService( Settings.EMPTY, mock(ThreadPool.class), - monitorService, + mock(MonitorService.class), mock(Discovery.class), transportService, mock(IndicesService.class), @@ -91,16 +67,16 @@ private NodeService createNodeService( mock(CircuitBreakerService.class), mock(ScriptService.class), null, // httpServerTransport - ingestService, - clusterService, + mock(IngestService.class), + mock(ClusterService.class), new SettingsFilter(Collections.emptyList()), null, // responseCollectorService - not needed when adaptiveSelection=false mock(SearchTransportService.class), mock(IndexingPressureService.class), null, // aggregationUsageService mock(SearchBackpressureService.class), - searchPipelineService, - null, // fileCache + mock(SearchPipelineService.class), + null, // nodeCacheService mock(TaskCancellationMonitoringService.class), null, // resourceUsageCollectorService mock(SegmentReplicationStatsTracker.class), @@ -115,10 +91,13 @@ private NodeService createNodeService( * Tests that stats() with nativeMemory=true and a non-null supplier * returns the stats from the supplier. */ - public void testStatsWithNativeMemoryTrueAndServicePresent() { - AnalyticsBackendNativeMemoryStats expectedStats = new AnalyticsBackendNativeMemoryStats(1024L, 2048L); - - NodeService nodeService = createNodeService(expectedStats); + public void testStatsWithNativeMemoryTrueAndSupplierPresent() { + NativeAllocatorPoolStats expected = new NativeAllocatorPoolStats( + 1024L, + 2048L, + List.of(new NativeAllocatorPoolStats.PoolStats("flight", 100L, 200L, 2048L)) + ); + NodeService nodeService = createNodeService(() -> expected); NodeStats nodeStats = nodeService.stats( CommonStatsFlags.NONE, @@ -150,21 +129,18 @@ public void testStatsWithNativeMemoryTrueAndServicePresent() { false, // admissionControl false, // cacheService false, // remoteStoreNodeStats - false, // nativeAllocator true // nativeMemory ); - assertNotNull(nodeStats.getAnalyticsBackendNativeMemoryStats()); - assertSame(expectedStats, nodeStats.getAnalyticsBackendNativeMemoryStats()); - assertEquals(1024L, nodeStats.getAnalyticsBackendNativeMemoryStats().getAllocatedBytes()); - assertEquals(2048L, nodeStats.getAnalyticsBackendNativeMemoryStats().getResidentBytes()); + assertNotNull("nativeAllocatorStats should be present when supplier returns non-null", nodeStats.getNativeAllocatorStats()); + assertSame(expected, nodeStats.getNativeAllocatorStats()); } /** - * Tests that stats() with nativeMemory=true and a null supplier - * returns null for the nativeMemoryStats field. + * Tests that stats() with nativeMemory=true and no supplier + * returns null for the nativeAllocatorStats field. */ - public void testStatsWithNativeMemoryTrueAndNullService() { + public void testStatsWithNativeMemoryTrueAndNoSupplier() { NodeService nodeService = createNodeService(null); NodeStats nodeStats = nodeService.stats( @@ -197,21 +173,23 @@ public void testStatsWithNativeMemoryTrueAndNullService() { false, // admissionControl false, // cacheService false, // remoteStoreNodeStats - false, // nativeAllocator true // nativeMemory ); - assertNull(nodeStats.getAnalyticsBackendNativeMemoryStats()); + assertNull("nativeAllocatorStats should be null when no supplier registered", nodeStats.getNativeAllocatorStats()); } /** * Tests that stats() with nativeMemory=false returns null for the - * nativeMemoryStats field regardless of whether the supplier is present. + * nativeAllocatorStats field regardless of whether the supplier is present. */ public void testStatsWithNativeMemoryFalse() { - AnalyticsBackendNativeMemoryStats expectedStats = new AnalyticsBackendNativeMemoryStats(4096L, 8192L); - - NodeService nodeService = createNodeService(expectedStats); + NativeAllocatorPoolStats expected = new NativeAllocatorPoolStats( + 4096L, + 8192L, + List.of(new NativeAllocatorPoolStats.PoolStats("flight", 100L, 200L, 2048L)) + ); + NodeService nodeService = createNodeService(() -> expected); NodeStats nodeStats = nodeService.stats( CommonStatsFlags.NONE, @@ -243,225 +221,9 @@ public void testStatsWithNativeMemoryFalse() { false, // admissionControl false, // cacheService false, // remoteStoreNodeStats - false, // nativeAllocator - false // nativeMemory - ); - - assertNull(nodeStats.getAnalyticsBackendNativeMemoryStats()); - } - - /** - * Integration test: verifies that the _nodes/stats/native_memory response format - * contains the expected "native_memory" object with "allocated_bytes" and "resident_bytes" fields. - * This ensures the response format is unchanged after the refactor. - */ - @SuppressWarnings("unchecked") - public void testNativeMemoryResponseFormatUnchanged() throws Exception { - AnalyticsBackendNativeMemoryStats expectedStats = new AnalyticsBackendNativeMemoryStats(123456789L, 987654321L); - - NodeService nodeService = createNodeService(expectedStats); - - NodeStats nodeStats = nodeService.stats( - CommonStatsFlags.NONE, - false, - false, - false, - false, - false, - false, - false, - false, - false, - false, - false, - false, - false, - false, - false, - false, - false, - false, - false, - false, // fileCacheDetailed - false, - false, - false, - false, - false, - false, - false, - false, - false, // nativeAllocator - true // nativeMemory - ); - - assertNotNull("nativeMemoryStats should be present", nodeStats.getAnalyticsBackendNativeMemoryStats()); - - // Render the parent NodeStats to JSON — NodeStats now opens the `native_memory` - // wrapper, emits `total_estimated_bytes` from OsProbe, then delegates to - // AnalyticsBackendNativeMemoryStats which renders only the `analytics_backend` block. - XContentBuilder builder = JsonXContent.contentBuilder(); - builder.startObject(); - nodeStats.toXContent(builder, ToXContent.EMPTY_PARAMS); - builder.endObject(); - String json = builder.toString(); - - Map root = XContentHelper.convertToMap(JsonXContent.jsonXContent, json, false); - - // Verify "native_memory" object is present - assertTrue("Response should contain 'native_memory' key", root.containsKey("native_memory")); - - @SuppressWarnings("unchecked") - Map nativeMemory = (Map) root.get("native_memory"); - assertNotNull("native_memory object should not be null", nativeMemory); - - // Verify nested "analytics_backend" with correct values - assertTrue("native_memory should contain 'analytics_backend'", nativeMemory.containsKey("analytics_backend")); - @SuppressWarnings("unchecked") - Map analyticsBackend = (Map) nativeMemory.get("analytics_backend"); - assertEquals(123456789L, ((Number) analyticsBackend.get("allocated_bytes")).longValue()); - assertEquals(987654321L, ((Number) analyticsBackend.get("resident_bytes")).longValue()); - } - - /** - * Integration test: verifies that when native stats are unavailable (null supplier), - * the response omits the native_memory object entirely. - */ - public void testNativeMemoryOmittedWhenUnavailable() throws Exception { - NodeService nodeService = createNodeService(null); - - NodeStats nodeStats = nodeService.stats( - CommonStatsFlags.NONE, - false, - false, - false, - false, - false, - false, - false, - false, - false, - false, - false, - false, - false, - false, - false, - false, - false, - false, - false, - false, // fileCacheDetailed - false, - false, - false, - false, - false, - false, - false, - false, - false, // nativeAllocator - true // nativeMemory - ); - - assertNull("nativeMemoryStats should be null when supplier is null", nodeStats.getAnalyticsBackendNativeMemoryStats()); - } - - /** - * Tests that {@code stats(... nativeAllocator=true ...)} invokes the constructor-injected - * {@code Supplier} and surfaces its return value on - * {@link NodeStats#getNativeAllocatorStats()}. Covers the supplier-invocation branch in - * {@code collectNativeAllocatorStats}. - */ - public void testStatsWithNativeAllocatorTrueAndSupplierPresent() { - NativeAllocatorPoolStats expected = new NativeAllocatorPoolStats( - 1024L, - 2048L, - 8192L, - List.of(new NativeAllocatorPoolStats.PoolStats("flight", 100L, 200L, 2048L)) - ); - NodeService nodeService = createNodeService(null, () -> expected); - - NodeStats nodeStats = nodeService.stats( - CommonStatsFlags.NONE, - false, - false, - false, - false, - false, - false, - false, - false, - false, - false, - false, - false, - false, - false, - false, - false, - false, - false, - false, - false, - false, - false, - false, - false, - false, - false, - false, - false, - true, // nativeAllocator false // nativeMemory ); - assertNotNull("nativeAllocatorStats should be present when supplier returns non-null", nodeStats.getNativeAllocatorStats()); - assertSame(expected, nodeStats.getNativeAllocatorStats()); - } - - /** - * Tests that {@code stats(... nativeAllocator=true ...)} returns {@code null} for the - * allocator stats when no supplier was injected at construction. - */ - public void testStatsWithNativeAllocatorTrueAndNoSupplier() { - NodeService nodeService = createNodeService(null); - // No supplier passed to the factory — defaults to null. - - NodeStats nodeStats = nodeService.stats( - CommonStatsFlags.NONE, - false, - false, - false, - false, - false, - false, - false, - false, - false, - false, - false, - false, - false, - false, - false, - false, - false, - false, - false, - false, - false, - false, - false, - false, - false, - false, - false, - false, - true, // nativeAllocator - false // nativeMemory - ); - - assertNull("nativeAllocatorStats should be null when no supplier registered", nodeStats.getNativeAllocatorStats()); + assertNull("nativeAllocatorStats should be null when nativeMemory=false", nodeStats.getNativeAllocatorStats()); } } diff --git a/server/src/test/java/org/opensearch/plugin/stats/NativeAllocatorPoolStatsTests.java b/server/src/test/java/org/opensearch/plugin/stats/NativeAllocatorPoolStatsTests.java index 8257f58f0466c..a49c784a34e51 100644 --- a/server/src/test/java/org/opensearch/plugin/stats/NativeAllocatorPoolStatsTests.java +++ b/server/src/test/java/org/opensearch/plugin/stats/NativeAllocatorPoolStatsTests.java @@ -25,7 +25,7 @@ public void testSerializationRoundTrip() throws IOException { new NativeAllocatorPoolStats.PoolStats("flight", 1000, 2000, 3000), new NativeAllocatorPoolStats.PoolStats("query", 4000, 5000, 6000) ); - NativeAllocatorPoolStats original = new NativeAllocatorPoolStats(10000, 20000, 30000, pools); + NativeAllocatorPoolStats original = new NativeAllocatorPoolStats(10000, 20000, pools); BytesStreamOutput out = new BytesStreamOutput(); original.writeTo(out); @@ -33,9 +33,8 @@ public void testSerializationRoundTrip() throws IOException { StreamInput in = out.bytes().streamInput(); NativeAllocatorPoolStats deserialized = new NativeAllocatorPoolStats(in); - assertEquals(original.getRootAllocatedBytes(), deserialized.getRootAllocatedBytes()); - assertEquals(original.getRootPeakBytes(), deserialized.getRootPeakBytes()); - assertEquals(original.getRootLimitBytes(), deserialized.getRootLimitBytes()); + assertEquals(original.getNativeAllocatedBytes(), deserialized.getNativeAllocatedBytes()); + assertEquals(original.getNativeResidentBytes(), deserialized.getNativeResidentBytes()); assertEquals(original.getPools().size(), deserialized.getPools().size()); for (int i = 0; i < pools.size(); i++) { @@ -49,7 +48,7 @@ public void testSerializationRoundTrip() throws IOException { } public void testEmptyPoolsSerialization() throws IOException { - NativeAllocatorPoolStats original = new NativeAllocatorPoolStats(0, 0, 16000000000L, List.of()); + NativeAllocatorPoolStats original = new NativeAllocatorPoolStats(-1, -1, List.of()); BytesStreamOutput out = new BytesStreamOutput(); original.writeTo(out); @@ -57,23 +56,16 @@ public void testEmptyPoolsSerialization() throws IOException { StreamInput in = out.bytes().streamInput(); NativeAllocatorPoolStats deserialized = new NativeAllocatorPoolStats(in); - assertEquals(0, deserialized.getRootAllocatedBytes()); - assertEquals(0, deserialized.getRootPeakBytes()); - assertEquals(16000000000L, deserialized.getRootLimitBytes()); + assertEquals(-1, deserialized.getNativeAllocatedBytes()); + assertEquals(-1, deserialized.getNativeResidentBytes()); assertTrue(deserialized.getPools().isEmpty()); } - /** - * Asserts the JSON shape: {@code root}/{@code pools.} blocks expose - * {@code allocated_bytes}, {@code peak_bytes}, and {@code limit_bytes}. Caller is - * responsible for the outer {@code native_allocator} wrapper, so this test does - * not expect it. - */ public void testToXContent() throws IOException { List pools = List.of( new NativeAllocatorPoolStats.PoolStats("flight", 1024, 1048576, 2147483648L) ); - NativeAllocatorPoolStats stats = new NativeAllocatorPoolStats(4096, 8192, 17179869184L, pools); + NativeAllocatorPoolStats stats = new NativeAllocatorPoolStats(4096, 8192, pools); XContentBuilder builder = JsonXContent.contentBuilder(); builder.startObject(); @@ -81,17 +73,12 @@ public void testToXContent() throws IOException { builder.endObject(); String json = builder.toString(); - assertTrue(json.contains("\"root\"")); + assertTrue(json.contains("\"allocated_bytes\"")); + assertTrue(json.contains("\"resident_bytes\"")); assertTrue(json.contains("\"pools\"")); assertTrue(json.contains("\"flight\"")); - assertTrue(json.contains("\"allocated_bytes\"")); - assertTrue(json.contains("\"peak_bytes\"")); assertTrue(json.contains("\"limit_bytes\"")); - - // Removed fields must NOT appear in the JSON. - assertFalse("child_count was dropped from the stats shape", json.contains("\"child_count\"")); - assertFalse("human-readable byte string was dropped", json.contains("\"allocated\":")); - assertFalse("human-readable byte string was dropped", json.contains("\"limit\":")); + assertFalse("root object should not exist", json.contains("\"root\"")); } public void testPoolStatsSerializationRoundTrip() throws IOException { diff --git a/test/framework/src/main/java/org/opensearch/cluster/MockInternalClusterInfoService.java b/test/framework/src/main/java/org/opensearch/cluster/MockInternalClusterInfoService.java index 9576112b8b12b..bd2842cdaa20d 100644 --- a/test/framework/src/main/java/org/opensearch/cluster/MockInternalClusterInfoService.java +++ b/test/framework/src/main/java/org/opensearch/cluster/MockInternalClusterInfoService.java @@ -154,7 +154,6 @@ List adjustNodesStats(List nodesStats) { nodeStats.getNodeCacheStats(), nodeStats.getRemoteStoreNodeStats(), nodeStats.getNativeAllocatorStats(), - nodeStats.getAnalyticsBackendNativeMemoryStats(), nodeStats.getTotalEstimatedNativeBytes() ); }).collect(Collectors.toList()); diff --git a/test/framework/src/main/java/org/opensearch/test/InternalTestCluster.java b/test/framework/src/main/java/org/opensearch/test/InternalTestCluster.java index e01f4d651d979..e92c1c6402a9f 100644 --- a/test/framework/src/main/java/org/opensearch/test/InternalTestCluster.java +++ b/test/framework/src/main/java/org/opensearch/test/InternalTestCluster.java @@ -2709,7 +2709,6 @@ public void ensureEstimatedStats() { false, false, false, - false, false ); assertThat( From 4a9d0ed19be6323790068ab3e75c5622b01fc252 Mon Sep 17 00:00:00 2001 From: A S K Kamal Nayan Date: Thu, 11 Jun 2026 16:21:59 +0530 Subject: [PATCH 03/14] Updated to stats to remove dependency from composite engine (#22078) Signed-off-by: Kamal Nayan Co-authored-by: Kamal Nayan Co-authored-by: Mohit Godwani <81609427+mgodwan@users.noreply.github.com> --- sandbox/libs/plugin-stats-spi/build.gradle | 4 +- .../plugin/stats/StatsRecorder.java | 16 ++ .../plugin/stats/StatsRecorderTests.java | 247 ++++++++++++++++++ .../index/LuceneIndexingExecutionEngine.java | 15 +- .../be/lucene/index/LuceneWriter.java | 22 +- .../be/lucene/stats/LuceneStatsProvider.java | 4 +- .../opensearch/composite/StatsFailureIT.java | 5 + .../parquet/bridge/NativeParquetWriter.java | 13 +- .../parquet/stats/ParquetStatsProvider.java | 4 +- .../parquet/writer/ParquetWriter.java | 10 +- 10 files changed, 297 insertions(+), 43 deletions(-) create mode 100644 sandbox/libs/plugin-stats-spi/src/test/java/org/opensearch/plugin/stats/StatsRecorderTests.java diff --git a/sandbox/libs/plugin-stats-spi/build.gradle b/sandbox/libs/plugin-stats-spi/build.gradle index 405dd7437a633..c12544b4caab6 100644 --- a/sandbox/libs/plugin-stats-spi/build.gradle +++ b/sandbox/libs/plugin-stats-spi/build.gradle @@ -14,9 +14,7 @@ java { sourceCompatibility = JavaVersion.toVersion(25); targetCompatibility = JavaVersion.toVersion(25) } -// no tests for now -testingConventions.enabled = false - dependencies { compileOnly project(':server') + testImplementation project(':test:framework') } diff --git a/sandbox/libs/plugin-stats-spi/src/main/java/org/opensearch/plugin/stats/StatsRecorder.java b/sandbox/libs/plugin-stats-spi/src/main/java/org/opensearch/plugin/stats/StatsRecorder.java index 249eaa65851e6..1ff03dc361da5 100644 --- a/sandbox/libs/plugin-stats-spi/src/main/java/org/opensearch/plugin/stats/StatsRecorder.java +++ b/sandbox/libs/plugin-stats-spi/src/main/java/org/opensearch/plugin/stats/StatsRecorder.java @@ -72,6 +72,22 @@ public static void recordTimeMillis(IORunnable work, LongConsumer timeRecorder) } } + /** + * Value-returning variant of {@link #recordTimeMillis(IORunnable, LongConsumer)}. + * Runs {@code work}, returns its result, and reports elapsed wall-clock millis to + * {@code timeRecorder}. Time is reported even if {@code work} throws. + */ + public static T recordTimeMillis(IOSupplier work, LongConsumer timeRecorder) throws IOException { + Objects.requireNonNull(work, "work"); + Objects.requireNonNull(timeRecorder, "timeRecorder"); + long start = System.nanoTime(); + try { + return work.get(); + } finally { + timeRecorder.accept(TimeUnit.NANOSECONDS.toMillis(System.nanoTime() - start)); + } + } + /** * Runs {@code work} and reports time + outcome: *

    diff --git a/sandbox/libs/plugin-stats-spi/src/test/java/org/opensearch/plugin/stats/StatsRecorderTests.java b/sandbox/libs/plugin-stats-spi/src/test/java/org/opensearch/plugin/stats/StatsRecorderTests.java new file mode 100644 index 0000000000000..3284753213d79 --- /dev/null +++ b/sandbox/libs/plugin-stats-spi/src/test/java/org/opensearch/plugin/stats/StatsRecorderTests.java @@ -0,0 +1,247 @@ +/* + * SPDX-License-Identifier: Apache-2.0 + * + * The OpenSearch Contributors require contributions made to + * this file be licensed under the Apache-2.0 license or a + * compatible open source license. + */ + +package org.opensearch.plugin.stats; + +import org.opensearch.test.OpenSearchTestCase; + +import java.io.IOException; +import java.util.concurrent.atomic.AtomicBoolean; +import java.util.concurrent.atomic.AtomicInteger; +import java.util.concurrent.atomic.AtomicLong; +import java.util.concurrent.atomic.AtomicReference; + +/** + * Unit tests for {@link StatsRecorder}. + * + *

    Covers all four public methods and their contracts: + *

      + *
    • time is always recorded (success and failure paths)
      • + *
    • {@code onSuccess}/{@code onFailure} are mutually exclusive
      • + *
    • the original exception propagates unchanged
      • + *
    • return values pass through
      • + *
    • null arguments are rejected
      • + *
    + */ +public class StatsRecorderTests extends OpenSearchTestCase { + + // ────────────────────────────────────────────────────────────────────── + // recordTimeMillis(IORunnable, ...) + // ────────────────────────────────────────────────────────────────────── + + public void testRecordTimeMillisVoidRunsWorkAndRecordsTime() throws IOException { + AtomicBoolean ran = new AtomicBoolean(false); + AtomicLong recorded = new AtomicLong(-1); + + StatsRecorder.recordTimeMillis(() -> ran.set(true), recorded::set); + + assertTrue("work should have run", ran.get()); + assertTrue("elapsed millis should be recorded (>= 0)", recorded.get() >= 0L); + } + + public void testRecordTimeMillisVoidRecordsTimeEvenOnThrow() { + AtomicLong recorded = new AtomicLong(-1); + IOException boom = new IOException("boom"); + + IOException thrown = expectThrows( + IOException.class, + () -> StatsRecorder.recordTimeMillis((StatsRecorder.IORunnable) () -> { throw boom; }, recorded::set) + ); + + assertSame("the original exception must propagate", boom, thrown); + assertTrue("elapsed millis must be recorded even when work throws", recorded.get() >= 0L); + } + + public void testRecordTimeMillisVoidRejectsNullArgs() { + expectThrows(NullPointerException.class, () -> StatsRecorder.recordTimeMillis((StatsRecorder.IORunnable) null, v -> {})); + expectThrows(NullPointerException.class, () -> StatsRecorder.recordTimeMillis(() -> {}, null)); + } + + // ────────────────────────────────────────────────────────────────────── + // recordTimeMillis(IOSupplier, ...) — value-returning + // ────────────────────────────────────────────────────────────────────── + + public void testRecordTimeMillisValueReturnsResultAndRecordsTime() throws IOException { + AtomicLong recorded = new AtomicLong(-1); + + String result = StatsRecorder.recordTimeMillis(() -> "hello", recorded::set); + + assertEquals("hello", result); + assertTrue("elapsed millis should be recorded", recorded.get() >= 0L); + } + + public void testRecordTimeMillisValuePassesThroughNullResult() throws IOException { + AtomicLong recorded = new AtomicLong(-1); + + String result = StatsRecorder.recordTimeMillis(() -> null, recorded::set); + + assertNull("null result must pass through", result); + assertTrue(recorded.get() >= 0L); + } + + public void testRecordTimeMillisValueRecordsTimeEvenOnThrow() { + AtomicLong recorded = new AtomicLong(-1); + IOException boom = new IOException("boom"); + + IOException thrown = expectThrows(IOException.class, () -> StatsRecorder.recordTimeMillis((StatsRecorder.IOSupplier) () -> { + throw boom; + }, recorded::set)); + + assertSame(boom, thrown); + assertTrue("elapsed millis must be recorded even when work throws", recorded.get() >= 0L); + } + + public void testRecordTimeMillisValueRejectsNullArgs() { + expectThrows(NullPointerException.class, () -> StatsRecorder.recordTimeMillis((StatsRecorder.IOSupplier) null, v -> {})); + expectThrows(NullPointerException.class, () -> StatsRecorder.recordTimeMillis(() -> "x", null)); + } + + // ────────────────────────────────────────────────────────────────────── + // recordOutcome(IOSupplier, ...) — value-returning + // ────────────────────────────────────────────────────────────────────── + + public void testRecordOutcomeValueSuccessFiresOnSuccessOnly() throws IOException { + AtomicLong recorded = new AtomicLong(-1); + AtomicInteger success = new AtomicInteger(0); + AtomicInteger failure = new AtomicInteger(0); + + String result = StatsRecorder.recordOutcome(() -> "ok", recorded::set, success::incrementAndGet, failure::incrementAndGet); + + assertEquals("ok", result); + assertEquals("onSuccess must fire exactly once", 1, success.get()); + assertEquals("onFailure must not fire on success", 0, failure.get()); + assertTrue("time must be recorded", recorded.get() >= 0L); + } + + public void testRecordOutcomeValueFailureFiresOnFailureOnly() { + AtomicLong recorded = new AtomicLong(-1); + AtomicInteger success = new AtomicInteger(0); + AtomicInteger failure = new AtomicInteger(0); + IOException boom = new IOException("boom"); + + IOException thrown = expectThrows(IOException.class, () -> StatsRecorder.recordOutcome((StatsRecorder.IOSupplier) () -> { + throw boom; + }, recorded::set, success::incrementAndGet, failure::incrementAndGet)); + + assertSame("original exception must propagate", boom, thrown); + assertEquals("onFailure must fire exactly once", 1, failure.get()); + assertEquals("onSuccess must not fire on failure", 0, success.get()); + assertTrue("time must be recorded even on failure", recorded.get() >= 0L); + } + + public void testRecordOutcomeValuePropagatesRuntimeException() { + AtomicInteger success = new AtomicInteger(0); + AtomicInteger failure = new AtomicInteger(0); + RuntimeException boom = new IllegalStateException("rt-boom"); + + RuntimeException thrown = expectThrows( + IllegalStateException.class, + () -> StatsRecorder.recordOutcome((StatsRecorder.IOSupplier) () -> { + throw boom; + }, v -> {}, success::incrementAndGet, failure::incrementAndGet) + ); + + assertSame(boom, thrown); + assertEquals("onFailure fires on RuntimeException too", 1, failure.get()); + assertEquals(0, success.get()); + } + + public void testRecordOutcomeValueRejectsNullArgs() { + expectThrows( + NullPointerException.class, + () -> StatsRecorder.recordOutcome((StatsRecorder.IOSupplier) null, v -> {}, () -> {}, () -> {}) + ); + expectThrows(NullPointerException.class, () -> StatsRecorder.recordOutcome(() -> "x", null, () -> {}, () -> {})); + expectThrows(NullPointerException.class, () -> StatsRecorder.recordOutcome(() -> "x", v -> {}, null, () -> {})); + expectThrows(NullPointerException.class, () -> StatsRecorder.recordOutcome(() -> "x", v -> {}, () -> {}, null)); + } + + // ────────────────────────────────────────────────────────────────────── + // recordOutcome(IORunnable, ...) — void + // ────────────────────────────────────────────────────────────────────── + + public void testRecordOutcomeVoidSuccessFiresOnSuccessOnly() throws IOException { + AtomicBoolean ran = new AtomicBoolean(false); + AtomicLong recorded = new AtomicLong(-1); + AtomicInteger success = new AtomicInteger(0); + AtomicInteger failure = new AtomicInteger(0); + + StatsRecorder.recordOutcome(() -> ran.set(true), recorded::set, success::incrementAndGet, failure::incrementAndGet); + + assertTrue("work should have run", ran.get()); + assertEquals(1, success.get()); + assertEquals(0, failure.get()); + assertTrue(recorded.get() >= 0L); + } + + public void testRecordOutcomeVoidFailureFiresOnFailureOnly() { + AtomicLong recorded = new AtomicLong(-1); + AtomicInteger success = new AtomicInteger(0); + AtomicInteger failure = new AtomicInteger(0); + IOException boom = new IOException("boom"); + + IOException thrown = expectThrows( + IOException.class, + () -> StatsRecorder.recordOutcome( + (StatsRecorder.IORunnable) () -> { throw boom; }, + recorded::set, + success::incrementAndGet, + failure::incrementAndGet + ) + ); + + assertSame(boom, thrown); + assertEquals(1, failure.get()); + assertEquals(0, success.get()); + assertTrue(recorded.get() >= 0L); + } + + public void testRecordOutcomeVoidRejectsNullWork() { + expectThrows( + NullPointerException.class, + () -> StatsRecorder.recordOutcome((StatsRecorder.IORunnable) null, v -> {}, () -> {}, () -> {}) + ); + } + + // ────────────────────────────────────────────────────────────────────── + // Ordering / interaction contracts + // ────────────────────────────────────────────────────────────────────── + + public void testTimeRecordedBeforeOnSuccessCallback() throws IOException { + // Contract: the time recorder runs in finally, before onSuccess is invoked. + AtomicReference order = new AtomicReference<>(""); + + StatsRecorder.recordOutcome( + () -> "v", + millis -> order.updateAndGet(s -> s + "T"), + () -> order.updateAndGet(s -> s + "S"), + () -> order.updateAndGet(s -> s + "F") + ); + + assertEquals("time (T) must be recorded before onSuccess (S)", "TS", order.get()); + } + + public void testOnFailureRunsBeforeExceptionPropagatesAndTimeStillRecorded() { + // Contract: onFailure runs inside catch (before rethrow); time recorder runs in finally. + AtomicReference order = new AtomicReference<>(""); + IOException boom = new IOException("boom"); + + expectThrows( + IOException.class, + () -> StatsRecorder.recordOutcome( + (StatsRecorder.IOSupplier) () -> { throw boom; }, + millis -> order.updateAndGet(s -> s + "T"), + () -> order.updateAndGet(s -> s + "S"), + () -> order.updateAndGet(s -> s + "F") + ) + ); + + // onFailure (F) fires in catch, then time (T) in finally; onSuccess (S) never fires. + assertEquals("expected onFailure then time, no onSuccess", "FT", order.get()); + } +} diff --git a/sandbox/plugins/analytics-backend-lucene/src/main/java/org/opensearch/be/lucene/index/LuceneIndexingExecutionEngine.java b/sandbox/plugins/analytics-backend-lucene/src/main/java/org/opensearch/be/lucene/index/LuceneIndexingExecutionEngine.java index b57cb84c379fe..ab99283d58c78 100644 --- a/sandbox/plugins/analytics-backend-lucene/src/main/java/org/opensearch/be/lucene/index/LuceneIndexingExecutionEngine.java +++ b/sandbox/plugins/analytics-backend-lucene/src/main/java/org/opensearch/be/lucene/index/LuceneIndexingExecutionEngine.java @@ -45,6 +45,7 @@ import org.opensearch.index.mapper.MapperService; import org.opensearch.index.store.Store; import org.opensearch.plugin.stats.DataFormatStatsProviderRegistry; +import org.opensearch.plugin.stats.StatsRecorder; import java.io.IOException; import java.nio.file.Files; @@ -328,15 +329,15 @@ public RefreshResult refresh(RefreshInput refreshInput) throws IOException { // Single batched addIndexes call for all source directories if (sourceDirectories.isEmpty() == false) { - long addIndexesStart = System.nanoTime(); try { - sharedWriter.addIndexes(sourceDirectories.toArray(new Directory[0])); - logger.debug( - "Incorporated {} Lucene segments into shared writer in a single addIndexes call", - sourceDirectories.size() - ); + StatsRecorder.recordTimeMillis(() -> { + sharedWriter.addIndexes(sourceDirectories.toArray(new Directory[0])); + logger.debug( + "Incorporated {} Lucene segments into shared writer in a single addIndexes call", + sourceDirectories.size() + ); + }, stats::addRefreshAddIndexesTimeMillis); } finally { - stats.addRefreshAddIndexesTimeMillis(TimeUnit.NANOSECONDS.toMillis(System.nanoTime() - addIndexesStart)); // Close all source directories for (Directory dir : sourceDirectories) { try { diff --git a/sandbox/plugins/analytics-backend-lucene/src/main/java/org/opensearch/be/lucene/index/LuceneWriter.java b/sandbox/plugins/analytics-backend-lucene/src/main/java/org/opensearch/be/lucene/index/LuceneWriter.java index caa46b006d7f5..77d92c7774256 100644 --- a/sandbox/plugins/analytics-backend-lucene/src/main/java/org/opensearch/be/lucene/index/LuceneWriter.java +++ b/sandbox/plugins/analytics-backend-lucene/src/main/java/org/opensearch/be/lucene/index/LuceneWriter.java @@ -62,7 +62,6 @@ import java.util.Optional; import java.util.Set; import java.util.concurrent.Executor; -import java.util.concurrent.TimeUnit; /** * Per-generation Lucene writer that creates segments in an isolated temporary directory. @@ -225,8 +224,7 @@ protected int maxBufferedDocs() { */ @Override public WriteResult addDoc(LuceneDocumentInput input) throws IOException { - long start = System.nanoTime(); - try { + return StatsRecorder.recordTimeMillis(() -> { if (state != WriterState.ACTIVE) { throw new IllegalStateException("addDoc requires ACTIVE state but was " + state); } @@ -250,9 +248,7 @@ public WriteResult addDoc(LuceneDocumentInput input) throws IOException { docCount++; stats.addDocsIndexed(1); return new WriteResult.Success(1L, 1L, currentDocId); - } finally { - stats.addIndexTimeMillis(TimeUnit.NANOSECONDS.toMillis(System.nanoTime() - start)); - } + }, stats::addIndexTimeMillis); } @Override @@ -307,9 +303,9 @@ public FileInfos flush(FlushInput flushInput) throws IOException { } long flushStart = System.nanoTime(); + long totalFlushDurationMs = 0; try { - long flushStartNanos = System.nanoTime(); - logger.info( + logger.debug( "flush: START generation={}, docCount={}, hasRowIdMapping={}", writerGeneration, docCount, @@ -359,7 +355,7 @@ public FileInfos flush(FlushInput flushInput) throws IOException { long forceMergeStartNanos = System.nanoTime(); StatsRecorder.recordTimeMillis(() -> indexWriter.forceMerge(1, true), stats::addFlushForceMergeTimeMillis); long forceMergeDurationMs = TimeValue.nsecToMSec(System.nanoTime() - forceMergeStartNanos); - logger.info( + logger.debug( "flush: forceMerge complete: generation={}, docCount={}, duration={}ms", writerGeneration, docCount, @@ -369,7 +365,7 @@ public FileInfos flush(FlushInput flushInput) throws IOException { long commitStartNanos = System.nanoTime(); indexWriter.commit(); long commitDurationMs = TimeValue.nsecToMSec(System.nanoTime() - commitStartNanos); - logger.info("flush: commit complete: generation={}, duration={}ms", writerGeneration, commitDurationMs); + logger.debug("flush: commit complete: generation={}, duration={}ms", writerGeneration, commitDurationMs); // Close the IndexWriter before rewriting segment metadata. // This prevents IndexFileDeleter from removing our rewritten segments_N @@ -426,8 +422,8 @@ assert assertRowIdsSequential(directory) : "___row_id__ doc values not sequentia directory.close(); flushed = true; - long totalFlushDurationMs = TimeValue.nsecToMSec(System.nanoTime() - flushStartNanos); - logger.info( + totalFlushDurationMs = TimeValue.nsecToMSec(System.nanoTime() - flushStart); + logger.debug( "flush: DONE generation={}, totalRows={}, forceMerge={}ms, commit={}ms, total={}ms", writerGeneration, docCount, @@ -439,7 +435,7 @@ assert assertRowIdsSequential(directory) : "___row_id__ doc values not sequentia return FileInfos.builder().putWriterFileSet(dataFormat, wfsBuilder.build()).build(); } finally { stats.incFlushTotal(); - stats.addFlushTimeMillis(TimeUnit.NANOSECONDS.toMillis(System.nanoTime() - flushStart)); + stats.addFlushTimeMillis(totalFlushDurationMs); } } diff --git a/sandbox/plugins/analytics-backend-lucene/src/main/java/org/opensearch/be/lucene/stats/LuceneStatsProvider.java b/sandbox/plugins/analytics-backend-lucene/src/main/java/org/opensearch/be/lucene/stats/LuceneStatsProvider.java index d18515f5d31d0..d76c9dc275c7f 100644 --- a/sandbox/plugins/analytics-backend-lucene/src/main/java/org/opensearch/be/lucene/stats/LuceneStatsProvider.java +++ b/sandbox/plugins/analytics-backend-lucene/src/main/java/org/opensearch/be/lucene/stats/LuceneStatsProvider.java @@ -8,6 +8,7 @@ package org.opensearch.be.lucene.stats; +import org.opensearch.be.lucene.LuceneDataFormat; import org.opensearch.common.annotation.ExperimentalApi; import org.opensearch.core.common.io.stream.Writeable; import org.opensearch.core.index.shard.ShardId; @@ -35,7 +36,8 @@ @ExperimentalApi public final class LuceneStatsProvider implements DataFormatStatsProvider { - public static final String FORMAT_NAME = "lucene"; + /** Canonical format name, sourced from {@link LuceneDataFormat} to avoid duplicating the literal. */ + public static final String FORMAT_NAME = LuceneDataFormat.LUCENE_FORMAT_NAME; private static volatile LuceneStatsProvider INSTANCE; diff --git a/sandbox/plugins/composite-engine/src/internalClusterTest/java/org/opensearch/composite/StatsFailureIT.java b/sandbox/plugins/composite-engine/src/internalClusterTest/java/org/opensearch/composite/StatsFailureIT.java index 87932868d4c25..58671a1aed1d1 100644 --- a/sandbox/plugins/composite-engine/src/internalClusterTest/java/org/opensearch/composite/StatsFailureIT.java +++ b/sandbox/plugins/composite-engine/src/internalClusterTest/java/org/opensearch/composite/StatsFailureIT.java @@ -80,6 +80,11 @@ public void tearDown() throws Exception { * (proving the old tracker was properly unregistered and a fresh one registered). * 4. New indexing after recovery is tracked correctly. */ + @AwaitsFix(bugUrl = "Intermittent deadlock: after the injected engine failure, a scheduled refresh thread parks " + + "permanently in LockablePool.checkoutAll (blocking ReentrantLock.lock with no timeout) because a " + + "writer lock is left held on the engine-failure path. The shard never recovers, so ensureGreen times " + + "out. Root cause is in DataFormatAwareEngine.refresh / LockablePool (server), not in stats. " + + "Re-enable once checkoutAll uses a bounded tryLock or the failure path releases writer locks.") public void testTrackerLifecycleAroundEngineFailure() throws Exception { // Start a single node to minimize interference with the JVM-wide failure countdown. internalCluster().startNode(); diff --git a/sandbox/plugins/parquet-data-format/src/main/java/org/opensearch/parquet/bridge/NativeParquetWriter.java b/sandbox/plugins/parquet-data-format/src/main/java/org/opensearch/parquet/bridge/NativeParquetWriter.java index 30a0e245d10b6..ee18effdb070a 100644 --- a/sandbox/plugins/parquet-data-format/src/main/java/org/opensearch/parquet/bridge/NativeParquetWriter.java +++ b/sandbox/plugins/parquet-data-format/src/main/java/org/opensearch/parquet/bridge/NativeParquetWriter.java @@ -14,7 +14,6 @@ import org.opensearch.plugin.stats.StatsRecorder; import java.io.IOException; -import java.util.concurrent.TimeUnit; import java.util.concurrent.atomic.AtomicBoolean; /** @@ -123,8 +122,7 @@ public void write(long arrayAddress, long schemaAddress) throws IOException { public ParquetFileMetadata flush() throws IOException { if (writerFlushed.compareAndSet(false, true)) { if (initialized) { - long startNanos = System.nanoTime(); - try { + StatsRecorder.recordOutcome(() -> { RustBridge.WriterFinalizeResult result = RustBridge.finalizeWriter(filePath); if (result != null) { metadata.set(result.metadata()); @@ -132,14 +130,7 @@ public ParquetFileMetadata flush() throws IOException { rowIdMapping.set(result.rowIdMapping()); } } - stats.incNativeFinalizeTotal(); - } catch (IOException e) { - stats.incNativeFinalizeFailures(); - throw e; - } finally { - long elapsed = TimeUnit.NANOSECONDS.toMillis(System.nanoTime() - startNanos); - stats.addNativeFinalizeTimeMillis(elapsed); - } + }, stats::addNativeFinalizeTimeMillis, stats::incNativeFinalizeTotal, stats::incNativeFinalizeFailures); } } return metadata.get(); diff --git a/sandbox/plugins/parquet-data-format/src/main/java/org/opensearch/parquet/stats/ParquetStatsProvider.java b/sandbox/plugins/parquet-data-format/src/main/java/org/opensearch/parquet/stats/ParquetStatsProvider.java index 02984fc21ed38..4357b7919bf05 100644 --- a/sandbox/plugins/parquet-data-format/src/main/java/org/opensearch/parquet/stats/ParquetStatsProvider.java +++ b/sandbox/plugins/parquet-data-format/src/main/java/org/opensearch/parquet/stats/ParquetStatsProvider.java @@ -13,6 +13,7 @@ import org.opensearch.common.annotation.ExperimentalApi; import org.opensearch.core.common.io.stream.Writeable; import org.opensearch.core.index.shard.ShardId; +import org.opensearch.parquet.engine.ParquetDataFormat; import org.opensearch.plugin.stats.DataFormatStatsProvider; import org.opensearch.plugin.stats.DataFormatStatsProviderRegistry; @@ -37,7 +38,8 @@ @ExperimentalApi public final class ParquetStatsProvider implements DataFormatStatsProvider { - public static final String FORMAT_NAME = "parquet"; + /** Canonical format name, sourced from {@link ParquetDataFormat} to avoid duplicating the literal. */ + public static final String FORMAT_NAME = ParquetDataFormat.PARQUET_DATA_FORMAT_NAME; private static final Logger logger = LogManager.getLogger(ParquetStatsProvider.class); diff --git a/sandbox/plugins/parquet-data-format/src/main/java/org/opensearch/parquet/writer/ParquetWriter.java b/sandbox/plugins/parquet-data-format/src/main/java/org/opensearch/parquet/writer/ParquetWriter.java index c342836f99dd2..24736d59524bb 100644 --- a/sandbox/plugins/parquet-data-format/src/main/java/org/opensearch/parquet/writer/ParquetWriter.java +++ b/sandbox/plugins/parquet-data-format/src/main/java/org/opensearch/parquet/writer/ParquetWriter.java @@ -27,11 +27,11 @@ import org.opensearch.parquet.memory.ArrowBufferPool; import org.opensearch.parquet.stats.ParquetShardStatsTracker; import org.opensearch.parquet.vsr.VSRManager; +import org.opensearch.plugin.stats.StatsRecorder; import org.opensearch.threadpool.ThreadPool; import java.io.IOException; import java.nio.file.Path; -import java.util.concurrent.TimeUnit; import java.util.function.Supplier; /** @@ -146,8 +146,7 @@ public WriteResult addDoc(ParquetDocumentInput d) throws IOException { if (state != WriterState.ACTIVE) { throw new IllegalStateException("Writer is not active, state=" + state); } - long startNanos = System.nanoTime(); - try { + return StatsRecorder.recordTimeMillis(() -> { // Schema mismatch is recoverable: the VSR rejected the doc pre-admission, so the // caller-driven rollback no-ops in the VSR and restores ACTIVE. try { @@ -159,10 +158,7 @@ public WriteResult addDoc(ParquetDocumentInput d) throws IOException { acceptedRows++; stats.addDocsIndexed(1); return new WriteResult.Success(1L, 1L, 1L); - } finally { - long elapsed = TimeUnit.NANOSECONDS.toMillis(System.nanoTime() - startNanos); - stats.addIndexTimeMillis(elapsed); - } + }, stats::addIndexTimeMillis); } @Override From e296762ba1c7739670a350dce4415c0ce50eee2c Mon Sep 17 00:00:00 2001 From: Mohit Kumar Date: Thu, 11 Jun 2026 20:50:48 +0530 Subject: [PATCH 04/14] Fix field resolution for dotted fields under disable_objects (#21929) * Fix field resolution for dotted fields under disable_objects --------- Signed-off-by: Mohit Kumar Signed-off-by: Mohit Kumar Co-authored-by: Mohit Kumar --- .../index/mapper/DocumentParser.java | 37 ++++- .../index/mapper/MappingLookup.java | 26 +++- .../opensearch/index/mapper/ObjectMapper.java | 8 + .../index/mapper/RootObjectMapper.java | 5 + .../index/mapper/DocumentParserTests.java | 147 ++++++++++++++++++ .../index/mapper/ParametrizedMapperTests.java | 41 +++++ 6 files changed, 252 insertions(+), 12 deletions(-) diff --git a/server/src/main/java/org/opensearch/index/mapper/DocumentParser.java b/server/src/main/java/org/opensearch/index/mapper/DocumentParser.java index f0b2a4f3c8368..ea72082fcb355 100644 --- a/server/src/main/java/org/opensearch/index/mapper/DocumentParser.java +++ b/server/src/main/java/org/opensearch/index/mapper/DocumentParser.java @@ -330,13 +330,25 @@ private static boolean shouldHandleAsDisableObjects(DocumentMapper docMapper, St /** * Handles flat field mapping by adding the mapper directly to the disable_objects parent. * Only FieldMappers are added; ObjectMappers are skipped as they conflict with disable_objects. + * If parentMappers is empty (first mapper case), the update is added directly; + * otherwise it is merged into the existing root entry. */ - private static void handleDisableObjectsMapping(List parentMappers, Mapper newMapper, DocumentMapper docMapper) { + private static void handleDisableObjectsMapping( + List parentMappers, + Mapper newMapper, + DocumentMapper docMapper, + Mapping mapping + ) { // If the mapper is an ObjectMapper, we cannot add it to a disable_objects parent // This can happen when intermediate object mappers are created during dynamic mapping // In disable_objects mode, we only want FieldMappers with dotted names if (newMapper instanceof ObjectMapper) { - // Skip ObjectMappers - they will be handled when their leaf FieldMappers are processed + if (parentMappers.isEmpty()) { + // For the first mapper case, seed with a root that contains just this ObjectMapper. + // ObjectMappers are skipped under disable_objects (their leaf FieldMappers handle it), + // but parentMappers needs an entry for subsequent merges. + parentMappers.add(createUpdate(mapping.root(), splitAndValidatePath(newMapper.name()), 0, newMapper)); + } return; } @@ -366,13 +378,11 @@ private static void handleDisableObjectsMapping(List parentMappers pathMappers.add(om); current = om; } else { - // Shouldn't happen if mapping exists break; } } // Build the update from the disable_objects parent back up to root - // Add the FieldMapper to the disable_objects parent (last in pathMappers) ObjectMapper disableObjectsParent = pathMappers.get(pathMappers.size() - 1); ObjectMapper update = disableObjectsParent.mappingUpdate(newMapper); @@ -381,8 +391,13 @@ private static void handleDisableObjectsMapping(List parentMappers update = pathMappers.get(i).mappingUpdate(update); } - // Merge with existing root update - parentMappers.set(0, parentMappers.get(0).merge(update)); + if (parentMappers.isEmpty()) { + // First mapper — add directly + parentMappers.add(update); + } else { + // Subsequent mapper — merge with existing root update + parentMappers.set(0, parentMappers.get(0).merge(update)); + } } /** Creates a Mapping containing any dynamically added fields, or returns null if there were no dynamic mappings. */ @@ -397,7 +412,13 @@ static Mapping createDynamicUpdate(Mapping mapping, DocumentMapper docMapper, Li Iterator dynamicMapperItr = dynamicMappers.iterator(); List parentMappers = new ArrayList<>(); Mapper firstUpdate = dynamicMapperItr.next(); - parentMappers.add(createUpdate(mapping.root(), splitAndValidatePath(firstUpdate.name()), 0, firstUpdate)); + String[] firstNameParts = splitAndValidatePath(firstUpdate.name()); + // Check if the first mapper should be handled as a disable_objects literal field + if (shouldHandleAsDisableObjects(docMapper, firstNameParts)) { + handleDisableObjectsMapping(parentMappers, firstUpdate, docMapper, mapping); + } else { + parentMappers.add(createUpdate(mapping.root(), firstNameParts, 0, firstUpdate)); + } Mapper previousMapper = null; while (dynamicMapperItr.hasNext()) { Mapper newMapper = dynamicMapperItr.next(); @@ -413,7 +434,7 @@ static Mapping createDynamicUpdate(Mapping mapping, DocumentMapper docMapper, Li // Check if this field should be handled as literal field if (shouldHandleAsDisableObjects(docMapper, nameParts)) { - handleDisableObjectsMapping(parentMappers, newMapper, docMapper); + handleDisableObjectsMapping(parentMappers, newMapper, docMapper, mapping); continue; // Skip the normal processing for this mapper } diff --git a/server/src/main/java/org/opensearch/index/mapper/MappingLookup.java b/server/src/main/java/org/opensearch/index/mapper/MappingLookup.java index e5498957b9926..da7eed023bb2c 100644 --- a/server/src/main/java/org/opensearch/index/mapper/MappingLookup.java +++ b/server/src/main/java/org/opensearch/index/mapper/MappingLookup.java @@ -59,6 +59,7 @@ public final class MappingLookup implements Iterable { private final Map fieldMappers; private final Map objectMappers; private final boolean hasNested; + private final boolean rootDisableObjects; private final FieldTypeLookup fieldTypeLookup; private final int metadataFieldCount; private final FieldNameAnalyzer indexAnalyzer; @@ -95,7 +96,8 @@ public static MappingLookup fromMapping(Mapping mapping, Analyzer defaultIndex, newFieldAliasMappers, mapping.metadataMappers.length, defaultIndex, - dynamicResolver + dynamicResolver, + mapping.root().disableObjects() ); } @@ -133,7 +135,7 @@ public MappingLookup( int metadataFieldCount, Analyzer defaultIndex ) { - this(mappers, objectMappers, aliasMappers, metadataFieldCount, defaultIndex, null); + this(mappers, objectMappers, aliasMappers, metadataFieldCount, defaultIndex, null, false); } MappingLookup( @@ -142,7 +144,8 @@ public MappingLookup( Collection aliasMappers, int metadataFieldCount, Analyzer defaultIndex, - DynamicPropertyFieldTypeResolver dynamicPropertyFieldTypes + DynamicPropertyFieldTypeResolver dynamicPropertyFieldTypes, + boolean rootDisableObjects ) { Map fieldMappers = new HashMap<>(); Map indexAnalyzers = new HashMap<>(); @@ -158,6 +161,7 @@ public MappingLookup( } } this.hasNested = hasNested; + this.rootDisableObjects = rootDisableObjects; for (FieldMapper mapper : mappers) { if (objects.containsKey(mapper.name())) { @@ -278,7 +282,21 @@ public Map objectMappers() { public boolean isMultiField(String field) { String sourceParent = parentObject(field); - return sourceParent != null && fieldMappers.containsKey(sourceParent); + if (sourceParent == null || !fieldMappers.containsKey(sourceParent)) { + return false; + } + // When disable_objects is true, dotted fields are independent flat fields, not multi-fields. + // Check if the root or any ancestor ObjectMapper has disable_objects enabled. + if (rootDisableObjects) { + return false; + } + for (String ancestor = parentObject(field); ancestor != null; ancestor = parentObject(ancestor)) { + ObjectMapper objectMapper = objectMappers.get(ancestor); + if (objectMapper != null && objectMapper.disableObjects()) { + return false; + } + } + return true; } public boolean isObjectField(String field) { diff --git a/server/src/main/java/org/opensearch/index/mapper/ObjectMapper.java b/server/src/main/java/org/opensearch/index/mapper/ObjectMapper.java index f1bbcf59190f0..587265d98675e 100644 --- a/server/src/main/java/org/opensearch/index/mapper/ObjectMapper.java +++ b/server/src/main/java/org/opensearch/index/mapper/ObjectMapper.java @@ -319,6 +319,14 @@ public static class TypeParser implements Mapper.TypeParser { public Mapper.Builder parse(String name, Map node, ParserContext parserContext) throws MapperParsingException { ObjectMapper.Builder builder = new Builder(name); parseNested(name, node, builder, parserContext); + // Parse disable_objects first so it's available when parseProperties is called. + // Map iteration order is not guaranteed, so disable_objects might come after properties + // in the map. Without this, parseProperties won't know to treat dotted field names as + // literal names and will incorrectly split them into intermediate object mappers. + Object disableObjectsNode = node.remove("disable_objects"); + if (disableObjectsNode != null) { + parseObjectOrDocumentTypeProperties("disable_objects", disableObjectsNode, parserContext, builder); + } Object compositeField = null; for (Iterator> iterator = node.entrySet().iterator(); iterator.hasNext();) { Map.Entry entry = iterator.next(); diff --git a/server/src/main/java/org/opensearch/index/mapper/RootObjectMapper.java b/server/src/main/java/org/opensearch/index/mapper/RootObjectMapper.java index 609f4b25d6840..6a6531c7b4213 100644 --- a/server/src/main/java/org/opensearch/index/mapper/RootObjectMapper.java +++ b/server/src/main/java/org/opensearch/index/mapper/RootObjectMapper.java @@ -195,6 +195,11 @@ public static class TypeParser extends ObjectMapper.TypeParser { public Mapper.Builder parse(String name, Map node, ParserContext parserContext) throws MapperParsingException { RootObjectMapper.Builder builder = new Builder(name); + // Parse disable_objects first so it's available when parseProperties is called. + Object disableObjectsNode = node.remove("disable_objects"); + if (disableObjectsNode != null) { + parseObjectOrDocumentTypeProperties("disable_objects", disableObjectsNode, parserContext, builder); + } Iterator> iterator = node.entrySet().iterator(); Object compositeField = null; Object contextAwareGoupingField = null; diff --git a/server/src/test/java/org/opensearch/index/mapper/DocumentParserTests.java b/server/src/test/java/org/opensearch/index/mapper/DocumentParserTests.java index 3013127110925..3bfee966b7d78 100644 --- a/server/src/test/java/org/opensearch/index/mapper/DocumentParserTests.java +++ b/server/src/test/java/org/opensearch/index/mapper/DocumentParserTests.java @@ -3639,6 +3639,153 @@ public void testParseDocumentWithDocumentInputAndNestedFields() throws Exception assertNotNull(doc.rootDoc().getField("obj.inner")); } + public void testDisableObjectsFieldTypeLookupWithPrefixConflict() throws Exception { + MapperService mapperService = createMapperService(topMapping(b -> { + b.field("dynamic", false); + b.startObject("properties"); + b.startObject("attributes").field("type", "object").field("dynamic", true).field("disable_objects", true).endObject(); + b.endObject(); + })); + + ParsedDocument doc1 = mapperService.documentMapper().parse(source(""" + {"attributes": {"address.city": "Austin", "address.state": "Texas"}} + """)); + // Dynamic mapping update expected since "address.city" and "address.state" are not pre-defined + assertNotNull("doc1 should produce dynamic mapping update", doc1.dynamicMappingsUpdate()); + merge(mapperService, dynamicMapping(doc1.dynamicMappingsUpdate())); + assertEquals("address.city should be text", "text", mapperService.fieldType("attributes.address.city").typeName()); + + ParsedDocument doc2 = mapperService.documentMapper().parse(source(""" + {"attributes": {"address": "US"}} + """)); + // Dynamic mapping update expected since "address" is not pre-defined + assertNotNull("doc2 should produce dynamic mapping update", doc2.dynamicMappingsUpdate()); + merge(mapperService, dynamicMapping(doc2.dynamicMappingsUpdate())); + + assertEquals("address should be text", "text", mapperService.fieldType("attributes.address").typeName()); + assertEquals( + "address.city should be text after prefix conflict", + "text", + mapperService.fieldType("attributes.address.city").typeName() + ); + assertEquals( + "address.state should be text after prefix conflict", + "text", + mapperService.fieldType("attributes.address.state").typeName() + ); + + assertNull(mapperService.fieldType("attributes.address.address.city")); + + MapperService ms2 = createMapperService(topMapping(b -> { + b.field("dynamic", false); + b.startObject("properties"); + b.startObject("attributes").field("type", "object").field("dynamic", true).field("disable_objects", true).endObject(); + b.endObject(); + })); + + ParsedDocument docA = ms2.documentMapper().parse(source(""" + {"attributes": {"address": "US"}} + """)); + // Dynamic mapping update expected since "address" is not pre-defined + assertNotNull("docA should produce dynamic mapping update", docA.dynamicMappingsUpdate()); + merge(ms2, dynamicMapping(docA.dynamicMappingsUpdate())); + + ParsedDocument docB = ms2.documentMapper().parse(source(""" + {"attributes": {"address.city": "Austin"}} + """)); + // Dynamic mapping update expected since "address.city" is not pre-defined + assertNotNull("docB should produce dynamic mapping update", docB.dynamicMappingsUpdate()); + merge(ms2, dynamicMapping(docB.dynamicMappingsUpdate())); + + assertEquals("address.city should be text in reverse order", "text", ms2.fieldType("attributes.address.city").typeName()); + assertEquals("address should be text in reverse order", "text", ms2.fieldType("attributes.address").typeName()); + } + + public void testRootLevelDisableObjectsPrefixConflict() throws Exception { + // Root-level disable_objects with dynamic fields that have prefix conflicts + MapperService mapperService = createMapperService(topMapping(b -> { + b.field("disable_objects", true); + b.field("dynamic", true); + })); + + // Dynamically add "address.city" + ParsedDocument doc1 = mapperService.documentMapper().parse(source(""" + {"address.city": "Austin"} + """)); + assertNotNull("doc1 should produce dynamic mapping update", doc1.dynamicMappingsUpdate()); + merge(mapperService, dynamicMapping(doc1.dynamicMappingsUpdate())); + + // Dynamically add "address" + ParsedDocument doc2 = mapperService.documentMapper().parse(source(""" + {"address": "US"} + """)); + assertNotNull("doc2 should produce dynamic mapping update", doc2.dynamicMappingsUpdate()); + merge(mapperService, dynamicMapping(doc2.dynamicMappingsUpdate())); + + // Assert both fields are resolvable + assertEquals("address.city should be text", "text", mapperService.fieldType("address.city").typeName()); + assertEquals("address should be text", "text", mapperService.fieldType("address").typeName()); + + // Assert address.city is NOT treated as a multi-field of address + assertFalse( + "address.city should not be a multi-field of address", + mapperService.documentMapper().mappers().isMultiField("address.city") + ); + } + + public void testDisableObjectsMappingRecoveryWithPrefixConflict() throws Exception { + MapperService mapperService = createMapperService(topMapping(b -> { + b.field("dynamic", false); + b.startArray("dynamic_templates"); + b.startObject().startObject("strings").field("match_mapping_type", "string"); + b.startObject("mapping").field("type", "text").field("copy_to", "event_all"); + b.startObject("fields") + .startObject("keyword") + .field("type", "keyword") + .field("ignore_above", 256) + .endObject() + .endObject() + .endObject() + .endObject() + .endObject(); + b.endArray(); + b.startObject("properties"); + b.startObject("event_all").field("type", "text").endObject(); + b.startObject("attributes").field("type", "object").field("dynamic", true).field("disable_objects", true).endObject(); + b.endObject(); + })); + + ParsedDocument doc1 = mapperService.documentMapper().parse(source(""" + {"attributes": {"address.city": "Austin", "address.state": "Texas"}} + """)); + // Dynamic mapping update expected since "address.city" and "address.state" are not pre-defined + assertNotNull("doc1 should produce dynamic mapping update", doc1.dynamicMappingsUpdate()); + merge(mapperService, dynamicMapping(doc1.dynamicMappingsUpdate())); + + ParsedDocument doc2 = mapperService.documentMapper().parse(source(""" + {"attributes": {"address": "US"}} + """)); + // Dynamic mapping update expected since "address" is not pre-defined + assertNotNull("doc2 should produce dynamic mapping update", doc2.dynamicMappingsUpdate()); + merge(mapperService, dynamicMapping(doc2.dynamicMappingsUpdate())); + + String mappingSource = mapperService.documentMapper().mappingSource().string(); + MapperService recoveredService = createMapperService(MapperService.SINGLE_MAPPING_NAME, mappingSource); + + assertEquals( + "address.city should be text after round-trip", + "text", + recoveredService.fieldType("attributes.address.city").typeName() + ); + assertEquals( + "address.state should be text after round-trip", + "text", + recoveredService.fieldType("attributes.address.state").typeName() + ); + assertEquals("address should be text after round-trip", "text", recoveredService.fieldType("attributes.address").typeName()); + assertNull("no doubled name after round-trip", recoveredService.fieldType("attributes.address.address.city")); + } + @LockFeatureFlag(FeatureFlags.PLUGGABLE_DATAFORMAT_EXPERIMENTAL_FLAG) public void testDynamicTextFieldWithoutKeywordMultiFieldForPluggableDataFormat() throws Exception { Settings pluggableSettings = Settings.builder() diff --git a/server/src/test/java/org/opensearch/index/mapper/ParametrizedMapperTests.java b/server/src/test/java/org/opensearch/index/mapper/ParametrizedMapperTests.java index b113b773824c9..3c8dccf65e11a 100644 --- a/server/src/test/java/org/opensearch/index/mapper/ParametrizedMapperTests.java +++ b/server/src/test/java/org/opensearch/index/mapper/ParametrizedMapperTests.java @@ -554,4 +554,45 @@ public void testCustomMergeValidation() throws IOException { assertThat(e.getMessage(), containsString("int_value")); } + public void testMergePreservesFullNameForDottedSimpleName() { + Builder builder1 = new Builder("address.city"); + builder1.required.setValue("value"); + ContentPath path1 = new ContentPath(); + path1.add("attributes"); + TestMapper mapper1 = (TestMapper) builder1.build(new Mapper.BuilderContext(Settings.EMPTY, path1)); + + assertEquals("address.city", mapper1.simpleName()); + assertEquals("attributes.address.city", mapper1.name()); + + Builder builder2 = new Builder("address.city"); + builder2.required.setValue("updated"); + ContentPath path2 = new ContentPath(); + path2.add("attributes"); + TestMapper mapper2 = (TestMapper) builder2.build(new Mapper.BuilderContext(Settings.EMPTY, path2)); + + TestMapper merged = (TestMapper) mapper1.merge(mapper2); + + assertEquals("address.city", merged.simpleName()); + assertEquals("attributes.address.city", merged.name()); + assertEquals("updated", merged.required); + } + + public void testMergePreservesFullNameForDeeplyDottedSimpleName() { + Builder builder1 = new Builder("a.b.c"); + builder1.required.setValue("v1"); + ContentPath path = new ContentPath(); + path.add("root"); + TestMapper mapper1 = (TestMapper) builder1.build(new Mapper.BuilderContext(Settings.EMPTY, path)); + + assertEquals("root.a.b.c", mapper1.name()); + + Builder builder2 = new Builder("a.b.c"); + builder2.required.setValue("v2"); + TestMapper mapper2 = (TestMapper) builder2.build(new Mapper.BuilderContext(Settings.EMPTY, path)); + + TestMapper merged = (TestMapper) mapper1.merge(mapper2); + assertEquals("root.a.b.c", merged.name()); + assertEquals("a.b.c", merged.simpleName()); + } + } From deabd86893b80825a68d36e047d589cbbe110cd0 Mon Sep 17 00:00:00 2001 From: Finn Date: Thu, 11 Jun 2026 09:49:21 -0700 Subject: [PATCH 05/14] Add data node profiling via ArrowBatchResponse metadata channel (#21972) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Extends the profile API to return per-shard DataFusion execution metrics (output_rows, elapsed_compute, scan_time, row_groups_pruned, etc.) in the profile output as 'data_node_metrics' per task. Uses the ArrowBatchResponse metadata channel (merged in #22003) to transmit metrics in-band on the last data batch. No sentinel frames, no transport SPI changes, no threading races. Data node side: - AnalyticsSearchService extracts metrics after stream exhaustion - channelResponseHandler buffers one batch ahead; attaches metrics to the final batch via FragmentExecutionArrowResponse(root, metadata) Coordinator side: - handleStreamResponse reads last.getMetadata() after the stream loop - StreamingResponseListener.onStreamComplete(bytes) passes to task - ShardFragmentStageExecution stores on StageTask.setDataNodeMetrics - QueryProfileBuilder parses JSON into TaskProfile.dataNodeMetrics Also includes: - Profile flag propagation (QueryContext → FragmentExecutionRequest) - Rust FFM: df_stream_get_metrics extracts ExecutionPlan.metrics() - Debug logging of metrics at shard level (guarded by isDebugEnabled) - Integration test: ExplainApiIT.testExplainTasksHaveDataNodeMetrics Signed-off-by: Finn Carroll --- .../analytics/exec/profile/TaskProfile.java | 21 +- .../exec/AnalyticsSearchService.java | 22 +- .../exec/AnalyticsSearchTransportService.java | 62 +++++- .../analytics/exec/DefaultPlanExecutor.java | 7 +- .../analytics/exec/QueryContext.java | 49 ++--- .../exec/StreamingResponseListener.java | 7 + .../FragmentExecutionArrowResponse.java | 4 + .../exec/action/FragmentExecutionRequest.java | 12 ++ .../exec/profile/QueryProfileBuilder.java | 27 ++- .../analytics/exec/stage/StageTask.java | 11 + .../shard/ShardFragmentStageExecution.java | 10 +- .../ShardFragmentStageExecutionFactory.java | 3 +- .../exec/stage/shard/ShardTaskRunner.java | 2 +- .../ppl/action/RestPPLQueryAction.java | 35 ++- .../opensearch/analytics/qa/ExplainApiIT.java | 91 ++++++++ .../resources/datasets/delegation/bulk.json | 200 ++++++++++++++++++ .../datasets/delegation/mapping.json | 12 ++ 17 files changed, 492 insertions(+), 83 deletions(-) create mode 100644 sandbox/qa/analytics-engine-rest/src/test/resources/datasets/delegation/bulk.json create mode 100644 sandbox/qa/analytics-engine-rest/src/test/resources/datasets/delegation/mapping.json diff --git a/sandbox/libs/analytics-api/src/main/java/org/opensearch/analytics/exec/profile/TaskProfile.java b/sandbox/libs/analytics-api/src/main/java/org/opensearch/analytics/exec/profile/TaskProfile.java index 0efcb47a7d927..88f8bf7ff67be 100644 --- a/sandbox/libs/analytics-api/src/main/java/org/opensearch/analytics/exec/profile/TaskProfile.java +++ b/sandbox/libs/analytics-api/src/main/java/org/opensearch/analytics/exec/profile/TaskProfile.java @@ -12,17 +12,23 @@ import org.opensearch.core.xcontent.XContentBuilder; import java.io.IOException; +import java.util.Map; /** - * Per-task profile snapshot. Captures target node, terminal state and - * wall-clock elapsed time. A "task" is one dispatch unit within a stage - * (one shard for SOURCE, one partition for HASH_PARTITIONED, one total for COORDINATOR). + * Per-task profile snapshot. Captures target node, terminal state, + * wall-clock elapsed time, and optional data-node execution metrics. * * @param node target node and shard the task ran on, or "(unknown)" if dispatch never happened * @param state terminal state — CREATED if the task was never dispatched * @param elapsedMs wall-clock time from dispatch to terminal, or 0 if never dispatched + * @param dataNodeMetrics execution metrics from the data node (DataFusion operator timings), or null if not profiled */ -public record TaskProfile(String node, String state, long elapsedMs) implements ToXContentObject { +public record TaskProfile(String node, String state, long elapsedMs, Map dataNodeMetrics) implements ToXContentObject { + + /** Convenience constructor without data node metrics. */ + public TaskProfile(String node, String state, long elapsedMs) { + this(node, state, elapsedMs, null); + } @Override public XContentBuilder toXContent(XContentBuilder builder, Params params) throws IOException { @@ -30,6 +36,13 @@ public XContentBuilder toXContent(XContentBuilder builder, Params params) throws builder.field("node", node); builder.field("state", state); builder.field("elapsed_ms", elapsedMs); + if (dataNodeMetrics != null && dataNodeMetrics.isEmpty() == false) { + builder.startObject("data_node_metrics"); + for (Map.Entry entry : dataNodeMetrics.entrySet()) { + builder.field(entry.getKey(), entry.getValue()); + } + builder.endObject(); + } builder.endObject(); return builder; } diff --git a/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/exec/AnalyticsSearchService.java b/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/exec/AnalyticsSearchService.java index bb7fd7501a2ce..8f0d58bbe9cf5 100644 --- a/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/exec/AnalyticsSearchService.java +++ b/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/exec/AnalyticsSearchService.java @@ -42,6 +42,7 @@ import org.opensearch.tasks.TaskResourceTrackingService; import java.io.IOException; +import java.nio.charset.StandardCharsets; import java.util.Iterator; import java.util.List; import java.util.Map; @@ -176,18 +177,24 @@ public void executeFragmentStreamingAsync( responseHandler.onBatch(batch); } long fragmentTookNanos = System.nanoTime() - startNanos; - // Extract and log DataFusion execution metrics at DEBUG level - if (LOGGER.isDebugEnabled()) { + // Extract DataFusion execution metrics only when needed + if (request.profile() || LOGGER.isDebugEnabled()) { byte[] metricsJson = exec.resources().getExecutionMetrics(); - if (metricsJson != null) { + if (LOGGER.isDebugEnabled() && metricsJson != null) { LOGGER.debug( "[FragmentMetrics] shard={} metrics={}", shard.shardId(), - new String(metricsJson, java.nio.charset.StandardCharsets.UTF_8) + new String(metricsJson, StandardCharsets.UTF_8) ); } + if (request.profile() && metricsJson != null) { + responseHandler.onCompleteWithMetrics(metricsJson); + } else { + responseHandler.onComplete(); + } + } else { + responseHandler.onComplete(); } - responseHandler.onComplete(); ResolvedFragment resolved = exec.resolved(); DelegationDescriptor delegation = resolved.plan().getDelegationDescriptor(); boolean usedSecondaryIndex = delegation != null; @@ -350,6 +357,11 @@ public interface StreamingFragmentResponseHandler { void onComplete(); + /** Called with execution metrics when profiling is enabled. Default delegates to onComplete(). */ + default void onCompleteWithMetrics(byte[] metrics) { + onComplete(); + } + void onFailure(Exception e); } diff --git a/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/exec/AnalyticsSearchTransportService.java b/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/exec/AnalyticsSearchTransportService.java index d9f62f0967080..ea830192c8b42 100644 --- a/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/exec/AnalyticsSearchTransportService.java +++ b/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/exec/AnalyticsSearchTransportService.java @@ -8,8 +8,9 @@ package org.opensearch.analytics.exec; -import org.apache.logging.log4j.LogManager; -import org.apache.logging.log4j.Logger; +import org.apache.arrow.memory.BufferAllocator; +import org.apache.arrow.vector.VectorSchemaRoot; +import org.apache.arrow.vector.types.pojo.Schema; import org.opensearch.analytics.backend.EngineResultBatch; import org.opensearch.analytics.exec.action.FetchByRowIdsAction; import org.opensearch.analytics.exec.action.FetchByRowIdsRequest; @@ -52,7 +53,6 @@ */ @Singleton public class AnalyticsSearchTransportService { - private static final Logger logger = LogManager.getLogger(AnalyticsSearchTransportService.class); private final StreamTransportService transportService; private final ClusterService clusterService; @@ -152,9 +152,17 @@ private static void registerFetchByRowIdsHandler( */ private static AnalyticsSearchService.StreamingFragmentResponseHandler channelResponseHandler(TransportChannel channel) { return new AnalyticsSearchService.StreamingFragmentResponseHandler() { + private Schema batchSchema = null; + private BufferAllocator batchAllocator = null; + @Override public void onBatch(EngineResultBatch batch) throws Exception { - channel.sendResponseBatch(new FragmentExecutionArrowResponse(batch.getArrowRoot())); + VectorSchemaRoot root = batch.getArrowRoot(); + if (batchSchema == null && root.getFieldVectors().isEmpty() == false) { + batchSchema = root.getSchema(); + batchAllocator = root.getFieldVectors().getFirst().getAllocator(); + } + channel.sendResponseBatch(new FragmentExecutionArrowResponse(root)); } @Override @@ -162,6 +170,21 @@ public void onComplete() { channel.completeStream(); } + @Override + public void onCompleteWithMetrics(byte[] metrics) { + if (batchSchema != null && batchAllocator != null) { + VectorSchemaRoot sentinel = VectorSchemaRoot.create(batchSchema, batchAllocator); + sentinel.setRowCount(0); + try { + channel.sendResponseBatch(new FragmentExecutionArrowResponse(sentinel, metrics)); + } catch (Exception e) { + // Stream may already be cancelled — close sentinel to prevent leak + sentinel.close(); + } + } + channel.completeStream(); + } + @Override public void onFailure(Exception e) { try { @@ -240,15 +263,32 @@ public void handleStreamResponse(StreamTransportResponse operationListeners; private final BufferAllocator allocator; private final boolean ownsAllocator; + private final boolean profile; private volatile ExecutorService localTaskExecutor; private boolean closed; // guarded by `this` /** @@ -66,42 +67,8 @@ public class QueryContext { */ private final Map> resolvedTargetsByStage = new HashMap<>(); + /** Full-parameter constructor. Tests use {@link #forTest} factories. */ public QueryContext( - QueryDAG dag, - ThreadPool threadPool, - AnalyticsQueryTask parentTask, - BufferAllocator allocator, - boolean ownsAllocator, - int maxConcurrentShardRequestsPerNode, - int maxShardsPerQuery - ) { - this(dag, threadPool, parentTask, maxConcurrentShardRequestsPerNode, maxShardsPerQuery, List.of(), allocator, ownsAllocator); - } - - public QueryContext( - QueryDAG dag, - ThreadPool threadPool, - AnalyticsQueryTask parentTask, - BufferAllocator allocator, - boolean ownsAllocator, - int maxConcurrentShardRequestsPerNode, - int maxShardsPerQuery, - List operationListeners - ) { - this( - dag, - threadPool, - parentTask, - maxConcurrentShardRequestsPerNode, - maxShardsPerQuery, - operationListeners, - allocator, - ownsAllocator - ); - } - - /** Full-parameter constructor. Private; tests use {@link #forTest} factories. */ - private QueryContext( QueryDAG dag, ThreadPool threadPool, AnalyticsQueryTask parentTask, @@ -109,7 +76,8 @@ private QueryContext( int maxShardsPerQuery, List operationListeners, BufferAllocator allocator, - boolean ownsAllocator + boolean ownsAllocator, + boolean profile ) { this.dag = dag; this.threadPool = threadPool; @@ -119,12 +87,18 @@ private QueryContext( this.operationListeners = operationListeners; this.allocator = allocator; this.ownsAllocator = ownsAllocator; + this.profile = profile; } public QueryDAG dag() { return dag; } + /** Whether profiling is enabled for this query (data nodes should collect and return metrics). */ + public boolean profile() { + return profile; + } + public Executor searchExecutor() { return threadPool != null ? threadPool.executor(ThreadPool.Names.SEARCH) : Runnable::run; } @@ -251,7 +225,8 @@ public static QueryContext forTest(QueryDAG dag, AnalyticsQueryTask parentTask, DEFAULT_MAX_SHARDS_PER_QUERY, operationListeners, testAllocator, - true + true, + false ); } } diff --git a/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/exec/StreamingResponseListener.java b/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/exec/StreamingResponseListener.java index f48118775d89d..2cbdd4b9c57a5 100644 --- a/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/exec/StreamingResponseListener.java +++ b/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/exec/StreamingResponseListener.java @@ -39,6 +39,13 @@ public interface StreamingResponseListener { */ boolean onStreamResponse(Resp response, boolean isLast); + /** + * Called after the stream is exhausted with any trailing metadata sent by the data node. + * + * @param trailingMetadata application metadata bytes, or {@code null} if none sent + */ + default void onStreamComplete(byte[] trailingMetadata) {} + /** * Called when the request fails. Terminal failure event. * diff --git a/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/exec/action/FragmentExecutionArrowResponse.java b/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/exec/action/FragmentExecutionArrowResponse.java index b99ffd619c773..cc0ed097534c7 100644 --- a/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/exec/action/FragmentExecutionArrowResponse.java +++ b/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/exec/action/FragmentExecutionArrowResponse.java @@ -26,6 +26,10 @@ public FragmentExecutionArrowResponse(VectorSchemaRoot root) { super(root); } + public FragmentExecutionArrowResponse(VectorSchemaRoot root, byte[] metadata) { + super(root, metadata); + } + public FragmentExecutionArrowResponse(StreamInput in) throws IOException { super(in); } diff --git a/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/exec/action/FragmentExecutionRequest.java b/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/exec/action/FragmentExecutionRequest.java index 1e1513b3aa189..a0076f32037b4 100644 --- a/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/exec/action/FragmentExecutionRequest.java +++ b/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/exec/action/FragmentExecutionRequest.java @@ -40,12 +40,18 @@ public class FragmentExecutionRequest extends ActionRequest implements ShardInvo private final int stageId; private final ShardId shardId; private final List planAlternatives; + private final boolean profile; public FragmentExecutionRequest(String queryId, int stageId, ShardId shardId, List planAlternatives) { + this(queryId, stageId, shardId, planAlternatives, false); + } + + public FragmentExecutionRequest(String queryId, int stageId, ShardId shardId, List planAlternatives, boolean profile) { this.queryId = queryId; this.stageId = stageId; this.shardId = shardId; this.planAlternatives = planAlternatives; + this.profile = profile; } public FragmentExecutionRequest(StreamInput in) throws IOException { @@ -58,6 +64,7 @@ public FragmentExecutionRequest(StreamInput in) throws IOException { for (int i = 0; i < numAlternatives; i++) { planAlternatives.add(new PlanAlternative(in)); } + this.profile = in.readBoolean(); } @Override @@ -70,6 +77,7 @@ public void writeTo(StreamOutput out) throws IOException { for (PlanAlternative alt : planAlternatives) { alt.writeTo(out); } + out.writeBoolean(profile); } public String getQueryId() { @@ -88,6 +96,10 @@ public List getPlanAlternatives() { return planAlternatives; } + public boolean profile() { + return profile; + } + @Override public Task createTask(long id, String type, String action, TaskId parentTaskId, Map headers) { String desc = "queryId[" + queryId + "] stageId[" + stageId + "] shardId[" + shardId + "]"; diff --git a/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/exec/profile/QueryProfileBuilder.java b/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/exec/profile/QueryProfileBuilder.java index 11700adc84660..7bbed540564b8 100644 --- a/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/exec/profile/QueryProfileBuilder.java +++ b/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/exec/profile/QueryProfileBuilder.java @@ -20,9 +20,14 @@ import org.opensearch.analytics.spi.FilterDelegationInstructionNode; import org.opensearch.analytics.spi.InstructionNode; import org.opensearch.analytics.spi.ShardScanWithDelegationInstructionNode; +import org.opensearch.common.xcontent.XContentType; +import org.opensearch.core.xcontent.DeprecationHandler; +import org.opensearch.core.xcontent.NamedXContentRegistry; import java.util.ArrayList; +import java.util.LinkedHashMap; import java.util.List; +import java.util.Map; /** * Snapshots an {@link ExecutionGraph} into a {@link QueryProfile}. Pure read — no @@ -99,11 +104,31 @@ private static List buildTaskProfiles(StageExecution exec) { long start = t.startedAtMs(); long end = t.finishedAtMs(); long elapsed = (start > 0 && end > 0) ? end - start : 0L; - out.add(new TaskProfile(describeTarget(t), t.state().name(), elapsed)); + Map metrics = parseDataNodeMetrics(t.dataNodeMetrics()); + out.add(new TaskProfile(describeTarget(t), t.state().name(), elapsed, metrics)); } return out; } + @SuppressWarnings("unchecked") + private static Map parseDataNodeMetrics(byte[] json) { + if (json == null || json.length == 0) return null; + try { + var parser = XContentType.JSON.xContent() + .createParser(NamedXContentRegistry.EMPTY, DeprecationHandler.IGNORE_DEPRECATIONS, json); + Map raw = parser.map(); + Map result = new LinkedHashMap<>(); + for (Map.Entry entry : raw.entrySet()) { + if (entry.getValue() instanceof Number n) { + result.put(entry.getKey(), n.longValue()); + } + } + return result.isEmpty() ? null : result; + } catch (Exception e) { + return null; + } + } + private static String describeTarget(StageTask task) { if (task instanceof ShardStageTask shardTask) { var target = shardTask.target(); diff --git a/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/exec/stage/StageTask.java b/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/exec/stage/StageTask.java index 99a1c2a92a298..19a9f69ce2a7c 100644 --- a/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/exec/stage/StageTask.java +++ b/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/exec/stage/StageTask.java @@ -33,6 +33,7 @@ public abstract class StageTask { private final AtomicReference state = new AtomicReference<>(StageTaskState.CREATED); private volatile long startedAtMs; private volatile long finishedAtMs; + private volatile byte[] dataNodeMetrics; protected StageTask(StageTaskId id) { this.id = id; @@ -46,6 +47,16 @@ public StageTaskState state() { return state.get(); } + /** Raw JSON metrics bytes received from the data node, or null if not profiled. */ + public byte[] dataNodeMetrics() { + return dataNodeMetrics; + } + + /** Set by the coordinator when metrics arrive from the data node. */ + public void setDataNodeMetrics(byte[] metrics) { + this.dataNodeMetrics = metrics; + } + /** Wall-clock millis stamped on the first successful transition to {@link StageTaskState#RUNNING}, or 0 if never dispatched. */ public long startedAtMs() { return startedAtMs; diff --git a/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/exec/stage/shard/ShardFragmentStageExecution.java b/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/exec/stage/shard/ShardFragmentStageExecution.java index 35100a6a5e46a..4317113103ca7 100644 --- a/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/exec/stage/shard/ShardFragmentStageExecution.java +++ b/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/exec/stage/shard/ShardFragmentStageExecution.java @@ -123,7 +123,8 @@ public ExchangeSource outputSource() { * offload: reordering would let isLast race ahead and drop earlier batches via the * stage-terminal short-circuit. Inline also preserves end-to-end backpressure. */ - StreamingResponseListener responseListenerFor(int sourceOrdinal, ActionListener listener) { + StreamingResponseListener responseListenerFor(ShardStageTask task, ActionListener listener) { + final int sourceOrdinal = ((ShardExecutionTarget) task.target()).ordinal(); return new StreamingResponseListener<>() { @Override public boolean onStreamResponse(FragmentExecutionArrowResponse response, boolean isLast) { @@ -162,6 +163,13 @@ public boolean onStreamResponse(FragmentExecutionArrowResponse response, boolean return true; } + @Override + public void onStreamComplete(byte[] trailingMetadata) { + if (trailingMetadata != null) { + task.setDataNodeMetrics(trailingMetadata); + } + } + @Override public void onFailure(Exception e) { listener.onFailure(new RuntimeException("Stage " + getStageId() + " failed", e)); diff --git a/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/exec/stage/shard/ShardFragmentStageExecutionFactory.java b/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/exec/stage/shard/ShardFragmentStageExecutionFactory.java index 48f1ac149ea9a..1cf9dab8d4e89 100644 --- a/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/exec/stage/shard/ShardFragmentStageExecutionFactory.java +++ b/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/exec/stage/shard/ShardFragmentStageExecutionFactory.java @@ -60,7 +60,8 @@ public StageExecution createExecution(Stage stage, ExchangeSink sink, QueryConte queryId, stageId, target.shardId(), - planAlternatives + planAlternatives, + config.profile() ); // Execution pulls the resolver off `stage` and calls resolve() lazily at start(). // This keeps target resolution out of the build phase so cancellation before diff --git a/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/exec/stage/shard/ShardTaskRunner.java b/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/exec/stage/shard/ShardTaskRunner.java index 4b00e9ac832ff..8aedbf3ee8d4d 100644 --- a/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/exec/stage/shard/ShardTaskRunner.java +++ b/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/exec/stage/shard/ShardTaskRunner.java @@ -55,7 +55,7 @@ public void run(ShardStageTask task, ActionListener listener) { transport.dispatchFragmentStreaming( request, target.node(), - stage.responseListenerFor(target.ordinal(), listener), + stage.responseListenerFor(task, listener), config.parentTask(), pending ); diff --git a/sandbox/plugins/test-ppl-frontend/src/main/java/org/opensearch/ppl/action/RestPPLQueryAction.java b/sandbox/plugins/test-ppl-frontend/src/main/java/org/opensearch/ppl/action/RestPPLQueryAction.java index 88461f72d906b..d2107eb6d018c 100644 --- a/sandbox/plugins/test-ppl-frontend/src/main/java/org/opensearch/ppl/action/RestPPLQueryAction.java +++ b/sandbox/plugins/test-ppl-frontend/src/main/java/org/opensearch/ppl/action/RestPPLQueryAction.java @@ -44,30 +44,27 @@ public List routes() { @Override protected RestChannelConsumer prepareRequest(RestRequest request, NodeClient client) throws IOException { - String queryText; + String queryText = null; + boolean profile = false; try (XContentParser parser = request.contentParser()) { - queryText = parseQueryText(parser); - } - boolean explain = request.path().endsWith("/_explain"); - PPLRequest pplRequest = new PPLRequest(queryText, explain); - return channel -> client.execute(UnifiedPPLExecuteAction.INSTANCE, pplRequest, new RestToXContentListener<>(channel)); - } - - private String parseQueryText(XContentParser parser) throws IOException { - String query = null; - parser.nextToken(); // START_OBJECT - while (parser.nextToken() != XContentParser.Token.END_OBJECT) { - String fieldName = parser.currentName(); parser.nextToken(); - if ("query".equals(fieldName)) { - query = parser.text(); - } else { - parser.skipChildren(); + while (parser.nextToken() != XContentParser.Token.END_OBJECT) { + String fieldName = parser.currentName(); + parser.nextToken(); + if ("query".equals(fieldName)) { + queryText = parser.text(); + } else if ("profile".equals(fieldName)) { + profile = parser.booleanValue(); + } else { + parser.skipChildren(); + } } } - if (query == null || query.isEmpty()) { + if (queryText == null || queryText.isEmpty()) { throw new IllegalArgumentException("Request body must contain a 'query' field"); } - return query; + boolean enableProfile = profile || request.path().endsWith("/_explain"); + PPLRequest pplRequest = new PPLRequest(queryText, enableProfile); + return channel -> client.execute(UnifiedPPLExecuteAction.INSTANCE, pplRequest, new RestToXContentListener<>(channel)); } } diff --git a/sandbox/qa/analytics-engine-rest/src/test/java/org/opensearch/analytics/qa/ExplainApiIT.java b/sandbox/qa/analytics-engine-rest/src/test/java/org/opensearch/analytics/qa/ExplainApiIT.java index 372cf8054a5da..2e016d00dfa99 100644 --- a/sandbox/qa/analytics-engine-rest/src/test/java/org/opensearch/analytics/qa/ExplainApiIT.java +++ b/sandbox/qa/analytics-engine-rest/src/test/java/org/opensearch/analytics/qa/ExplainApiIT.java @@ -32,8 +32,10 @@ public class ExplainApiIT extends AnalyticsRestTestCase { private static final Dataset DATASET = new Dataset("calcs", "calcs"); private static final Dataset CLICKBENCH = ClickBenchTestHelper.DATASET; + private static final Dataset DELEGATION = new Dataset("delegation", "delegation"); private static boolean dataProvisioned = false; private static boolean clickBenchProvisioned = false; + private static boolean delegationProvisioned = false; @Override protected void onBeforeQuery() throws IOException { @@ -50,6 +52,13 @@ private void ensureClickBenchProvisioned() throws IOException { } } + private void ensureDelegationProvisioned() throws IOException { + if (delegationProvisioned == false) { + DatasetProvisioner.provision(client(), DELEGATION); + delegationProvisioned = true; + } + } + @SuppressWarnings("unchecked") public void testExplainReturnsProfileWithStages() throws IOException { Map result = executeExplain("source=" + DATASET.indexName + " | fields str0, num0"); @@ -177,10 +186,92 @@ public void testExplainTotalElapsedIsPositive() throws IOException { assertTrue("planning_time_ms is non-negative", planningTime >= 0); } + @SuppressWarnings("unchecked") + public void testProfileReturnsDataNodeMetrics() throws IOException { + ensureClickBenchProvisioned(); + Map result = executeWithProfile( + "source=" + CLICKBENCH.indexName + " | stats avg(AdvEngineID) by RegionID" + ); + + Map profile = (Map) result.get("profile"); + assertNotNull("profile present", profile); + List> stages = (List>) profile.get("stages"); + + Map shardStage = stages.stream() + .filter(s -> "SHARD_FRAGMENT".equals(s.get("execution_type"))) + .findFirst() + .orElseThrow(() -> new AssertionError("no SHARD_FRAGMENT stage")); + + List> tasks = (List>) shardStage.get("tasks"); + assertNotNull("tasks present", tasks); + + boolean anyMetrics = false; + for (Map task : tasks) { + Map metrics = (Map) task.get("data_node_metrics"); + if (metrics != null) { + anyMetrics = true; + assertNotNull("output_rows present", metrics.get("output_rows")); + assertNotNull("elapsed_compute present", metrics.get("elapsed_compute")); + assertNotNull("output_batches present", metrics.get("output_batches")); + assertTrue("output_rows non-negative", ((Number) metrics.get("output_rows")).longValue() >= 0); + } + } + assertTrue("at least one task has data_node_metrics", anyMetrics); + } + + @SuppressWarnings("unchecked") + public void testProfileDelegatedQueryHasFFMCollectorCalls() throws IOException { + ensureDelegationProvisioned(); + Map result = executeWithProfile( + "source=" + DELEGATION.indexName + " | where status = \"active\" | fields status, value" + ); + + Map profile = (Map) result.get("profile"); + assertNotNull("profile present", profile); + List> stages = (List>) profile.get("stages"); + + Map shardStage = stages.stream() + .filter(s -> "SHARD_FRAGMENT".equals(s.get("execution_type"))) + .findFirst() + .orElseThrow(() -> new AssertionError("no SHARD_FRAGMENT stage")); + + List> tasks = (List>) shardStage.get("tasks"); + assertNotNull("tasks present", tasks); + assertFalse("has tasks", tasks.isEmpty()); + + // Find a task with data_node_metrics (some shards may be empty on multi-node clusters) + Map metrics = null; + for (Map t : tasks) { + Map m = (Map) t.get("data_node_metrics"); + if (m != null && m.containsKey("ffm_collector_calls")) { + metrics = m; + break; + } + } + assertNotNull("at least one task has data_node_metrics with ffm_collector_calls", metrics); + + // Verify IndexedTableExec custom metrics proving Lucene delegation occurred + assertTrue( + "ffm_collector_calls > 0 (Lucene delegation occurred)", + ((Number) metrics.get("ffm_collector_calls")).longValue() > 0 + ); + assertNotNull("rows_matched present", metrics.get("rows_matched")); + assertEquals("rows_matched equals 10 (10% of 100 docs)", 10L, ((Number) metrics.get("rows_matched")).longValue()); + assertNotNull("row_groups_processed present", metrics.get("row_groups_processed")); + assertNotNull("index_query_time present", metrics.get("index_query_time")); + } + private Map executeExplain(String ppl) throws IOException { Request request = new Request("POST", "/_analytics/ppl/_explain"); request.setJsonEntity("{\"query\": \"" + escapeJson(ppl) + "\"}"); Response response = client().performRequest(request); return assertOkAndParse(response, "EXPLAIN: " + ppl); } + + private Map executeWithProfile(String ppl) throws IOException { + Request request = new Request("POST", "/_analytics/ppl"); + request.setJsonEntity("{\"query\": \"" + escapeJson(ppl) + "\", \"profile\": true}"); + Response response = client().performRequest(request); + return assertOkAndParse(response, "PROFILE: " + ppl); + } } diff --git a/sandbox/qa/analytics-engine-rest/src/test/resources/datasets/delegation/bulk.json b/sandbox/qa/analytics-engine-rest/src/test/resources/datasets/delegation/bulk.json new file mode 100644 index 0000000000000..c81324903917c --- /dev/null +++ b/sandbox/qa/analytics-engine-rest/src/test/resources/datasets/delegation/bulk.json @@ -0,0 +1,200 @@ +{"index": {}} +{"status": "active", "value": 0} +{"index": {}} +{"status": "inactive", "value": 1} +{"index": {}} +{"status": "inactive", "value": 2} +{"index": {}} +{"status": "inactive", "value": 3} +{"index": {}} +{"status": "inactive", "value": 4} +{"index": {}} +{"status": "inactive", "value": 5} +{"index": {}} +{"status": "inactive", "value": 6} +{"index": {}} +{"status": "inactive", "value": 7} +{"index": {}} +{"status": "inactive", "value": 8} +{"index": {}} +{"status": "inactive", "value": 9} +{"index": {}} +{"status": "active", "value": 10} +{"index": {}} +{"status": "inactive", "value": 11} +{"index": {}} +{"status": "inactive", "value": 12} +{"index": {}} +{"status": "inactive", "value": 13} +{"index": {}} +{"status": "inactive", "value": 14} +{"index": {}} +{"status": "inactive", "value": 15} +{"index": {}} +{"status": "inactive", "value": 16} +{"index": {}} +{"status": "inactive", "value": 17} +{"index": {}} +{"status": "inactive", "value": 18} +{"index": {}} +{"status": "inactive", "value": 19} +{"index": {}} +{"status": "active", "value": 20} +{"index": {}} +{"status": "inactive", "value": 21} +{"index": {}} +{"status": "inactive", "value": 22} +{"index": {}} +{"status": "inactive", "value": 23} +{"index": {}} +{"status": "inactive", "value": 24} +{"index": {}} +{"status": "inactive", "value": 25} +{"index": {}} +{"status": "inactive", "value": 26} +{"index": {}} +{"status": "inactive", "value": 27} +{"index": {}} +{"status": "inactive", "value": 28} +{"index": {}} +{"status": "inactive", "value": 29} +{"index": {}} +{"status": "active", "value": 30} +{"index": {}} +{"status": "inactive", "value": 31} +{"index": {}} +{"status": "inactive", "value": 32} +{"index": {}} +{"status": "inactive", "value": 33} +{"index": {}} +{"status": "inactive", "value": 34} +{"index": {}} +{"status": "inactive", "value": 35} +{"index": {}} +{"status": "inactive", "value": 36} +{"index": {}} +{"status": "inactive", "value": 37} +{"index": {}} +{"status": "inactive", "value": 38} +{"index": {}} +{"status": "inactive", "value": 39} +{"index": {}} +{"status": "active", "value": 40} +{"index": {}} +{"status": "inactive", "value": 41} +{"index": {}} +{"status": "inactive", "value": 42} +{"index": {}} +{"status": "inactive", "value": 43} +{"index": {}} +{"status": "inactive", "value": 44} +{"index": {}} +{"status": "inactive", "value": 45} +{"index": {}} +{"status": "inactive", "value": 46} +{"index": {}} +{"status": "inactive", "value": 47} +{"index": {}} +{"status": "inactive", "value": 48} +{"index": {}} +{"status": "inactive", "value": 49} +{"index": {}} +{"status": "active", "value": 50} +{"index": {}} +{"status": "inactive", "value": 51} +{"index": {}} +{"status": "inactive", "value": 52} +{"index": {}} +{"status": "inactive", "value": 53} +{"index": {}} +{"status": "inactive", "value": 54} +{"index": {}} +{"status": "inactive", "value": 55} +{"index": {}} +{"status": "inactive", "value": 56} +{"index": {}} +{"status": "inactive", "value": 57} +{"index": {}} +{"status": "inactive", "value": 58} +{"index": {}} +{"status": "inactive", "value": 59} +{"index": {}} +{"status": "active", "value": 60} +{"index": {}} +{"status": "inactive", "value": 61} +{"index": {}} +{"status": "inactive", "value": 62} +{"index": {}} +{"status": "inactive", "value": 63} +{"index": {}} +{"status": "inactive", "value": 64} +{"index": {}} +{"status": "inactive", "value": 65} +{"index": {}} +{"status": "inactive", "value": 66} +{"index": {}} +{"status": "inactive", "value": 67} +{"index": {}} +{"status": "inactive", "value": 68} +{"index": {}} +{"status": "inactive", "value": 69} +{"index": {}} +{"status": "active", "value": 70} +{"index": {}} +{"status": "inactive", "value": 71} +{"index": {}} +{"status": "inactive", "value": 72} +{"index": {}} +{"status": "inactive", "value": 73} +{"index": {}} +{"status": "inactive", "value": 74} +{"index": {}} +{"status": "inactive", "value": 75} +{"index": {}} +{"status": "inactive", "value": 76} +{"index": {}} +{"status": "inactive", "value": 77} +{"index": {}} +{"status": "inactive", "value": 78} +{"index": {}} +{"status": "inactive", "value": 79} +{"index": {}} +{"status": "active", "value": 80} +{"index": {}} +{"status": "inactive", "value": 81} +{"index": {}} +{"status": "inactive", "value": 82} +{"index": {}} +{"status": "inactive", "value": 83} +{"index": {}} +{"status": "inactive", "value": 84} +{"index": {}} +{"status": "inactive", "value": 85} +{"index": {}} +{"status": "inactive", "value": 86} +{"index": {}} +{"status": "inactive", "value": 87} +{"index": {}} +{"status": "inactive", "value": 88} +{"index": {}} +{"status": "inactive", "value": 89} +{"index": {}} +{"status": "active", "value": 90} +{"index": {}} +{"status": "inactive", "value": 91} +{"index": {}} +{"status": "inactive", "value": 92} +{"index": {}} +{"status": "inactive", "value": 93} +{"index": {}} +{"status": "inactive", "value": 94} +{"index": {}} +{"status": "inactive", "value": 95} +{"index": {}} +{"status": "inactive", "value": 96} +{"index": {}} +{"status": "inactive", "value": 97} +{"index": {}} +{"status": "inactive", "value": 98} +{"index": {}} +{"status": "inactive", "value": 99} diff --git a/sandbox/qa/analytics-engine-rest/src/test/resources/datasets/delegation/mapping.json b/sandbox/qa/analytics-engine-rest/src/test/resources/datasets/delegation/mapping.json new file mode 100644 index 0000000000000..868ef5baed4e0 --- /dev/null +++ b/sandbox/qa/analytics-engine-rest/src/test/resources/datasets/delegation/mapping.json @@ -0,0 +1,12 @@ +{ + "settings": { + "number_of_shards": 1, + "number_of_replicas": 0 + }, + "mappings": { + "properties": { + "status": { "type": "keyword" }, + "value": { "type": "integer" } + } + } +} From 10b8f6cb43bd5ed53aebb2300580c9947809b8d5 Mon Sep 17 00:00:00 2001 From: Somesh Gupta <35426854+aasom143@users.noreply.github.com> Date: Thu, 11 Jun 2026 23:01:04 +0530 Subject: [PATCH 06/14] feat: Add StatsPruneTree for segment, RG, and subtree-level pruning (#22051) Resolved comment, realted to sample StatsPruneTree, comment, log info to debug Passing StatsPruneTree in PredicateOnlyEvaluator for RG level pruning Signed-off-by: Somesh Gupta --- .../rust/src/indexed_executor.rs | 28 +- .../src/indexed_table/eval/bitmap_tree.rs | 231 +++++++++++- .../rust/src/indexed_table/eval/mod.rs | 19 + .../indexed_table/eval/predicate_evaluator.rs | 79 +++- .../indexed_table/eval/single_collector.rs | 86 ++++- .../rust/src/indexed_table/page_pruner.rs | 342 +++++++++++++++++- .../rust/src/indexed_table/stream.rs | 7 + .../rust/src/indexed_table/table_provider.rs | 32 +- .../tests_e2e/constant_predicate.rs | 5 +- .../tests_e2e/dynamic_filter_pushdown.rs | 4 +- .../tests_e2e/fuzz/delegation.rs | 5 +- .../indexed_table/tests_e2e/fuzz/harness.rs | 12 +- .../rust/src/indexed_table/tests_e2e/mod.rs | 5 +- .../indexed_table/tests_e2e/multi_segment.rs | 20 +- .../indexed_table/tests_e2e/null_columns.rs | 5 +- .../indexed_table/tests_e2e/page_pruning.rs | 8 +- .../tests_e2e/qtf_fetch_phase.rs | 5 +- .../tests_e2e/row_id_emission.rs | 20 +- .../indexed_table/tests_e2e/schema_drift.rs | 10 +- .../tests_e2e/streaming_at_scale.rs | 10 +- 20 files changed, 860 insertions(+), 73 deletions(-) diff --git a/sandbox/plugins/analytics-backend-datafusion/rust/src/indexed_executor.rs b/sandbox/plugins/analytics-backend-datafusion/rust/src/indexed_executor.rs index c22bd03a5b1d3..3169a4ef6190d 100644 --- a/sandbox/plugins/analytics-backend-datafusion/rust/src/indexed_executor.rs +++ b/sandbox/plugins/analytics-backend-datafusion/rust/src/indexed_executor.rs @@ -75,7 +75,7 @@ use crate::api::ShardView; use crate::datafusion_query_config::DatafusionQueryConfig; use crate::indexed_table::bool_tree::residual_bool_to_physical_expr; use crate::indexed_table::metrics::StreamMetrics; -use crate::indexed_table::page_pruner::{build_pruning_predicate, PagePruneMetrics}; +use crate::indexed_table::page_pruner::{build_pruning_predicate, PagePruneMetrics, StatsPruneTree}; /// Execute an indexed query. @@ -597,6 +597,22 @@ async unsafe fn execute_indexed_with_context_inner( FilterClass::Tree => None, }; + let prune_tree_config = extraction.as_ref().and_then(|e| { + let mut leaf_exprs: Vec> = Vec::new(); + collect_predicate_exprs(&e.tree, &mut leaf_exprs); + let leaf_predicates: HashMap> = leaf_exprs + .iter() + .filter_map(|expr| { + build_pruning_predicate(expr, schema.clone()) + .map(|pp| (Arc::as_ptr(expr) as *const () as usize, pp)) + }) + .collect(); + if leaf_predicates.is_empty() { + return None; + } + Some((e.tree.clone(), Arc::new(leaf_predicates), schema.clone())) + }); + let predicate_columns = collect_predicate_column_indices(extraction.as_ref()); let factory: EvaluatorFactory = match classification { @@ -615,7 +631,7 @@ async unsafe fn execute_indexed_with_context_inner( .and_then(|expr| build_pruning_predicate(expr, Arc::clone(&schema_for_pruner))); Arc::new( - move |segment: &SegmentFileInfo, _chunk, stream_metrics: &StreamMetrics| { + move |segment: &SegmentFileInfo, _chunk, stream_metrics: &StreamMetrics, stats_prune_tree: Option<&StatsPruneTree>| { let pruner = Arc::new(PagePruner::new( &schema_for_pruner, Arc::clone(&segment.metadata), @@ -626,6 +642,7 @@ async unsafe fn execute_indexed_with_context_inner( residual_pruning_predicate.clone(), residual_expr.clone(), Some(PagePruneMetrics::from_stream_metrics(stream_metrics)), + stats_prune_tree.cloned(), )); Ok(eval) }, @@ -691,7 +708,7 @@ async unsafe fn execute_indexed_with_context_inner( let bloom_schema = schema.clone(); let bloom_on_read = query_config.bloom_filter_on_read; Arc::new( - move |segment: &SegmentFileInfo, chunk, stream_metrics: &StreamMetrics| { + move |segment: &SegmentFileInfo, chunk, stream_metrics: &StreamMetrics, stats_prune_tree: Option<&StatsPruneTree>| { let collector_opt: Option> = match &correctness_provider { Some(provider) => { let collector = FfmSegmentCollector::create( @@ -747,6 +764,7 @@ async unsafe fn execute_indexed_with_context_inner( Arc::new(crate::indexed_table::eval::single_collector::FfmDelegatedBackendCollectorFactory), context_id, bloom_config, + stats_prune_tree.cloned(), )); Ok(eval) }, @@ -799,7 +817,7 @@ async unsafe fn execute_indexed_with_context_inner( ); Arc::new( - move |segment: &SegmentFileInfo, chunk, stream_metrics: &StreamMetrics| { + move |segment: &SegmentFileInfo, chunk, stream_metrics: &StreamMetrics, stats_prune_tree: Option<&StatsPruneTree>| { // Build one collector per Collector leaf for this chunk. let mut per_leaf: Vec<(i32, Arc)> = Vec::with_capacity(providers.len()); @@ -843,6 +861,7 @@ async unsafe fn execute_indexed_with_context_inner( stream_metrics, )), collector_strategy, + stats_prune_tree: stats_prune_tree.cloned(), }); Ok(eval) }, @@ -870,6 +889,7 @@ async unsafe fn execute_indexed_with_context_inner( query_config: Arc::clone(&query_config), predicate_columns, emit_row_ids, + prune_tree_config, })); ctx.register_table(&table_name, provider)?; diff --git a/sandbox/plugins/analytics-backend-datafusion/rust/src/indexed_table/eval/bitmap_tree.rs b/sandbox/plugins/analytics-backend-datafusion/rust/src/indexed_table/eval/bitmap_tree.rs index b278dfe706c3a..ce241b62dbcc8 100644 --- a/sandbox/plugins/analytics-backend-datafusion/rust/src/indexed_table/eval/bitmap_tree.rs +++ b/sandbox/plugins/analytics-backend-datafusion/rust/src/indexed_table/eval/bitmap_tree.rs @@ -71,7 +71,7 @@ use roaring::RoaringBitmap; use super::{LeafBitmapSource, RgEvalContext, TreeEvaluator, TreePrefetch}; use crate::indexed_table::bool_tree::ResolvedNode; -use crate::indexed_table::page_pruner::{PagePruneMetrics, PagePruner}; +use crate::indexed_table::page_pruner::{PagePruneMetrics, PagePruner, StatsPruneTree}; use crate::indexed_table::row_selection::{packed_bits_to_boolean_array, PositionMap}; use datafusion::physical_optimizer::pruning::PruningPredicate; @@ -88,6 +88,7 @@ impl TreeEvaluator for BitmapTreeEvaluator { page_pruner: &PagePruner, pruning_predicates: &HashMap>, page_prune_metrics: Option<&PagePruneMetrics>, + stats_prune_tree: Option<&StatsPruneTree>, ) -> Result { let mut per_leaf = Vec::new(); let mut dfs_counter = 0usize; @@ -104,6 +105,7 @@ impl TreeEvaluator for BitmapTreeEvaluator { &mut dfs_counter, &mut per_leaf, /* under_all_and_path */ true, + stats_prune_tree, )?; Ok(TreePrefetch { candidates, @@ -181,7 +183,27 @@ fn prefetch_node( dfs: &mut usize, out: &mut Vec<(usize, RoaringBitmap)>, under_all_and_path: bool, + stats_prune_tree: Option<&StatsPruneTree>, ) -> Result { + // RG-level subtree pruning: if this subtree provably can't match + // the current RG, skip the entire tree walk. Since collectors are + // always-true in the resolution, a false here means a Predicate + // under AND proved no match — collector bitmaps are irrelevant. + if let Some(spt) = stats_prune_tree { + if let Some(&false) = spt.rg_can_match.get(ctx.rg_idx) { + native_bridge_common::log_debug!( + "BitmapTree: skipping subtree for RG {} — pruned by RG-level stats", + ctx.rg_idx + ); + if under_all_and_path { + skip_dfs(node, dfs); + } else { + skip_dfs_with_empty_bitmaps(node, dfs, out); + } + return Ok(RoaringBitmap::new()); + } + } + match node { ResolvedNode::And(children) => { let mut indices: Vec = (0..children.len()).collect(); @@ -207,7 +229,8 @@ fn prefetch_node( page_prune_metrics, dfs, out, - under_all_and_path, // AND preserves the all-AND path + under_all_and_path, + stats_prune_tree.and_then(|spt| spt.children.get(i)), )?; result_bitmap = Some(match result_bitmap { None => child_bitmap, @@ -271,6 +294,7 @@ fn prefetch_node( out, // OR breaks all-AND propagation for its subtree. false, + stats_prune_tree.and_then(|spt| spt.children.get(val)), )?; result_bitmap |= &filtered_bitmap; @@ -304,6 +328,7 @@ fn prefetch_node( dfs, out, /* under_all_and_path */ false, + stats_prune_tree.and_then(|spt| spt.children.first()), )?; // Candidate-stage is a superset. Inverting a superset does // NOT yield a superset of the true NOT — it yields a subset @@ -440,6 +465,36 @@ fn skip_dfs(node: &ResolvedNode, dfs: &mut usize) { } } +/// Advance `dfs` and push empty bitmaps for each Collector leaf without +/// making FFM calls. Used when a subtree is pruned by RG-level stats but +/// refinement may still run (OR/NOT ancestor) — ensures the side-table +/// has entries so refinement doesn't panic on missing keys. +fn skip_dfs_with_empty_bitmaps( + node: &ResolvedNode, + dfs: &mut usize, + out: &mut Vec<(usize, RoaringBitmap)>, +) { + match node { + ResolvedNode::And(children) | ResolvedNode::Or(children) => { + for c in children { + skip_dfs_with_empty_bitmaps(c, dfs, out); + } + } + ResolvedNode::Not(child) => skip_dfs_with_empty_bitmaps(child, dfs, out), + ResolvedNode::Collector { collector, .. } => { + *dfs += 1; + let key = Arc::as_ptr(collector) as *const () as usize; + out.push((key, RoaringBitmap::new())); + } + ResolvedNode::Predicate(_) => *dfs += 1, + ResolvedNode::DelegationPossible { .. } => { + unimplemented!( + "invariant violation: DelegationPossible reached skip_dfs_with_empty_bitmaps." + ) + } + } +} + fn predicate_page_bitmap( expr: &Arc, ctx: &RgEvalContext, @@ -1155,7 +1210,7 @@ mod tests { }; let pruner = empty_pruner(); let result = BitmapTreeEvaluator - .prefetch(&tree, &test_ctx(), &leaves, &pruner, &HashMap::new(), None) + .prefetch(&tree, &test_ctx(), &leaves, &pruner, &HashMap::new(), None, None) .unwrap(); assert_eq!(result.candidates, bm(&[3, 4])); assert_eq!(result.per_leaf.len(), 2); @@ -1169,7 +1224,7 @@ mod tests { }; let pruner = empty_pruner(); let result = BitmapTreeEvaluator - .prefetch(&tree, &test_ctx(), &leaves, &pruner, &HashMap::new(), None) + .prefetch(&tree, &test_ctx(), &leaves, &pruner, &HashMap::new(), None, None) .unwrap(); assert_eq!(result.candidates, bm(&[1, 2, 3])); } @@ -1182,7 +1237,7 @@ mod tests { }; let pruner = empty_pruner(); let result = BitmapTreeEvaluator - .prefetch(&tree, &test_ctx(), &leaves, &pruner, &HashMap::new(), None) + .prefetch(&tree, &test_ctx(), &leaves, &pruner, &HashMap::new(), None, None) .unwrap(); // Universe is [0, 16). Minus {0,1,2} = {3..15} let expected: RoaringBitmap = (3u32..16).collect(); @@ -1197,7 +1252,7 @@ mod tests { }; let pruner = empty_pruner(); let state = BitmapTreeEvaluator - .prefetch(&tree, &test_ctx(), &leaves, &pruner, &HashMap::new(), None) + .prefetch(&tree, &test_ctx(), &leaves, &pruner, &HashMap::new(), None, None) .unwrap(); let schema = Arc::new(Schema::new(vec![Field::new("a", DataType::Int32, false)])); @@ -1457,7 +1512,7 @@ mod tests { let pruner = empty_pruner(); let result = BitmapTreeEvaluator - .prefetch(&tree, &test_ctx(), &leaves, &pruner, &HashMap::new(), None) + .prefetch(&tree, &test_ctx(), &leaves, &pruner, &HashMap::new(), None, None) .unwrap(); assert!(result.candidates.is_empty()); } @@ -1497,7 +1552,7 @@ mod tests { let pruner = empty_pruner(); let result = BitmapTreeEvaluator - .prefetch(&tree, &test_ctx(), &leaves, &pruner, &HashMap::new(), None) + .prefetch(&tree, &test_ctx(), &leaves, &pruner, &HashMap::new(), None, None) .unwrap(); // OR contributes {5} from standalone_leaf → non-empty candidates. assert!(!result.candidates.is_empty()); @@ -1535,7 +1590,7 @@ mod tests { let pruner = empty_pruner(); let result = BitmapTreeEvaluator - .prefetch(&tree, &test_ctx(), &leaves, &pruner, &HashMap::new(), None) + .prefetch(&tree, &test_ctx(), &leaves, &pruner, &HashMap::new(), None, None) .unwrap(); // NOT inverts empty AND → universe. assert_eq!(result.candidates.len(), 16); @@ -1566,7 +1621,9 @@ mod tests { subtree_cost(&test_predicate_node(), &ctx, &pruner, &pp), ctx.cost_predicate * COST_SCALE ); - assert_eq!(subtree_cost(&collector_leaf(0), &ctx, &pruner, &pp), ctx.cost_collector * COST_SCALE); + assert_eq!( + subtree_cost(&collector_leaf(0), &ctx, &pruner, &pp), ctx.cost_collector * COST_SCALE + ); } #[test] @@ -1575,7 +1632,9 @@ mod tests { let pruner = empty_pruner(); let pp = HashMap::new(); let wrapped = ResolvedNode::Not(Box::new(test_predicate_node())); - assert_eq!(subtree_cost(&wrapped, &ctx, &pruner, &pp), ctx.cost_predicate * COST_SCALE); + assert_eq!( + subtree_cost(&wrapped, &ctx, &pruner, &pp), ctx.cost_predicate * COST_SCALE + ); } #[test] @@ -1721,4 +1780,154 @@ mod tests { let bm = RoaringBitmap::new(); assert_eq!(ranges_from_bitmap(&bm, &ctx), vec![]); } + + // ── StatsPruneTree subtree pruning in prefetch ───────────────────── + + fn prune_tree_leaf(can_match: Vec) -> StatsPruneTree { + StatsPruneTree { + rg_can_match: can_match, + children: vec![], + } + } + + fn prune_tree_and( + children: Vec, + ) -> StatsPruneTree { + let mut rg_can_match = vec![true; children[0].rg_can_match.len()]; + for c in &children { + for (r, v) in rg_can_match.iter_mut().zip(c.rg_can_match.iter()) { + *r &= v; + } + } + StatsPruneTree { + rg_can_match, + children, + } + } + + fn prune_tree_or( + children: Vec, + ) -> StatsPruneTree { + let mut rg_can_match = vec![false; children[0].rg_can_match.len()]; + for c in &children { + for (r, v) in rg_can_match.iter_mut().zip(c.rg_can_match.iter()) { + *r |= v; + } + } + StatsPruneTree { + rg_can_match, + children, + } + } + + #[test] + fn stats_prune_tree_root_and_false_skips_entire_rg() { + let tree = ResolvedNode::And(vec![collector_leaf(0), collector_leaf(1)]); + let leaves = FixedLeafBitmaps { + bitmaps: vec![bm(&[1, 2, 3]), bm(&[3, 4, 5])], + }; + let pruner = empty_pruner(); + let spt = prune_tree_and(vec![ + prune_tree_leaf(vec![false]), + prune_tree_leaf(vec![true]), + ]); + let result = BitmapTreeEvaluator + .prefetch(&tree, &test_ctx(), &leaves, &pruner, &HashMap::new(), None, Some(&spt)) + .unwrap(); + assert!(result.candidates.is_empty()); + } + + #[test] + fn stats_prune_tree_or_skips_false_child() { + let tree = ResolvedNode::Or(vec![collector_leaf(0), collector_leaf(1)]); + let leaves = FixedLeafBitmaps { + bitmaps: vec![bm(&[10, 11, 12]), bm(&[3, 4])], + }; + let pruner = empty_pruner(); + let spt = prune_tree_or(vec![ + prune_tree_leaf(vec![false]), + prune_tree_leaf(vec![true]), + ]); + let result = BitmapTreeEvaluator + .prefetch(&tree, &test_ctx(), &leaves, &pruner, &HashMap::new(), None, Some(&spt)) + .unwrap(); + assert_eq!(result.candidates, bm(&[3, 4])); + } + + #[test] + fn stats_prune_tree_and_child_false_short_circuits() { + let tree = ResolvedNode::And(vec![collector_leaf(0), collector_leaf(1)]); + let leaves = FixedLeafBitmaps { + bitmaps: vec![bm(&[1, 2, 3]), bm(&[99])], + }; + let pruner = empty_pruner(); + let spt = prune_tree_and(vec![ + prune_tree_leaf(vec![true]), + prune_tree_leaf(vec![false]), + ]); + let result = BitmapTreeEvaluator + .prefetch(&tree, &test_ctx(), &leaves, &pruner, &HashMap::new(), None, Some(&spt)) + .unwrap(); + assert!(result.candidates.is_empty()); + } + + #[test] + fn stats_prune_tree_nested_or_under_and() { + // AND(OR(collector0, collector1), collector2) + // Cost sort: collector2 (cost=10k) first, OR subtree (cost=20k) second. + // DFS order after sort: collector2=0, OR-child0=1(skipped), OR-child1=2. + let tree = ResolvedNode::And(vec![ + ResolvedNode::Or(vec![collector_leaf(0), collector_leaf(1)]), + collector_leaf(2), + ]); + let leaves = FixedLeafBitmaps { + bitmaps: vec![bm(&[3, 4, 5]), bm(&[99]), bm(&[3, 4, 5, 6])], + }; + let pruner = empty_pruner(); + let or_spt = prune_tree_or(vec![ + prune_tree_leaf(vec![false]), + prune_tree_leaf(vec![true]), + ]); + let spt = prune_tree_and(vec![or_spt, prune_tree_leaf(vec![true])]); + let result = BitmapTreeEvaluator + .prefetch(&tree, &test_ctx(), &leaves, &pruner, &HashMap::new(), None, Some(&spt)) + .unwrap(); + // collector2 (dfs=0) → {3,4,5}; OR-child1 (dfs=2) → {3,4,5,6}; OR = {3,4,5,6} + // AND = {3,4,5} ∩ {3,4,5,6} = {3,4,5} + assert_eq!(result.candidates, bm(&[3, 4, 5])); + } + + #[test] + fn stats_prune_tree_or_all_children_false() { + let tree = ResolvedNode::And(vec![ + ResolvedNode::Or(vec![collector_leaf(0), collector_leaf(1)]), + collector_leaf(2), + ]); + let leaves = FixedLeafBitmaps { + bitmaps: vec![bm(&[1, 2]), bm(&[3, 4]), bm(&[5, 6])], + }; + let pruner = empty_pruner(); + let or_spt = prune_tree_or(vec![ + prune_tree_leaf(vec![false]), + prune_tree_leaf(vec![false]), + ]); + let spt = prune_tree_and(vec![or_spt, prune_tree_leaf(vec![true])]); + let result = BitmapTreeEvaluator + .prefetch(&tree, &test_ctx(), &leaves, &pruner, &HashMap::new(), None, Some(&spt)) + .unwrap(); + assert!(result.candidates.is_empty()); + } + + #[test] + fn stats_prune_tree_none_evaluates_normally() { + let tree = ResolvedNode::And(vec![collector_leaf(0), collector_leaf(1)]); + let leaves = FixedLeafBitmaps { + bitmaps: vec![bm(&[1, 2, 3, 4]), bm(&[3, 4, 5])], + }; + let pruner = empty_pruner(); + let result = BitmapTreeEvaluator + .prefetch(&tree, &test_ctx(), &leaves, &pruner, &HashMap::new(), None, None) + .unwrap(); + assert_eq!(result.candidates, bm(&[3, 4])); + } } diff --git a/sandbox/plugins/analytics-backend-datafusion/rust/src/indexed_table/eval/mod.rs b/sandbox/plugins/analytics-backend-datafusion/rust/src/indexed_table/eval/mod.rs index 49ea62df2cdb2..e00764887b9e1 100644 --- a/sandbox/plugins/analytics-backend-datafusion/rust/src/indexed_table/eval/mod.rs +++ b/sandbox/plugins/analytics-backend-datafusion/rust/src/indexed_table/eval/mod.rs @@ -83,6 +83,7 @@ pub(super) fn remap_expr_to_batch( use super::bool_tree::ResolvedNode; use super::page_pruner::PagePruneMetrics; use super::page_pruner::PagePruner; +use super::page_pruner::StatsPruneTree; use super::row_selection::PositionMap; use super::stream::RowGroupInfo; use datafusion::arrow::buffer::Buffer; @@ -272,6 +273,7 @@ pub trait TreeEvaluator: Send + Sync { page_pruner: &PagePruner, pruning_predicates: &HashMap>, page_prune_metrics: Option<&PagePruneMetrics>, + stats_prune_tree: Option<&StatsPruneTree>, ) -> Result; /// Refinement stage: produce the exact per-row `BooleanArray` for one @@ -348,6 +350,8 @@ pub struct TreeBitsetSource { /// recommended here — multiple FFM calls per collector per RG can /// be expensive in multi-collector trees. pub collector_strategy: CollectorCallStrategy, + /// Precomputed per-subtree RG match vectors. Built once at construction. + pub stats_prune_tree: Option, } impl RowGroupBitsetSource for TreeBitsetSource { @@ -358,6 +362,18 @@ impl RowGroupBitsetSource for TreeBitsetSource { max_doc: i32, ) -> Result, String> { let t = Instant::now(); + + // RG-level early-exit: precomputed from column stats at construction. + if let Some(ref ann) = self.stats_prune_tree { + if let Some(&false) = ann.rg_can_match.get(rg.index) { + native_bridge_common::log_debug!( + "BitmapTree: skipping RG {} — pruned by RG-level stats", + rg.index + ); + return Ok(None); + } + } + let ctx = RgEvalContext { rg_idx: rg.index, rg_first_row: rg.first_row, @@ -405,6 +421,7 @@ impl RowGroupBitsetSource for TreeBitsetSource { // inflate counts. We compute final page-level metrics below // after the bitmap tree is fully resolved. None, + self.stats_prune_tree.as_ref(), ) .map_err(|e| format!("TreeBitsetSource::prefetch_rg(rg={}): {}", rg.index, e))?; if prefetch.candidates.is_empty() { @@ -755,6 +772,7 @@ mod tests { _page_pruner: &PagePruner, _pruning_predicates: &HashMap>, _page_prune_metrics: Option<&PagePruneMetrics>, + _stats_prune_tree: Option<&StatsPruneTree>, ) -> Result { Ok(TreePrefetch { candidates: roaring::RoaringBitmap::new(), @@ -803,6 +821,7 @@ mod tests { pruning_predicates: std::sync::Arc::new(HashMap::new()), page_prune_metrics: None, collector_strategy: CollectorCallStrategy::TightenOuterBounds, + stats_prune_tree: None, }; assert!(!source.needs_row_mask()); } diff --git a/sandbox/plugins/analytics-backend-datafusion/rust/src/indexed_table/eval/predicate_evaluator.rs b/sandbox/plugins/analytics-backend-datafusion/rust/src/indexed_table/eval/predicate_evaluator.rs index 8152643673efb..ef76ece7c69b8 100644 --- a/sandbox/plugins/analytics-backend-datafusion/rust/src/indexed_table/eval/predicate_evaluator.rs +++ b/sandbox/plugins/analytics-backend-datafusion/rust/src/indexed_table/eval/predicate_evaluator.rs @@ -23,7 +23,7 @@ use roaring::RoaringBitmap; use super::eval_helpers::{compute_page_ranges, evaluate_residual, universe_bitmap_from_page_ranges}; use super::{PrefetchedRg, RowGroupBitsetSource}; -use crate::indexed_table::page_pruner::{PagePruneMetrics, PagePruner}; +use crate::indexed_table::page_pruner::{PagePruneMetrics, PagePruner, StatsPruneTree}; use crate::indexed_table::row_selection::{bitmap_to_packed_bits, PositionMap}; use crate::indexed_table::stream::RowGroupInfo; @@ -35,6 +35,7 @@ pub struct PredicateOnlyEvaluator { pruning_predicate: Option>, residual_expr: Option>, page_prune_metrics: Option, + stats_prune_tree: Option, } impl PredicateOnlyEvaluator { @@ -43,12 +44,14 @@ impl PredicateOnlyEvaluator { pruning_predicate: Option>, residual_expr: Option>, page_prune_metrics: Option, + stats_prune_tree: Option, ) -> Self { Self { page_pruner, pruning_predicate, residual_expr, page_prune_metrics, + stats_prune_tree, } } } @@ -62,6 +65,17 @@ impl RowGroupBitsetSource for PredicateOnlyEvaluator { ) -> Result, String> { let t = Instant::now(); + // RG-level early-exit: precomputed from column stats at construction. + if let Some(ref spt) = self.stats_prune_tree { + if let Some(&false) = spt.rg_can_match.get(rg.index) { + native_bridge_common::log_debug!( + "PredicateOnly: skipping RG {} — pruned by RG-level stats", + rg.index + ); + return Ok(None); + } + } + let page_ranges = compute_page_ranges( self.pruning_predicate.as_ref(), &self.page_pruner, @@ -102,3 +116,66 @@ impl RowGroupBitsetSource for PredicateOnlyEvaluator { Ok(Some(evaluate_residual(residual, batch, batch_len)?)) } } + +#[cfg(test)] +mod tests { + use super::*; + use crate::indexed_table::page_pruner::PagePruner; + use datafusion::arrow::array::Int32Array; + use datafusion::arrow::datatypes::{DataType, Field, Schema}; + use datafusion::parquet::arrow::arrow_reader::{ArrowReaderMetadata, ArrowReaderOptions}; + use datafusion::parquet::arrow::ArrowWriter; + use std::sync::Arc; + use tempfile::NamedTempFile; + + fn minimal_page_pruner() -> Arc { + let schema = Arc::new(Schema::new(vec![Field::new("a", DataType::Int32, false)])); + let batch = datafusion::arrow::record_batch::RecordBatch::try_new( + schema.clone(), + vec![Arc::new(Int32Array::from(vec![0i32; 8]))], + ) + .unwrap(); + let tmp = NamedTempFile::new().unwrap(); + let mut writer = ArrowWriter::try_new(tmp.reopen().unwrap(), schema.clone(), None).unwrap(); + writer.write(&batch).unwrap(); + writer.close().unwrap(); + let file = tmp.reopen().unwrap(); + let options = ArrowReaderOptions::new().with_page_index(true); + let meta = ArrowReaderMetadata::load(&file, options).unwrap(); + Arc::new(PagePruner::new(meta.schema(), meta.metadata().clone())) + } + + #[test] + fn stats_prune_tree_skips_rg_when_false() { + let pruner = minimal_page_pruner(); + let spt = StatsPruneTree { + rg_can_match: vec![false], + children: vec![], + }; + let eval = PredicateOnlyEvaluator::new(pruner, None, None, None, Some(spt)); + let rg = RowGroupInfo { index: 0, first_row: 0, num_rows: 8 }; + assert!(eval.prefetch_rg(&rg, 0, 8).unwrap().is_none()); + } + + #[test] + fn stats_prune_tree_allows_rg_when_true() { + let pruner = minimal_page_pruner(); + let spt = StatsPruneTree { + rg_can_match: vec![true], + children: vec![], + }; + let eval = PredicateOnlyEvaluator::new(pruner, None, None, None, Some(spt)); + let rg = RowGroupInfo { index: 0, first_row: 0, num_rows: 8 }; + let prefetched = eval.prefetch_rg(&rg, 0, 8).unwrap().expect("should have candidates"); + assert_eq!(prefetched.candidates.len(), 8); + } + + #[test] + fn stats_prune_tree_none_does_not_prune() { + let pruner = minimal_page_pruner(); + let eval = PredicateOnlyEvaluator::new(pruner, None, None, None, None); + let rg = RowGroupInfo { index: 0, first_row: 0, num_rows: 8 }; + let prefetched = eval.prefetch_rg(&rg, 0, 8).unwrap().expect("should have candidates"); + assert_eq!(prefetched.candidates.len(), 8); + } +} diff --git a/sandbox/plugins/analytics-backend-datafusion/rust/src/indexed_table/eval/single_collector.rs b/sandbox/plugins/analytics-backend-datafusion/rust/src/indexed_table/eval/single_collector.rs index ab897f4360ff4..5a573dd4db9f6 100644 --- a/sandbox/plugins/analytics-backend-datafusion/rust/src/indexed_table/eval/single_collector.rs +++ b/sandbox/plugins/analytics-backend-datafusion/rust/src/indexed_table/eval/single_collector.rs @@ -33,7 +33,7 @@ use roaring::RoaringBitmap; use super::{PrefetchedRg, RowGroupBitsetSource}; use crate::indexed_table::ffm_callbacks::{create_provider, FfmSegmentCollector, ProviderHandle}; use crate::indexed_table::index::RowGroupDocsCollector; -use crate::indexed_table::page_pruner::{PagePruneMetrics, PagePruner}; +use crate::indexed_table::page_pruner::{PagePruneMetrics, PagePruner, StatsPruneTree}; use crate::indexed_table::row_selection::{ bitmap_to_packed_bits, packed_bits_to_boolean_array, row_selection_to_bitmap, PositionMap, }; @@ -94,7 +94,6 @@ impl DelegatedBackendCollectorFactory for FfmDelegatedBackendCollectorFactory { } } - /// Per-RG state the evaluator keeps for refinement. In row-granular /// mode parquet narrowed fully via `with_predicate` + `RowSelection` /// and nothing is needed here. In block-granular mode we need the @@ -177,6 +176,8 @@ pub struct SingleCollectorEvaluator { context_id: i64, /// Bloom filter pruning config. None = disabled. bloom_config: Option, + /// Precomputed per-RG/subtree match status from RG-level column stats. + stats_prune_tree: Option, } /// Resources needed for per-RG bloom filter pruning. @@ -204,6 +205,7 @@ impl SingleCollectorEvaluator { delegated_backend_collector_factory: Arc, context_id: i64, bloom_config: Option, + stats_prune_tree: Option, ) -> Self { Self { collector, @@ -218,6 +220,7 @@ impl SingleCollectorEvaluator { delegated_backend_collector_factory, context_id, bloom_config, + stats_prune_tree, } } } @@ -258,6 +261,17 @@ impl RowGroupBitsetSource for SingleCollectorEvaluator { ) -> Result, String> { let t = Instant::now(); + // RG-level early-exit: precomputed from column stats at construction. + if let Some(ref spt) = self.stats_prune_tree { + if let Some(&false) = spt.rg_can_match.get(rg.index) { + native_bridge_common::log_debug!( + "SingleCollector: skipping RG {} — pruned by RG-level stats", + rg.index + ); + return Ok(None); + } + } + // Page-prune to discover which row ranges survive. let page_ranges: Option> = self.pruning_predicate.as_ref().and_then(|pp| { self.page_pruner @@ -690,7 +704,7 @@ mod tests { docs: vec![0, 3, 7], }) as Arc; let pruner = minimal_page_pruner(); - let eval = SingleCollectorEvaluator::new(Some(collector), pruner, None, None, None, None, CollectorCallStrategy::FullRange, Arc::new(HashMap::new()), 0, Arc::new(FfmDelegatedBackendCollectorFactory), 0, None); + let eval = SingleCollectorEvaluator::new(Some(collector), pruner, None, None, None, None, CollectorCallStrategy::FullRange, Arc::new(HashMap::new()), 0, Arc::new(FfmDelegatedBackendCollectorFactory), 0, None, None); let rg = RowGroupInfo { index: 0, @@ -706,7 +720,7 @@ mod tests { fn on_batch_mask_returns_none_for_path_b() { let collector = Arc::new(StubCollector { docs: vec![0] }) as Arc; let pruner = minimal_page_pruner(); - let eval = SingleCollectorEvaluator::new(Some(collector), pruner, None, None, None, None, CollectorCallStrategy::FullRange, Arc::new(HashMap::new()), 0, Arc::new(FfmDelegatedBackendCollectorFactory), 0, None); + let eval = SingleCollectorEvaluator::new(Some(collector), pruner, None, None, None, None, CollectorCallStrategy::FullRange, Arc::new(HashMap::new()), 0, Arc::new(FfmDelegatedBackendCollectorFactory), 0, None, None); let schema = Arc::new(Schema::new(vec![Field::new("a", DataType::Int32, false)])); let batch = datafusion::arrow::record_batch::RecordBatch::try_new( schema, @@ -734,7 +748,7 @@ mod tests { // (it's the only post-decode filter we have on this path). let collector = Arc::new(StubCollector { docs: vec![0] }) as Arc; let pruner = minimal_page_pruner(); - let eval = SingleCollectorEvaluator::new(Some(collector), pruner, None, None, None, None, CollectorCallStrategy::FullRange, Arc::new(HashMap::new()), 0, Arc::new(FfmDelegatedBackendCollectorFactory), 0, None); + let eval = SingleCollectorEvaluator::new(Some(collector), pruner, None, None, None, None, CollectorCallStrategy::FullRange, Arc::new(HashMap::new()), 0, Arc::new(FfmDelegatedBackendCollectorFactory), 0, None, None); assert!(eval.needs_row_mask()); } @@ -742,7 +756,7 @@ mod tests { fn empty_match_returns_none() { let collector = Arc::new(StubCollector { docs: vec![] }) as Arc; let pruner = minimal_page_pruner(); - let eval = SingleCollectorEvaluator::new(Some(collector), pruner, None, None, None, None, CollectorCallStrategy::FullRange, Arc::new(HashMap::new()), 0, Arc::new(FfmDelegatedBackendCollectorFactory), 0, None); + let eval = SingleCollectorEvaluator::new(Some(collector), pruner, None, None, None, None, CollectorCallStrategy::FullRange, Arc::new(HashMap::new()), 0, Arc::new(FfmDelegatedBackendCollectorFactory), 0, None, None); let rg = RowGroupInfo { index: 0, first_row: 0, @@ -762,7 +776,7 @@ mod tests { docs: vec![0, 3, 7], }) as Arc; let pruner = minimal_page_pruner(); - let eval = SingleCollectorEvaluator::new(Some(collector), pruner, None, None, None, None, CollectorCallStrategy::FullRange, Arc::new(HashMap::new()), 0, Arc::new(FfmDelegatedBackendCollectorFactory), 0, None); + let eval = SingleCollectorEvaluator::new(Some(collector), pruner, None, None, None, None, CollectorCallStrategy::FullRange, Arc::new(HashMap::new()), 0, Arc::new(FfmDelegatedBackendCollectorFactory), 0, None, None); let rg = RowGroupInfo { index: 0, @@ -774,8 +788,64 @@ mod tests { assert_eq!(got, vec![0u32, 3, 7]); } + #[test] + fn stats_prune_tree_skips_rg_when_false() { + let collector = Arc::new(StubCollector { + docs: vec![0, 3, 7], + }) as Arc; + let pruner = minimal_page_pruner(); + let spt = StatsPruneTree { + rg_can_match: vec![false], + children: vec![], + }; + let eval = SingleCollectorEvaluator::new(Some(collector), pruner, None, None, None, None, CollectorCallStrategy::FullRange, Arc::new(HashMap::new()), 0, Arc::new(FfmDelegatedBackendCollectorFactory), 0, None, Some(spt)); + let rg = RowGroupInfo { + index: 0, + first_row: 0, + num_rows: 8, + }; + assert!(eval.prefetch_rg(&rg, 0, 8).unwrap().is_none()); + } + + #[test] + fn stats_prune_tree_allows_rg_when_true() { + let collector = Arc::new(StubCollector { + docs: vec![0, 3, 7], + }) as Arc; + let pruner = minimal_page_pruner(); + let spt = StatsPruneTree { + rg_can_match: vec![true], + children: vec![], + }; + let eval = SingleCollectorEvaluator::new(Some(collector), pruner, None, None, None, None, CollectorCallStrategy::FullRange, Arc::new(HashMap::new()), 0, Arc::new(FfmDelegatedBackendCollectorFactory), 0, None, Some(spt)); + let rg = RowGroupInfo { + index: 0, + first_row: 0, + num_rows: 8, + }; + let prefetched = eval.prefetch_rg(&rg, 0, 8).unwrap().expect("should have matches"); + let got: Vec = prefetched.candidates.iter().collect(); + assert_eq!(got, vec![0u32, 3, 7]); + } + + #[test] + fn stats_prune_tree_none_does_not_prune() { + let collector = Arc::new(StubCollector { + docs: vec![1, 5], + }) as Arc; + let pruner = minimal_page_pruner(); + let eval = SingleCollectorEvaluator::new(Some(collector), pruner, None, None, None, None, CollectorCallStrategy::FullRange, Arc::new(HashMap::new()), 0, Arc::new(FfmDelegatedBackendCollectorFactory), 0, None, None); + let rg = RowGroupInfo { + index: 0, + first_row: 0, + num_rows: 8, + }; + let prefetched = eval.prefetch_rg(&rg, 0, 8).unwrap().expect("should have matches"); + let got: Vec = prefetched.candidates.iter().collect(); + assert_eq!(got, vec![1u32, 5]); + } + // Keep the `fmt` import used #[allow(dead_code)] fn _use(_: &dyn fmt::Debug) {} } - diff --git a/sandbox/plugins/analytics-backend-datafusion/rust/src/indexed_table/page_pruner.rs b/sandbox/plugins/analytics-backend-datafusion/rust/src/indexed_table/page_pruner.rs index 177cbb5d1c494..19f8ea7037baa 100644 --- a/sandbox/plugins/analytics-backend-datafusion/rust/src/indexed_table/page_pruner.rs +++ b/sandbox/plugins/analytics-backend-datafusion/rust/src/indexed_table/page_pruner.rs @@ -50,6 +50,7 @@ use datafusion::parquet::file::metadata::ParquetMetaData; #[cfg(test)] use datafusion::physical_expr::expressions::{BinaryExpr, Column as PhysColumn, Literal}; use datafusion::physical_expr::PhysicalExpr; +use datafusion::physical_expr::utils::collect_columns; use datafusion::physical_optimizer::pruning::{PruningPredicate, PruningStatistics}; /// Per-row-group page pruner. Owns schema + metadata references; the @@ -84,7 +85,7 @@ impl PagePruner { metrics: Option<&PagePruneMetrics>, ) -> Option { let columns = - datafusion::physical_expr::utils::collect_columns(pruning_predicate.orig_expr()); + collect_columns(pruning_predicate.orig_expr()); if columns.is_empty() { return None; } @@ -184,7 +185,7 @@ impl PagePruner { let keep = match pruning_predicate.prune(&stats) { Ok(k) => k, Err(e) => { - log::debug!("page pruning error for rg {}: {}", rg_idx, e); + native_bridge_common::log_debug!("page pruning error for rg {}: {}", rg_idx, e); if let Some(m) = metrics { if let Some(ref c) = m.page_pruning_unavailable { c.add(1); @@ -272,12 +273,12 @@ pub fn build_pruning_predicate( let pruning_predicate = match PruningPredicate::try_new(Arc::clone(expr), schema) { Ok(pp) => pp, Err(e) => { - log::debug!("PruningPredicate::try_new failed for {:?}: {}", expr, e); + native_bridge_common::log_debug!("PruningPredicate::try_new failed for {:?}: {}", expr, e); return None; } }; if pruning_predicate.always_true() { - log::trace!("PruningPredicate collapsed to always_true for {:?}", expr); + native_bridge_common::log_debug!("PruningPredicate collapsed to always_true for {:?}", expr); return None; } Some(Arc::new(pruning_predicate)) @@ -503,6 +504,217 @@ fn to_row_selection(keep: Vec, row_counts: &[usize]) -> RowSelection { RowSelection::from(out) } +/// Evaluate a single predicate against RG-level column stats for the given RGs. +/// Returns a `Vec` aligned to `rg_indices`. +/// +/// Conservative on any error: returns `true` (can-match) for every RG rather +/// than falsely pruning. This means a stats extraction failure, missing column, +/// or `PruningPredicate::prune` error will never cause data loss — the RG will +/// simply not be skipped. +fn eval_leaf( + pruning_predicate: &PruningPredicate, + metadata: &ParquetMetaData, + schema: &SchemaRef, + rg_indices: &[usize], +) -> Vec { + let num = rg_indices.len(); + if num == 0 { + return vec![]; + } + let columns = collect_columns(pruning_predicate.orig_expr()); + if columns.is_empty() { + return vec![true; num]; + } + let arrow_schema = schema.as_ref(); + let rg_metas: Vec<_> = rg_indices + .iter() + .filter_map(|&idx| metadata.row_groups().get(idx)) + .collect(); + if rg_metas.len() != num { + return vec![true; num]; + } + let mut col_stats: HashMap)> = HashMap::new(); + for col in &columns { + if arrow_schema.index_of(col.name()).is_err() { + continue; + } + let converter = match StatisticsConverter::try_new( + col.name(), arrow_schema, metadata.file_metadata().schema_descr(), + ) { + Ok(c) => c, + Err(_) => continue, + }; + let min_arr = match converter.row_group_mins(rg_metas.iter().copied()) { + Ok(arr) => arr, + Err(_) => continue, + }; + let max_arr = match converter.row_group_maxes(rg_metas.iter().copied()) { + Ok(arr) => arr, + Err(_) => continue, + }; + let null_counts = converter + .row_group_null_counts(rg_metas.iter().copied()) + .ok() + .map(|arr| Arc::new(arr) as ArrayRef); + col_stats.insert(col.name().to_string(), (min_arr, max_arr, null_counts)); + } + if col_stats.is_empty() { + return vec![true; num]; + } + let stats = RgLevelStats { + col_stats, + num_rgs: num, + row_counts: rg_metas.iter().map(|m| m.num_rows()).collect(), + }; + match pruning_predicate.prune(&stats) { + Ok(keep) => keep, + Err(_) => vec![true; num], + } +} + +struct RgLevelStats { + col_stats: HashMap)>, + num_rgs: usize, + row_counts: Vec, +} + +impl PruningStatistics for RgLevelStats { + fn min_values(&self, column: &Column) -> Option { + self.col_stats.get(column.name()).map(|(m, _, _)| Arc::clone(m)) + } + fn max_values(&self, column: &Column) -> Option { + self.col_stats.get(column.name()).map(|(_, m, _)| Arc::clone(m)) + } + fn num_containers(&self) -> usize { + self.num_rgs + } + fn null_counts(&self, column: &Column) -> Option { + self.col_stats.get(column.name()).and_then(|(_, _, n)| n.clone()) + } + fn row_counts(&self) -> Option { + let arr = Int64Array::from_iter_values(self.row_counts.iter().copied()); + Some(Arc::new(arr) as ArrayRef) + } + fn contained( + &self, + _column: &Column, + _values: &std::collections::HashSet, + ) -> Option { + None + } +} + +/// Precomputed tree of per-RG match vectors built from column statistics. +/// Used for segment-level, RG-level, and subtree-level pruning. +/// +/// # Example: 5 RGs, 9 leaves +/// +/// Legend: `df` = Predicate leaf (parquet column stats), `luc` = Collector (always true). +/// Each node shows its `rg_can_match` bitset: `[RG0, RG1, RG2, RG3, RG4]`. +/// +/// ```text +/// AND +/// [00001] +/// / | \ +/// OR OR AND +/// [01101] [11111] [00011] +/// / \ / | \ / \ +/// AND df AND luc df df NOT +/// [01100] [00001] [00100] [11111] [11110] [00011] [11111] +/// / \ / \ | +/// df df df df df +/// [11100] [01111] [11100] [00111] [11000] +/// ``` +/// +/// Bottom-up: +/// - AND(df[11100], df[01111]) = [01100] +/// - OR₁(AND[01100], df[00001]) = [01101] +/// - AND(df[11100], df[00111]) = [00100] +/// - OR₂(AND[00100], luc[11111], df[11110]) = [11111] ← luc dominates +/// - NOT(df[11000]) = [11111] ← conservative: can't prune through negation +/// - AND₃(df[00011], NOT[11111]) = [00011] +/// - Root AND(OR₁[01101], OR₂[11111], AND₃[00011]) = [00001] +/// +/// Result: only RG4 survives. NOT is always [11111] because inverting +/// stats-based pruning is unsound (superset inverted = subset). +/// Subtree vectors enable short-circuiting: for RG0, AND₃=[0....] so the +/// bitmap-tree walk skips that entire subtree (no FFM collector calls). +#[derive(Clone)] +pub struct StatsPruneTree { + /// Per-RG: `false` means this subtree provably can't match that RG. + pub rg_can_match: Vec, + pub children: Vec, +} + +impl StatsPruneTree { + pub fn build_from_bool_node( + node: &super::bool_tree::BoolNode, + leaf_predicates: &HashMap>, + metadata: &ParquetMetaData, + schema: &SchemaRef, + rg_indices: &[usize], + ) -> Self { + use super::bool_tree::BoolNode; + let num = rg_indices.len(); + match node { + BoolNode::And(children) => { + let annotated_children: Vec<_> = children + .iter() + .map(|c| Self::build_from_bool_node(c, leaf_predicates, metadata, schema, rg_indices)) + .collect(); + let mut rg_can_match = vec![true; num]; + for child in &annotated_children { + for (r, c) in rg_can_match.iter_mut().zip(child.rg_can_match.iter()) { + *r &= c; + } + } + Self { rg_can_match, children: annotated_children } + } + BoolNode::Or(children) => { + let annotated_children: Vec<_> = children + .iter() + .map(|c| Self::build_from_bool_node(c, leaf_predicates, metadata, schema, rg_indices)) + .collect(); + let mut rg_can_match = vec![false; num]; + for child in &annotated_children { + for (r, c) in rg_can_match.iter_mut().zip(child.rg_can_match.iter()) { + *r |= c; + } + } + Self { rg_can_match, children: annotated_children } + } + BoolNode::Not(child) => { + let annotated_child = Self::build_from_bool_node(child, leaf_predicates, metadata, schema, rg_indices); + // NOT is conservatively all-true: negating stats-based pruning is unsound + // because stats give a superset, and inverting a superset is a subset. + native_bridge_common::log_debug!("StatsPruneTree: NOT node → all-true (conservative, cannot prune through negation)"); + Self { rg_can_match: vec![true; num], children: vec![annotated_child] } + } + BoolNode::Predicate(expr) => { + let key = Arc::as_ptr(expr) as *const () as usize; + let rg_can_match = match leaf_predicates.get(&key) { + Some(pp) => eval_leaf(pp, metadata, schema, rg_indices), + None => vec![true; num], + }; + Self { rg_can_match, children: vec![] } + } + BoolNode::DelegationPossible { original_expr, .. } => { + let key = Arc::as_ptr(original_expr) as *const () as usize; + let rg_can_match = match leaf_predicates.get(&key) { + Some(pp) => eval_leaf(pp, metadata, schema, rg_indices), + None => vec![true; num], + }; + Self { rg_can_match, children: vec![] } + } + BoolNode::Collector { .. } => { + // Collectors have no column stats — always all-true. + native_bridge_common::log_debug!("StatsPruneTree: Collector node → all-true (no column stats available)"); + Self { rg_can_match: vec![true; num], children: vec![] } + } + } + } +} + #[cfg(test)] mod tests { use super::*; @@ -1203,4 +1415,126 @@ mod tests { ); assert_eq!(count_rows_kept(&sel1), 0, "RG1: no page has values < 4"); } + + /// Matches the tree from the doc comment on `StatsPruneTree`: + /// + /// ```text + /// AND + /// [00001] + /// / | \ + /// OR OR AND + /// [01101] [11111] [00011] + /// / \ / | \ / \ + /// AND df AND luc df df NOT + /// [01100] [00001] [00100] [11111] [11110] [00011] [11111] + /// / \ / \ | + /// df df df df df + /// [11100] [01111] [11100] [00111] [11000] + /// ``` + /// + /// Uses 5 RGs (price ranges [0..9],[10..19],[20..29],[30..39],[40..49]) + /// to verify the full shape and bitset propagation. + #[test] + fn stats_prune_tree_five_leaf_shape() { + use crate::indexed_table::bool_tree::BoolNode; + + // 5 RGs: price [0..9], [10..19], [20..29], [30..39], [40..49] + let schema = Arc::new(Schema::new(vec![Field::new("price", DataType::Int32, false)])); + let tmp = NamedTempFile::new().unwrap(); + let props = WriterProperties::builder() + .set_max_row_group_size(10) + .set_statistics_enabled(EnabledStatistics::Chunk) + .build(); + let mut w = ArrowWriter::try_new(tmp.reopen().unwrap(), schema.clone(), Some(props)).unwrap(); + for i in 0..5i32 { + let vals: Vec = (i * 10..(i + 1) * 10).collect(); + let batch = RecordBatch::try_new(schema.clone(), vec![Arc::new(Int32Array::from(vals))]).unwrap(); + w.write(&batch).unwrap(); + } + w.close().unwrap(); + let meta = ArrowReaderMetadata::load(&tmp.reopen().unwrap(), ArrowReaderOptions::new()).unwrap(); + let arc_meta = meta.metadata().clone(); + assert_eq!(arc_meta.num_row_groups(), 5); + let rg_indices: Vec = (0..5).collect(); + + // Leaf predicates and their expected rg_can_match bitsets: + let p1 = pred_leaf("price", Operator::Lt, 30, &schema); // [11100] + let p2 = pred_leaf("price", Operator::GtEq, 10, &schema); // [01111] + let p3 = pred_leaf("price", Operator::Gt, 45, &schema); // [00001] + let p4 = pred_leaf("price", Operator::Lt, 25, &schema); // [11100] + let p5 = pred_leaf("price", Operator::GtEq, 20, &schema); // [00111] + let p6 = pred_leaf("price", Operator::Lt, 40, &schema); // [11110] + let p7 = pred_leaf("price", Operator::Gt, 30, &schema); // [00011] + let p8 = pred_leaf("price", Operator::Lt, 15, &schema); // [11000] + + // Tree: AND(OR(AND(p1,p2), p3), OR(AND(p4,p5), Collector, p6), AND(p7, NOT(p8))) + let collector = BoolNode::Collector { annotation_id: 0 }; + let tree = BoolNode::And(vec![ + BoolNode::Or(vec![ + BoolNode::And(vec![p1.clone(), p2.clone()]), + p3.clone(), + ]), + BoolNode::Or(vec![ + BoolNode::And(vec![p4.clone(), p5.clone()]), + collector, + p6.clone(), + ]), + BoolNode::And(vec![ + p7.clone(), + BoolNode::Not(Box::new(p8.clone())), + ]), + ]); + + // Build PruningPredicates for all df leaves. + let mut leaf_predicates: HashMap> = HashMap::new(); + for node in [&p1, &p2, &p3, &p4, &p5, &p6, &p7, &p8] { + if let BoolNode::Predicate(expr) = node { + let key = Arc::as_ptr(expr) as *const () as usize; + let pp = build_pruning_predicate(expr, schema.clone()).unwrap(); + leaf_predicates.insert(key, pp); + } + } + + let spt = StatsPruneTree::build_from_bool_node( + &tree, &leaf_predicates, &arc_meta, &schema, &rg_indices, + ); + + // Verify bottom-up: + // AND(p1[11100], p2[01111]) = [01100] + // OR₁(AND[01100], p3[00001]) = [01101] + // AND(p4[11100], p5[00111]) = [00100] + // OR₂(AND[00100], luc[11111], p6[11110]) = [11111] ← luc dominates + // NOT(p8[11000]) = [11111] ← conservative + // AND₃(p7[00011], NOT[11111]) = [00011] + // Root AND(OR₁[01101], OR₂[11111], AND₃[00011]) = [00001] + + // Root + assert_eq!(spt.rg_can_match, vec![false, false, false, false, true]); + assert_eq!(spt.children.len(), 3); + + // Child 0: OR₁ + assert_eq!(spt.children[0].rg_can_match, vec![false, true, true, false, true]); + assert_eq!(spt.children[0].children.len(), 2); + // OR₁/AND + assert_eq!(spt.children[0].children[0].rg_can_match, vec![false, true, true, false, false]); + // OR₁/AND/p1 + assert_eq!(spt.children[0].children[0].children[0].rg_can_match, vec![true, true, true, false, false]); + // OR₁/AND/p2 + assert_eq!(spt.children[0].children[0].children[1].rg_can_match, vec![false, true, true, true, true]); + // OR₁/p3 + assert_eq!(spt.children[0].children[1].rg_can_match, vec![false, false, false, false, true]); + + // Child 1: OR₂ — luc makes it all-true + assert_eq!(spt.children[1].rg_can_match, vec![true, true, true, true, true]); + + // Child 2: AND₃ + assert_eq!(spt.children[2].rg_can_match, vec![false, false, false, true, true]); + assert_eq!(spt.children[2].children.len(), 2); + // AND₃/p7 + assert_eq!(spt.children[2].children[0].rg_can_match, vec![false, false, false, true, true]); + // AND₃/NOT → always true + assert_eq!(spt.children[2].children[1].rg_can_match, vec![true, true, true, true, true]); + // AND₃/NOT/p8 + assert_eq!(spt.children[2].children[1].children[0].rg_can_match, vec![true, true, false, false, false]); + } } diff --git a/sandbox/plugins/analytics-backend-datafusion/rust/src/indexed_table/stream.rs b/sandbox/plugins/analytics-backend-datafusion/rust/src/indexed_table/stream.rs index f345c2707008f..2b1f2e5890b8b 100644 --- a/sandbox/plugins/analytics-backend-datafusion/rust/src/indexed_table/stream.rs +++ b/sandbox/plugins/analytics-backend-datafusion/rust/src/indexed_table/stream.rs @@ -33,6 +33,7 @@ use std::pin::Pin; use std::sync::Arc; use std::task::{Context, Poll}; +use native_bridge_common::log_debug; use datafusion::arrow::array::{Array, BooleanArray, UInt64Array}; use datafusion::arrow::compute::filter_record_batch; use datafusion::arrow::datatypes::{DataType, Field, Schema, SchemaRef}; @@ -482,6 +483,8 @@ struct IndexedStream { mask_offset: usize, batch_offset: usize, finished: bool, + /// Wall-clock start time for per-segment timing. Set on first poll. + stream_start: Option, metadata: Arc, predicate: Option>, initialized: bool, @@ -575,6 +578,7 @@ impl IndexedStream { mask_offset: 0, batch_offset: 0, finished: false, + stream_start: None, metadata, predicate, initialized: false, @@ -777,6 +781,7 @@ impl Stream for IndexedStream { let poll_start = Instant::now(); if !self.initialized { + self.stream_start = Some(Instant::now()); self.index_reader.init_prefetch(); self.initialized = true; } @@ -809,6 +814,7 @@ impl IndexedStream { // 2. If upstream is done and coalescer has drained, we're done. if self.coalescer_finished && self.batch_coalescer.is_empty() { + log_debug!("[scf-segment-done] file={} row_groups={} elapsed={:?}", self.object_path.filename().unwrap_or("?"), self.index_reader.row_groups.len(), self.stream_start.unwrap().elapsed()); return Poll::Ready(None); } @@ -828,6 +834,7 @@ impl IndexedStream { if self.coalescer_finished { // Unreachable in practice — step 1 already drained or // step 2 already returned. Defensive. + log_debug!("[scf-segment-done] file={} row_groups={} elapsed={:?}", self.object_path.filename().unwrap_or("?"), self.index_reader.row_groups.len(), self.stream_start.unwrap().elapsed()); return Poll::Ready(None); } diff --git a/sandbox/plugins/analytics-backend-datafusion/rust/src/indexed_table/table_provider.rs b/sandbox/plugins/analytics-backend-datafusion/rust/src/indexed_table/table_provider.rs index ae7644dccf522..033119297347c 100644 --- a/sandbox/plugins/analytics-backend-datafusion/rust/src/indexed_table/table_provider.rs +++ b/sandbox/plugins/analytics-backend-datafusion/rust/src/indexed_table/table_provider.rs @@ -40,16 +40,19 @@ use datafusion::physical_plan::execution_plan::{Boundedness, EmissionType}; use datafusion::physical_plan::metrics::{ExecutionPlanMetricsSet, MetricsSet}; use datafusion::physical_plan::{DisplayAs, DisplayFormatType, ExecutionPlan, PlanProperties}; use datafusion_common::DataFusionError; +use datafusion::physical_optimizer::pruning::PruningPredicate; use datafusion::physical_plan::stream::RecordBatchStreamAdapter; use futures::StreamExt; +use super::bool_tree::BoolNode; use super::eval::RowGroupBitsetSource; use super::metrics::PartitionMetrics; use super::partitioning::{compute_assignments, PartitionAssignment, SegmentChunk, SegmentLayout}; use super::stream::{FilterStrategy, IndexedExec, RowGroupInfo}; use crate::datafusion_query_config::DatafusionQueryConfig; use crate::indexed_table::metrics::StreamMetrics; +use crate::indexed_table::page_pruner::StatsPruneTree; use std::collections::HashSet; /// Info about a segment and its corresponding parquet file. @@ -95,6 +98,7 @@ pub type EvaluatorFactory = Arc< &SegmentFileInfo, &SegmentChunk, &StreamMetrics, + Option<&StatsPruneTree>, ) -> Result, String> + Send + Sync, @@ -133,6 +137,13 @@ pub struct IndexedTableConfig { /// from position (global_base + rg.first_row + position_in_rg) instead of /// being read from parquet. Other projected columns are read normally. pub emit_row_ids: bool, + /// Query-level data for building StatsPruneTree per segment. + /// (BoolNode tree, prebuilt PruningPredicates keyed by Arc ptr, schema) + pub prune_tree_config: Option<( + BoolNode, + Arc>>, + SchemaRef, + )>, } /// Table provider. Returns a `QueryShardExec` that fans out across chunks. @@ -493,8 +504,24 @@ impl ExecutionPlan for QueryShardExec { continue; } + // Build stats prune tree for segment/RG/subtree-level pruning. + let stats_prune_tree = self.config.prune_tree_config.as_ref().map(|(tree, preds, schema)| { + let rg_indices: Vec = row_groups.iter().map(|rg| rg.index).collect(); + StatsPruneTree::build_from_bool_node( + tree, preds, &segment.metadata, schema, &rg_indices, + ) + }); + + // Segment-level skip: if no RG in the chunk can match, skip entirely. + if let Some(ref spt) = stats_prune_tree { + if !spt.rg_can_match.iter().any(|&k| k) { + native_bridge_common::log_debug!("[segment-skip] skipping chunk — pruned by segment-level stats"); + continue; + } + } + // Build evaluator for this chunk. - let evaluator = (self.config.evaluator_factory)(segment, chunk, &stream_metrics) + let evaluator = (self.config.evaluator_factory)(segment, chunk, &stream_metrics, stats_prune_tree.as_ref()) .map_err(|e| DataFusionError::External(e.into()))?; let props = Arc::new(PlanProperties::new( @@ -631,13 +658,14 @@ mod tests { store: Arc::new(object_store::local::LocalFileSystem::new()), store_url: datafusion::execution::object_store::ObjectStoreUrl::local_filesystem(), // Evaluator factory would never be invoked for this test (no segments). - evaluator_factory: Arc::new(|_, _, _| unreachable!()), + evaluator_factory: Arc::new(|_, _, _, _| unreachable!()), pushdown_predicate: None, query_config: std::sync::Arc::new( crate::datafusion_query_config::DatafusionQueryConfig::test_default(), ), predicate_columns: vec![], emit_row_ids: false, + prune_tree_config: None, } } diff --git a/sandbox/plugins/analytics-backend-datafusion/rust/src/indexed_table/tests_e2e/constant_predicate.rs b/sandbox/plugins/analytics-backend-datafusion/rust/src/indexed_table/tests_e2e/constant_predicate.rs index 4534d1a532f2a..d40e53a83296e 100644 --- a/sandbox/plugins/analytics-backend-datafusion/rust/src/indexed_table/tests_e2e/constant_predicate.rs +++ b/sandbox/plugins/analytics-backend-datafusion/rust/src/indexed_table/tests_e2e/constant_predicate.rs @@ -80,13 +80,14 @@ async fn run_constant_residual(residual: Arc) -> usize { let factory: EvaluatorFactory = { let schema = schema.clone(); let residual = Arc::clone(&residual); - Arc::new(move |segment: &SegmentFileInfo, _chunk, stream_metrics| { + Arc::new(move |segment: &SegmentFileInfo, _chunk, stream_metrics, _stats_prune_tree| { let pruner = Arc::new(PagePruner::new(&schema, Arc::clone(&segment.metadata))); let eval: Arc = Arc::new(PredicateOnlyEvaluator::new( pruner, None, Some(Arc::clone(&residual)), Some(PagePruneMetrics::from_stream_metrics(stream_metrics)), + None, )); Ok(eval) }) @@ -109,7 +110,7 @@ async fn run_constant_residual(residual: Arc) -> usize { pushdown_predicate: None, query_config: std::sync::Arc::new(qc), predicate_columns: vec![], - emit_row_ids: false, + emit_row_ids: false, prune_tree_config: None, })); let ctx = SessionContext::new(); diff --git a/sandbox/plugins/analytics-backend-datafusion/rust/src/indexed_table/tests_e2e/dynamic_filter_pushdown.rs b/sandbox/plugins/analytics-backend-datafusion/rust/src/indexed_table/tests_e2e/dynamic_filter_pushdown.rs index 947bd92af5d57..9d58a8363bfa0 100644 --- a/sandbox/plugins/analytics-backend-datafusion/rust/src/indexed_table/tests_e2e/dynamic_filter_pushdown.rs +++ b/sandbox/plugins/analytics-backend-datafusion/rust/src/indexed_table/tests_e2e/dynamic_filter_pushdown.rs @@ -135,7 +135,7 @@ async fn run_indexed( let factory: super::super::table_provider::EvaluatorFactory = { let schema = schema.clone(); - Arc::new(move |segment, _chunk, _stream_metrics| { + Arc::new(move |segment, _chunk, _stream_metrics, _stats_prune_tree| { let pruner = Arc::new(PagePruner::new(&schema, Arc::clone(&segment.metadata))); let collector: Arc = Arc::new(MatchAllCollector); let eval: Arc = Arc::new( @@ -154,6 +154,7 @@ async fn run_indexed( ), 0, None, + None, ), ); Ok(eval) @@ -183,6 +184,7 @@ async fn run_indexed( query_config: std::sync::Arc::new(qc), predicate_columns: vec![], emit_row_ids: false, + prune_tree_config: None, })); let ctx = SessionContext::new(); diff --git a/sandbox/plugins/analytics-backend-datafusion/rust/src/indexed_table/tests_e2e/fuzz/delegation.rs b/sandbox/plugins/analytics-backend-datafusion/rust/src/indexed_table/tests_e2e/fuzz/delegation.rs index 0bf0b60789af4..7edc2f3e7f4c7 100644 --- a/sandbox/plugins/analytics-backend-datafusion/rust/src/indexed_table/tests_e2e/fuzz/delegation.rs +++ b/sandbox/plugins/analytics-backend-datafusion/rust/src/indexed_table/tests_e2e/fuzz/delegation.rs @@ -405,7 +405,7 @@ pub(in crate::indexed_table::tests_e2e) async fn execute_delegation_tree( let factory = Arc::clone(&factory); let provider_locks = Arc::clone(&provider_locks); let schema = loaded.schema.clone(); - Arc::new(move |segment, _chunk, stream_metrics| { + Arc::new(move |segment, _chunk, stream_metrics, _stats_prune_tree| { let pruner = Arc::new(PagePruner::new(&schema, Arc::clone(&segment.metadata))); let eval: Arc = Arc::new(SingleCollectorEvaluator::new( Some(Arc::clone(&correctness)), @@ -420,6 +420,7 @@ pub(in crate::indexed_table::tests_e2e) async fn execute_delegation_tree( Arc::clone(&factory), 0, None, + None, )); Ok(eval) }) @@ -467,7 +468,7 @@ pub(in crate::indexed_table::tests_e2e) async fn execute_delegation_tree( pushdown_predicate: Some(Arc::clone(&residual_physical)), query_config: Arc::new(qc), predicate_columns: pred_cols, - emit_row_ids: false, + emit_row_ids: false, prune_tree_config: None, })); let ctx = SessionContext::new(); diff --git a/sandbox/plugins/analytics-backend-datafusion/rust/src/indexed_table/tests_e2e/fuzz/harness.rs b/sandbox/plugins/analytics-backend-datafusion/rust/src/indexed_table/tests_e2e/fuzz/harness.rs index d1c7f3b576538..c0be72ee1e9d5 100644 --- a/sandbox/plugins/analytics-backend-datafusion/rust/src/indexed_table/tests_e2e/fuzz/harness.rs +++ b/sandbox/plugins/analytics-backend-datafusion/rust/src/indexed_table/tests_e2e/fuzz/harness.rs @@ -234,7 +234,7 @@ pub(in crate::indexed_table::tests_e2e) async fn execute_tree_with_plan_pushdown }) .collect(), ); - Arc::new(move |segment, _chunk, stream_metrics| { + Arc::new(move |segment, _chunk, stream_metrics, _stats_prune_tree| { let resolved = tree.resolve(&per_leaf)?; let pruner = Arc::new(PagePruner::new(&schema, Arc::clone(&segment.metadata))); let eval: Arc = Arc::new(TreeBitsetSource { @@ -264,6 +264,7 @@ pub(in crate::indexed_table::tests_e2e) async fn execute_tree_with_plan_pushdown crate::indexed_table::eval::CollectorCallStrategy::FullRange, crate::indexed_table::eval::CollectorCallStrategy::PageRangeSplit, ][seed as usize % 3], + stats_prune_tree: None, }); Ok(eval) }) @@ -290,7 +291,7 @@ pub(in crate::indexed_table::tests_e2e) async fn execute_tree_with_plan_pushdown pushdown_predicate: None, query_config: Arc::new(qc), predicate_columns: collect_predicate_column_indices(&bool_tree), - emit_row_ids: false, + emit_row_ids: false, prune_tree_config: None, })); let ctx = SessionContext::new(); @@ -384,7 +385,7 @@ pub(in crate::indexed_table::tests_e2e) async fn execute_tree_single_collector( let schema = schema.clone(); let residual_pp = residual_pp.clone(); let residual_physical = residual_physical.clone(); - Arc::new(move |segment, _chunk, stream_metrics| { + Arc::new(move |segment, _chunk, stream_metrics, _stats_prune_tree| { let pruner = Arc::new(PagePruner::new(&schema, Arc::clone(&segment.metadata))); let eval: Arc = Arc::new(SingleCollectorEvaluator::new( Some(Arc::clone(&collector)), @@ -403,6 +404,7 @@ pub(in crate::indexed_table::tests_e2e) async fn execute_tree_single_collector( std::sync::Arc::new(crate::indexed_table::eval::single_collector::FfmDelegatedBackendCollectorFactory), 0, None, + None, )); let _ = segment; Ok(eval) @@ -471,7 +473,7 @@ async fn run_single_collector_query( pushdown_predicate, query_config: Arc::new(qc), predicate_columns: pred_cols, - emit_row_ids: false, + emit_row_ids: false, prune_tree_config: None, })); let ctx = SessionContext::new(); ctx.register_table("t", provider).unwrap(); @@ -680,7 +682,7 @@ async fn run_with_factory_plan( pushdown_predicate, query_config: Arc::new(qc), predicate_columns: vec![], - emit_row_ids: false, + emit_row_ids: false, prune_tree_config: None, })); let ctx = SessionContext::new(); ctx.register_table("t", provider).unwrap(); diff --git a/sandbox/plugins/analytics-backend-datafusion/rust/src/indexed_table/tests_e2e/mod.rs b/sandbox/plugins/analytics-backend-datafusion/rust/src/indexed_table/tests_e2e/mod.rs index 7ada00de3f869..839558c5dc1bf 100644 --- a/sandbox/plugins/analytics-backend-datafusion/rust/src/indexed_table/tests_e2e/mod.rs +++ b/sandbox/plugins/analytics-backend-datafusion/rust/src/indexed_table/tests_e2e/mod.rs @@ -248,7 +248,7 @@ async fn run_tree_and_plan( let per_leaf = per_leaf.clone(); let tree = Arc::clone(&tree); let schema = schema.clone(); - Arc::new(move |segment, _chunk, _stream_metrics| { + Arc::new(move |segment, _chunk, _stream_metrics, _stats_prune_tree| { let resolved = tree.resolve(&per_leaf)?; let pruner = Arc::new(PagePruner::new(&schema, Arc::clone(&segment.metadata))); let eval: Arc = Arc::new(TreeBitsetSource { @@ -270,6 +270,7 @@ async fn run_tree_and_plan( ), ), collector_strategy: crate::indexed_table::eval::CollectorCallStrategy::TightenOuterBounds, + stats_prune_tree: None, }); Ok(eval) }) @@ -296,7 +297,7 @@ async fn run_tree_and_plan( pushdown_predicate: None, query_config: std::sync::Arc::new(qc), predicate_columns: vec![], - emit_row_ids: false, + emit_row_ids: false, prune_tree_config: None, })); let ctx = SessionContext::new(); diff --git a/sandbox/plugins/analytics-backend-datafusion/rust/src/indexed_table/tests_e2e/multi_segment.rs b/sandbox/plugins/analytics-backend-datafusion/rust/src/indexed_table/tests_e2e/multi_segment.rs index 40fbee4f0564e..9c294066a20c2 100644 --- a/sandbox/plugins/analytics-backend-datafusion/rust/src/indexed_table/tests_e2e/multi_segment.rs +++ b/sandbox/plugins/analytics-backend-datafusion/rust/src/indexed_table/tests_e2e/multi_segment.rs @@ -144,7 +144,7 @@ async fn run_two_segment_query( { let per_segment_matches = Arc::clone(&per_segment_matches); let schema = schema.clone(); - Arc::new(move |segment, _chunk, _stream_metrics| { + Arc::new(move |segment, _chunk, _stream_metrics, _stats_prune_tree| { let matching = per_segment_matches .get(segment.writer_generation as usize) .cloned() @@ -161,6 +161,7 @@ async fn run_two_segment_query( std::sync::Arc::new(crate::indexed_table::eval::single_collector::FfmDelegatedBackendCollectorFactory), 0, None, + None, ), ); Ok(eval) @@ -184,7 +185,7 @@ async fn run_two_segment_query( pushdown_predicate: None, query_config: std::sync::Arc::new(qc), predicate_columns: vec![], - emit_row_ids: false, + emit_row_ids: false, prune_tree_config: None, })); let ctx = SessionContext::new(); @@ -346,7 +347,7 @@ async fn run_two_segment_query_witness( let schema = schema.clone(); let in_flight = Arc::clone(&in_flight); let max_in_flight = Arc::clone(&max_in_flight); - Arc::new(move |segment, _chunk, _stream_metrics| { + Arc::new(move |segment, _chunk, _stream_metrics, _stats_prune_tree| { let matching = per_segment_matches .get(segment.writer_generation as usize) .cloned() @@ -367,6 +368,7 @@ async fn run_two_segment_query_witness( std::sync::Arc::new(crate::indexed_table::eval::single_collector::FfmDelegatedBackendCollectorFactory), 0, None, + None, ), ); Ok(eval) @@ -390,7 +392,7 @@ async fn run_two_segment_query_witness( pushdown_predicate: None, query_config: std::sync::Arc::new(qc), predicate_columns: vec![], - emit_row_ids: false, + emit_row_ids: false, prune_tree_config: None, })); let ctx = SessionContext::new(); @@ -554,7 +556,7 @@ async fn run_segments(specs: Vec, num_partitions: usize) -> Vec<(i32, S { let per_segment_matches = Arc::clone(&per_segment_matches); let schema = schema.clone(); - Arc::new(move |segment, _chunk, _stream_metrics| { + Arc::new(move |segment, _chunk, _stream_metrics, _stats_prune_tree| { let matching = per_segment_matches .get(segment.writer_generation as usize) .cloned() @@ -571,6 +573,7 @@ async fn run_segments(specs: Vec, num_partitions: usize) -> Vec<(i32, S std::sync::Arc::new(crate::indexed_table::eval::single_collector::FfmDelegatedBackendCollectorFactory), 0, None, + None, ), ); Ok(eval) @@ -594,7 +597,7 @@ async fn run_segments(specs: Vec, num_partitions: usize) -> Vec<(i32, S pushdown_predicate: None, query_config: std::sync::Arc::new(qc), predicate_columns: vec![], - emit_row_ids: false, + emit_row_ids: false, prune_tree_config: None, })); let ctx = SessionContext::new(); @@ -1046,7 +1049,7 @@ async fn run_wide_segments( let factory: super::super::table_provider::EvaluatorFactory = { let tree = Arc::clone(&tree); let schema = schema.clone(); - Arc::new(move |segment, _chunk, _stream_metrics| { + Arc::new(move |segment, _chunk, _stream_metrics, _stats_prune_tree| { // One (provider_key, collector) per Collector leaf — our trees // here use 1 collector leaf, so one pair. let leaf_count = tree.collector_leaf_count(); @@ -1074,6 +1077,7 @@ async fn run_wide_segments( pruning_predicates: std::sync::Arc::new(std::collections::HashMap::new()), page_prune_metrics: None, collector_strategy: crate::indexed_table::eval::CollectorCallStrategy::TightenOuterBounds, + stats_prune_tree: None, }, ); Ok(eval) @@ -1097,7 +1101,7 @@ async fn run_wide_segments( pushdown_predicate: None, query_config: std::sync::Arc::new(qc), predicate_columns: vec![], - emit_row_ids: false, + emit_row_ids: false, prune_tree_config: None, })); let ctx = SessionContext::new(); diff --git a/sandbox/plugins/analytics-backend-datafusion/rust/src/indexed_table/tests_e2e/null_columns.rs b/sandbox/plugins/analytics-backend-datafusion/rust/src/indexed_table/tests_e2e/null_columns.rs index 9674000c51b11..545a559b4d2e1 100644 --- a/sandbox/plugins/analytics-backend-datafusion/rust/src/indexed_table/tests_e2e/null_columns.rs +++ b/sandbox/plugins/analytics-backend-datafusion/rust/src/indexed_table/tests_e2e/null_columns.rs @@ -346,7 +346,7 @@ async fn assert_engine_matches_reference_null(name: &str, tree: NT) { let per_leaf = per_leaf.clone(); let tree = Arc::clone(&tree); let schema = schema.clone(); - Arc::new(move |segment, _chunk, _stream_metrics| { + Arc::new(move |segment, _chunk, _stream_metrics, _stats_prune_tree| { let resolved = tree.resolve(&per_leaf)?; let pruner = Arc::new(PagePruner::new(&schema, Arc::clone(&segment.metadata))); let eval: Arc = Arc::new(TreeBitsetSource { @@ -361,6 +361,7 @@ async fn assert_engine_matches_reference_null(name: &str, tree: NT) { page_prune_metrics: None, collector_strategy: crate::indexed_table::eval::CollectorCallStrategy::TightenOuterBounds, + stats_prune_tree: None, }); Ok(eval) }) @@ -381,7 +382,7 @@ async fn assert_engine_matches_reference_null(name: &str, tree: NT) { pushdown_predicate: None, query_config: std::sync::Arc::new(qc), predicate_columns: vec![], - emit_row_ids: false, + emit_row_ids: false, prune_tree_config: None, })); let ctx = SessionContext::new(); ctx.register_table("t", provider).unwrap(); diff --git a/sandbox/plugins/analytics-backend-datafusion/rust/src/indexed_table/tests_e2e/page_pruning.rs b/sandbox/plugins/analytics-backend-datafusion/rust/src/indexed_table/tests_e2e/page_pruning.rs index bafd6a0c400a2..3598b8f833ef2 100644 --- a/sandbox/plugins/analytics-backend-datafusion/rust/src/indexed_table/tests_e2e/page_pruning.rs +++ b/sandbox/plugins/analytics-backend-datafusion/rust/src/indexed_table/tests_e2e/page_pruning.rs @@ -307,7 +307,7 @@ async fn run_bitmap_tree(tree: BoolNode) -> (Vec, Arc) { let tree = Arc::clone(&tree); let schema = schema.clone(); let pp_map = Arc::clone(&pp_map); - Arc::new(move |segment, _chunk, sm| { + Arc::new(move |segment, _chunk, sm, _stats_prune_tree| { let resolved = tree.resolve(&per_leaf)?; let pruner = Arc::new(PagePruner::new(&schema, Arc::clone(&segment.metadata))); let eval: Arc = Arc::new(TreeBitsetSource { @@ -326,6 +326,7 @@ async fn run_bitmap_tree(tree: BoolNode) -> (Vec, Arc) { ), collector_strategy: crate::indexed_table::eval::CollectorCallStrategy::TightenOuterBounds, + stats_prune_tree: None, }); Ok(eval) }) @@ -348,7 +349,7 @@ async fn run_single_collector( let schema = schema.clone(); let residual_pp = residual_pp.clone(); let residual_expr = Arc::clone(&residual_expr); - Arc::new(move |segment, _chunk, sm| { + Arc::new(move |segment, _chunk, sm, _stats_prune_tree| { let pruner = Arc::new(PagePruner::new(&schema, Arc::clone(&segment.metadata))); let eval: Arc = Arc::new(SingleCollectorEvaluator::new( Some(collector_for_tag(collector_tag)), @@ -363,6 +364,7 @@ async fn run_single_collector( std::sync::Arc::new(crate::indexed_table::eval::single_collector::FfmDelegatedBackendCollectorFactory), 0, None, + None, )); Ok(eval) }) @@ -393,7 +395,7 @@ async fn execute_and_collect( pushdown_predicate: None, query_config: Arc::new(qc), predicate_columns: vec![], - emit_row_ids: false, + emit_row_ids: false, prune_tree_config: None, })); let ctx = SessionContext::new(); diff --git a/sandbox/plugins/analytics-backend-datafusion/rust/src/indexed_table/tests_e2e/qtf_fetch_phase.rs b/sandbox/plugins/analytics-backend-datafusion/rust/src/indexed_table/tests_e2e/qtf_fetch_phase.rs index 35d71c045ee13..a430ba842d94f 100644 --- a/sandbox/plugins/analytics-backend-datafusion/rust/src/indexed_table/tests_e2e/qtf_fetch_phase.rs +++ b/sandbox/plugins/analytics-backend-datafusion/rust/src/indexed_table/tests_e2e/qtf_fetch_phase.rs @@ -76,7 +76,7 @@ async fn query_phase(tree: BoolNode) -> Vec { let per_leaf = per_leaf.clone(); let tree = Arc::clone(&tree); let schema = schema.clone(); - Arc::new(move |segment, _chunk, _stream_metrics| { + Arc::new(move |segment, _chunk, _stream_metrics, _stats_prune_tree| { let resolved = tree.resolve(&per_leaf)?; let pruner = Arc::new(PagePruner::new(&schema, Arc::clone(&segment.metadata))); let eval: Arc = Arc::new(TreeBitsetSource { @@ -99,6 +99,7 @@ async fn query_phase(tree: BoolNode) -> Vec { ), collector_strategy: crate::indexed_table::eval::CollectorCallStrategy::TightenOuterBounds, + stats_prune_tree: None, }); Ok(eval) }) @@ -122,7 +123,7 @@ async fn query_phase(tree: BoolNode) -> Vec { qc }), predicate_columns: vec![0, 1, 2, 3], - emit_row_ids: true, + emit_row_ids: true, prune_tree_config: None, })); let ctx = SessionContext::new(); diff --git a/sandbox/plugins/analytics-backend-datafusion/rust/src/indexed_table/tests_e2e/row_id_emission.rs b/sandbox/plugins/analytics-backend-datafusion/rust/src/indexed_table/tests_e2e/row_id_emission.rs index a01c0b26aed90..0719b6d380dd7 100644 --- a/sandbox/plugins/analytics-backend-datafusion/rust/src/indexed_table/tests_e2e/row_id_emission.rs +++ b/sandbox/plugins/analytics-backend-datafusion/rust/src/indexed_table/tests_e2e/row_id_emission.rs @@ -66,7 +66,7 @@ async fn run_tree_row_ids(tree: BoolNode) -> Vec { let per_leaf = per_leaf.clone(); let tree = Arc::clone(&tree); let schema = schema.clone(); - Arc::new(move |segment, _chunk, _stream_metrics| { + Arc::new(move |segment, _chunk, _stream_metrics, _stats_prune_tree| { let resolved = tree.resolve(&per_leaf)?; let pruner = Arc::new(PagePruner::new(&schema, Arc::clone(&segment.metadata))); let eval: Arc = Arc::new(TreeBitsetSource { @@ -89,6 +89,7 @@ async fn run_tree_row_ids(tree: BoolNode) -> Vec { ), collector_strategy: crate::indexed_table::eval::CollectorCallStrategy::TightenOuterBounds, + stats_prune_tree: None, }); Ok(eval) }) @@ -112,7 +113,7 @@ async fn run_tree_row_ids(tree: BoolNode) -> Vec { qc }), predicate_columns: vec![0, 1, 2, 3], - emit_row_ids: true, + emit_row_ids: true, prune_tree_config: None, })); let ctx = SessionContext::new(); @@ -230,7 +231,7 @@ async fn run_tree_row_ids_with_global_base(tree: BoolNode, global_base: u64) -> let per_leaf = per_leaf.clone(); let tree = Arc::clone(&tree); let schema = schema.clone(); - Arc::new(move |segment, _chunk, _stream_metrics| { + Arc::new(move |segment, _chunk, _stream_metrics, _stats_prune_tree| { let resolved = tree.resolve(&per_leaf)?; let pruner = Arc::new(PagePruner::new(&schema, Arc::clone(&segment.metadata))); let eval: Arc = Arc::new(TreeBitsetSource { @@ -253,6 +254,7 @@ async fn run_tree_row_ids_with_global_base(tree: BoolNode, global_base: u64) -> ), collector_strategy: crate::indexed_table::eval::CollectorCallStrategy::TightenOuterBounds, + stats_prune_tree: None, }); Ok(eval) }) @@ -276,7 +278,7 @@ async fn run_tree_row_ids_with_global_base(tree: BoolNode, global_base: u64) -> qc }), predicate_columns: vec![0, 1, 2, 3], - emit_row_ids: true, + emit_row_ids: true, prune_tree_config: None, })); let ctx = SessionContext::new(); @@ -492,7 +494,7 @@ async fn test_row_id_with_data_columns() { let per_leaf = per_leaf.clone(); let tree = Arc::clone(&tree); let schema = schema.clone(); - Arc::new(move |segment, _chunk, _stream_metrics| { + Arc::new(move |segment, _chunk, _stream_metrics, _stats_prune_tree| { let resolved = tree.resolve(&per_leaf)?; let pruner = Arc::new(PagePruner::new(&schema, Arc::clone(&segment.metadata))); let eval: Arc = Arc::new(TreeBitsetSource { @@ -515,6 +517,7 @@ async fn test_row_id_with_data_columns() { ), collector_strategy: crate::indexed_table::eval::CollectorCallStrategy::TightenOuterBounds, + stats_prune_tree: None, }); Ok(eval) }) @@ -538,7 +541,7 @@ async fn test_row_id_with_data_columns() { qc }), predicate_columns: vec![0, 1, 2, 3], - emit_row_ids: true, + emit_row_ids: true, prune_tree_config: None, })); let ctx = SessionContext::new(); @@ -739,7 +742,7 @@ async fn run_two_segments_row_ids(tree: BoolNode) -> Vec { let per_leaf = per_leaf.clone(); let tree = Arc::clone(&tree); let schema = schema.clone(); - Arc::new(move |segment, _chunk, _stream_metrics| { + Arc::new(move |segment, _chunk, _stream_metrics, _stats_prune_tree| { let resolved = tree.resolve(&per_leaf)?; let pruner = Arc::new(PagePruner::new(&schema, Arc::clone(&segment.metadata))); let eval: Arc = Arc::new(TreeBitsetSource { @@ -762,6 +765,7 @@ async fn run_two_segments_row_ids(tree: BoolNode) -> Vec { ), collector_strategy: crate::indexed_table::eval::CollectorCallStrategy::TightenOuterBounds, + stats_prune_tree: None, }); Ok(eval) }) @@ -785,7 +789,7 @@ async fn run_two_segments_row_ids(tree: BoolNode) -> Vec { qc }), predicate_columns: vec![0, 1, 2, 3], - emit_row_ids: true, + emit_row_ids: true, prune_tree_config: None, })); let ctx = SessionContext::new(); diff --git a/sandbox/plugins/analytics-backend-datafusion/rust/src/indexed_table/tests_e2e/schema_drift.rs b/sandbox/plugins/analytics-backend-datafusion/rust/src/indexed_table/tests_e2e/schema_drift.rs index e35fed2e83443..d356388a7505a 100644 --- a/sandbox/plugins/analytics-backend-datafusion/rust/src/indexed_table/tests_e2e/schema_drift.rs +++ b/sandbox/plugins/analytics-backend-datafusion/rust/src/indexed_table/tests_e2e/schema_drift.rs @@ -103,7 +103,7 @@ async fn run_missing_col_tree(tree_bool: BoolNode) -> usize { let factory: super::super::table_provider::EvaluatorFactory = { let tree = Arc::clone(&tree); let schema = schema.clone(); - Arc::new(move |segment, _chunk, _stream_metrics| { + Arc::new(move |segment, _chunk, _stream_metrics, _stats_prune_tree| { let resolved = tree.resolve(&[])?; let pruner = Arc::new(PagePruner::new(&schema, Arc::clone(&segment.metadata))); let eval: Arc = Arc::new(TreeBitsetSource { @@ -117,6 +117,7 @@ async fn run_missing_col_tree(tree_bool: BoolNode) -> usize { pruning_predicates: std::sync::Arc::new(std::collections::HashMap::new()), page_prune_metrics: None, collector_strategy: crate::indexed_table::eval::CollectorCallStrategy::TightenOuterBounds, + stats_prune_tree: None, }); Ok(eval) }) @@ -136,7 +137,7 @@ async fn run_missing_col_tree(tree_bool: BoolNode) -> usize { pushdown_predicate: None, query_config: std::sync::Arc::new(qc), predicate_columns: vec![], - emit_row_ids: false, + emit_row_ids: false, prune_tree_config: None, })); let ctx = SessionContext::new(); ctx.register_table("t", provider).unwrap(); @@ -409,7 +410,7 @@ async fn query_with_mismatched_schema( let factory: super::super::table_provider::EvaluatorFactory = { let tree = Arc::clone(&tree); let schema = table_schema.clone(); - Arc::new(move |segment, _chunk, _stream_metrics| { + Arc::new(move |segment, _chunk, _stream_metrics, _stats_prune_tree| { let resolved = tree.resolve(&[])?; let pruner = Arc::new(PagePruner::new(&schema, Arc::clone(&segment.metadata))); let eval: Arc = Arc::new(TreeBitsetSource { @@ -424,6 +425,7 @@ async fn query_with_mismatched_schema( page_prune_metrics: None, collector_strategy: crate::indexed_table::eval::CollectorCallStrategy::TightenOuterBounds, + stats_prune_tree: None, }); Ok(eval) }) @@ -443,7 +445,7 @@ async fn query_with_mismatched_schema( pushdown_predicate: None, query_config: std::sync::Arc::new(qc), predicate_columns: vec![], - emit_row_ids: false, + emit_row_ids: false, prune_tree_config: None, })); let ctx = SessionContext::new(); ctx.register_table("t", provider).unwrap(); diff --git a/sandbox/plugins/analytics-backend-datafusion/rust/src/indexed_table/tests_e2e/streaming_at_scale.rs b/sandbox/plugins/analytics-backend-datafusion/rust/src/indexed_table/tests_e2e/streaming_at_scale.rs index d70ed285e5db9..21a1a3bd31c28 100644 --- a/sandbox/plugins/analytics-backend-datafusion/rust/src/indexed_table/tests_e2e/streaming_at_scale.rs +++ b/sandbox/plugins/analytics-backend-datafusion/rust/src/indexed_table/tests_e2e/streaming_at_scale.rs @@ -419,7 +419,7 @@ async fn run_large( let per_leaf = per_leaf.clone(); let tree = Arc::clone(&tree); let schema = schema.clone(); - Arc::new(move |segment, _chunk, _stream_metrics| { + Arc::new(move |segment, _chunk, _stream_metrics, _stats_prune_tree| { let resolved = tree.resolve(&per_leaf)?; let pruner = Arc::new(PagePruner::new(&schema, Arc::clone(&segment.metadata))); let eval: Arc = Arc::new(TreeBitsetSource { @@ -433,6 +433,7 @@ async fn run_large( pruning_predicates: std::sync::Arc::new(std::collections::HashMap::new()), page_prune_metrics: None, collector_strategy: crate::indexed_table::eval::CollectorCallStrategy::TightenOuterBounds, + stats_prune_tree: None, }); Ok(eval) }) @@ -453,7 +454,7 @@ async fn run_large( pushdown_predicate: None, query_config: std::sync::Arc::new(qc), predicate_columns: vec![], - emit_row_ids: false, + emit_row_ids: false, prune_tree_config: None, })); let ctx = SessionContext::new(); @@ -871,7 +872,7 @@ async fn run_large_partitioned( let per_leaf = per_leaf.clone(); let tree = Arc::clone(&tree); let schema = schema.clone(); - Arc::new(move |segment, _chunk, _stream_metrics| { + Arc::new(move |segment, _chunk, _stream_metrics, _stats_prune_tree| { let resolved = tree.resolve(&per_leaf)?; let pruner = Arc::new(PagePruner::new(&schema, Arc::clone(&segment.metadata))); let eval: Arc = Arc::new(TreeBitsetSource { @@ -885,6 +886,7 @@ async fn run_large_partitioned( pruning_predicates: std::sync::Arc::new(std::collections::HashMap::new()), page_prune_metrics: None, collector_strategy: crate::indexed_table::eval::CollectorCallStrategy::TightenOuterBounds, + stats_prune_tree: None, }); Ok(eval) }) @@ -904,7 +906,7 @@ async fn run_large_partitioned( pushdown_predicate: None, query_config: std::sync::Arc::new(qc), predicate_columns: vec![], - emit_row_ids: false, + emit_row_ids: false, prune_tree_config: None, })); let ctx = SessionContext::new(); ctx.register_table("t", provider).unwrap(); From 405c486d5e232dc3a050cdf98823cb741fe1d4f4 Mon Sep 17 00:00:00 2001 From: Andriy Redko Date: Thu, 11 Jun 2026 13:54:38 -0400 Subject: [PATCH 07/14] Update OpenTelemetry to 1.63.0 (#22111) Signed-off-by: Andriy Redko --- gradle/libs.versions.toml | 2 +- .../telemetry-otel/licenses/opentelemetry-api-1.62.0.jar.sha1 | 1 - .../telemetry-otel/licenses/opentelemetry-api-1.63.0.jar.sha1 | 1 + .../licenses/opentelemetry-api-incubator-1.62.0-alpha.jar.sha1 | 1 - .../licenses/opentelemetry-api-incubator-1.63.0-alpha.jar.sha1 | 1 + .../licenses/opentelemetry-common-1.62.0.jar.sha1 | 1 - .../licenses/opentelemetry-common-1.63.0.jar.sha1 | 1 + .../licenses/opentelemetry-context-1.62.0.jar.sha1 | 1 - .../licenses/opentelemetry-context-1.63.0.jar.sha1 | 1 + .../licenses/opentelemetry-exporter-common-1.62.0.jar.sha1 | 1 - .../licenses/opentelemetry-exporter-common-1.63.0.jar.sha1 | 1 + .../licenses/opentelemetry-exporter-logging-1.62.0.jar.sha1 | 1 - .../licenses/opentelemetry-exporter-logging-1.63.0.jar.sha1 | 1 + .../licenses/opentelemetry-exporter-otlp-1.62.0.jar.sha1 | 1 - .../licenses/opentelemetry-exporter-otlp-1.63.0.jar.sha1 | 1 + .../licenses/opentelemetry-exporter-otlp-common-1.62.0.jar.sha1 | 1 - .../licenses/opentelemetry-exporter-otlp-common-1.63.0.jar.sha1 | 1 + .../opentelemetry-exporter-sender-okhttp-1.62.0.jar.sha1 | 1 - .../opentelemetry-exporter-sender-okhttp-1.63.0.jar.sha1 | 1 + .../telemetry-otel/licenses/opentelemetry-sdk-1.62.0.jar.sha1 | 1 - .../telemetry-otel/licenses/opentelemetry-sdk-1.63.0.jar.sha1 | 1 + .../licenses/opentelemetry-sdk-common-1.62.0.jar.sha1 | 1 - .../licenses/opentelemetry-sdk-common-1.63.0.jar.sha1 | 1 + .../licenses/opentelemetry-sdk-logs-1.62.0.jar.sha1 | 1 - .../licenses/opentelemetry-sdk-logs-1.63.0.jar.sha1 | 1 + .../licenses/opentelemetry-sdk-metrics-1.62.0.jar.sha1 | 1 - .../licenses/opentelemetry-sdk-metrics-1.63.0.jar.sha1 | 1 + .../licenses/opentelemetry-sdk-trace-1.62.0.jar.sha1 | 1 - .../licenses/opentelemetry-sdk-trace-1.63.0.jar.sha1 | 1 + 29 files changed, 15 insertions(+), 15 deletions(-) delete mode 100644 plugins/telemetry-otel/licenses/opentelemetry-api-1.62.0.jar.sha1 create mode 100644 plugins/telemetry-otel/licenses/opentelemetry-api-1.63.0.jar.sha1 delete mode 100644 plugins/telemetry-otel/licenses/opentelemetry-api-incubator-1.62.0-alpha.jar.sha1 create mode 100644 plugins/telemetry-otel/licenses/opentelemetry-api-incubator-1.63.0-alpha.jar.sha1 delete mode 100644 plugins/telemetry-otel/licenses/opentelemetry-common-1.62.0.jar.sha1 create mode 100644 plugins/telemetry-otel/licenses/opentelemetry-common-1.63.0.jar.sha1 delete mode 100644 plugins/telemetry-otel/licenses/opentelemetry-context-1.62.0.jar.sha1 create mode 100644 plugins/telemetry-otel/licenses/opentelemetry-context-1.63.0.jar.sha1 delete mode 100644 plugins/telemetry-otel/licenses/opentelemetry-exporter-common-1.62.0.jar.sha1 create mode 100644 plugins/telemetry-otel/licenses/opentelemetry-exporter-common-1.63.0.jar.sha1 delete mode 100644 plugins/telemetry-otel/licenses/opentelemetry-exporter-logging-1.62.0.jar.sha1 create mode 100644 plugins/telemetry-otel/licenses/opentelemetry-exporter-logging-1.63.0.jar.sha1 delete mode 100644 plugins/telemetry-otel/licenses/opentelemetry-exporter-otlp-1.62.0.jar.sha1 create mode 100644 plugins/telemetry-otel/licenses/opentelemetry-exporter-otlp-1.63.0.jar.sha1 delete mode 100644 plugins/telemetry-otel/licenses/opentelemetry-exporter-otlp-common-1.62.0.jar.sha1 create mode 100644 plugins/telemetry-otel/licenses/opentelemetry-exporter-otlp-common-1.63.0.jar.sha1 delete mode 100644 plugins/telemetry-otel/licenses/opentelemetry-exporter-sender-okhttp-1.62.0.jar.sha1 create mode 100644 plugins/telemetry-otel/licenses/opentelemetry-exporter-sender-okhttp-1.63.0.jar.sha1 delete mode 100644 plugins/telemetry-otel/licenses/opentelemetry-sdk-1.62.0.jar.sha1 create mode 100644 plugins/telemetry-otel/licenses/opentelemetry-sdk-1.63.0.jar.sha1 delete mode 100644 plugins/telemetry-otel/licenses/opentelemetry-sdk-common-1.62.0.jar.sha1 create mode 100644 plugins/telemetry-otel/licenses/opentelemetry-sdk-common-1.63.0.jar.sha1 delete mode 100644 plugins/telemetry-otel/licenses/opentelemetry-sdk-logs-1.62.0.jar.sha1 create mode 100644 plugins/telemetry-otel/licenses/opentelemetry-sdk-logs-1.63.0.jar.sha1 delete mode 100644 plugins/telemetry-otel/licenses/opentelemetry-sdk-metrics-1.62.0.jar.sha1 create mode 100644 plugins/telemetry-otel/licenses/opentelemetry-sdk-metrics-1.63.0.jar.sha1 delete mode 100644 plugins/telemetry-otel/licenses/opentelemetry-sdk-trace-1.62.0.jar.sha1 create mode 100644 plugins/telemetry-otel/licenses/opentelemetry-sdk-trace-1.63.0.jar.sha1 diff --git a/gradle/libs.versions.toml b/gradle/libs.versions.toml index 483412f03b386..875d0bca23966 100644 --- a/gradle/libs.versions.toml +++ b/gradle/libs.versions.toml @@ -97,7 +97,7 @@ jzlib = "1.1.3" resteasy = "6.2.4.Final" # opentelemetry dependencies -opentelemetry = "1.62.0" +opentelemetry = "1.63.0" opentelemetrysemconv = "1.41.0" # arrow dependencies diff --git a/plugins/telemetry-otel/licenses/opentelemetry-api-1.62.0.jar.sha1 b/plugins/telemetry-otel/licenses/opentelemetry-api-1.62.0.jar.sha1 deleted file mode 100644 index 02ab255e34d5e..0000000000000 --- a/plugins/telemetry-otel/licenses/opentelemetry-api-1.62.0.jar.sha1 +++ /dev/null @@ -1 +0,0 @@ -c4ee83d77005567852a72e08b945ebb023be1daa \ No newline at end of file diff --git a/plugins/telemetry-otel/licenses/opentelemetry-api-1.63.0.jar.sha1 b/plugins/telemetry-otel/licenses/opentelemetry-api-1.63.0.jar.sha1 new file mode 100644 index 0000000000000..0f145e077247a --- /dev/null +++ b/plugins/telemetry-otel/licenses/opentelemetry-api-1.63.0.jar.sha1 @@ -0,0 +1 @@ +39c923c0f236417ec8c4e2091f9e3032b4b6fb91 \ No newline at end of file diff --git a/plugins/telemetry-otel/licenses/opentelemetry-api-incubator-1.62.0-alpha.jar.sha1 b/plugins/telemetry-otel/licenses/opentelemetry-api-incubator-1.62.0-alpha.jar.sha1 deleted file mode 100644 index 88802c8009d0e..0000000000000 --- a/plugins/telemetry-otel/licenses/opentelemetry-api-incubator-1.62.0-alpha.jar.sha1 +++ /dev/null @@ -1 +0,0 @@ -91f3bcf6b93261cbaf32dd156e0007aa5fa5b25a \ No newline at end of file diff --git a/plugins/telemetry-otel/licenses/opentelemetry-api-incubator-1.63.0-alpha.jar.sha1 b/plugins/telemetry-otel/licenses/opentelemetry-api-incubator-1.63.0-alpha.jar.sha1 new file mode 100644 index 0000000000000..aa0966ff936aa --- /dev/null +++ b/plugins/telemetry-otel/licenses/opentelemetry-api-incubator-1.63.0-alpha.jar.sha1 @@ -0,0 +1 @@ +2a88356ec37eb66666dc6798e4221b26af74e24d \ No newline at end of file diff --git a/plugins/telemetry-otel/licenses/opentelemetry-common-1.62.0.jar.sha1 b/plugins/telemetry-otel/licenses/opentelemetry-common-1.62.0.jar.sha1 deleted file mode 100644 index db25f474db864..0000000000000 --- a/plugins/telemetry-otel/licenses/opentelemetry-common-1.62.0.jar.sha1 +++ /dev/null @@ -1 +0,0 @@ -e6468bd64a94429b68761f7c13e143c3fdfaafc7 \ No newline at end of file diff --git a/plugins/telemetry-otel/licenses/opentelemetry-common-1.63.0.jar.sha1 b/plugins/telemetry-otel/licenses/opentelemetry-common-1.63.0.jar.sha1 new file mode 100644 index 0000000000000..5c57ec9f8ba92 --- /dev/null +++ b/plugins/telemetry-otel/licenses/opentelemetry-common-1.63.0.jar.sha1 @@ -0,0 +1 @@ +450f8d552f33d51b19457d1336d8f7bdaec35e01 \ No newline at end of file diff --git a/plugins/telemetry-otel/licenses/opentelemetry-context-1.62.0.jar.sha1 b/plugins/telemetry-otel/licenses/opentelemetry-context-1.62.0.jar.sha1 deleted file mode 100644 index 4608aebd30520..0000000000000 --- a/plugins/telemetry-otel/licenses/opentelemetry-context-1.62.0.jar.sha1 +++ /dev/null @@ -1 +0,0 @@ -365cee4d1f365e4d4a05654742b50aa436c2dd8e \ No newline at end of file diff --git a/plugins/telemetry-otel/licenses/opentelemetry-context-1.63.0.jar.sha1 b/plugins/telemetry-otel/licenses/opentelemetry-context-1.63.0.jar.sha1 new file mode 100644 index 0000000000000..0b6d3c50f42e3 --- /dev/null +++ b/plugins/telemetry-otel/licenses/opentelemetry-context-1.63.0.jar.sha1 @@ -0,0 +1 @@ +4af6d513cedbf78dafe104f82f8dce7c7a205f07 \ No newline at end of file diff --git a/plugins/telemetry-otel/licenses/opentelemetry-exporter-common-1.62.0.jar.sha1 b/plugins/telemetry-otel/licenses/opentelemetry-exporter-common-1.62.0.jar.sha1 deleted file mode 100644 index 9114878e88cef..0000000000000 --- a/plugins/telemetry-otel/licenses/opentelemetry-exporter-common-1.62.0.jar.sha1 +++ /dev/null @@ -1 +0,0 @@ -2dafa6ae65cbf1aa321cd644d200f3ff8465284d \ No newline at end of file diff --git a/plugins/telemetry-otel/licenses/opentelemetry-exporter-common-1.63.0.jar.sha1 b/plugins/telemetry-otel/licenses/opentelemetry-exporter-common-1.63.0.jar.sha1 new file mode 100644 index 0000000000000..20fd2fafeff60 --- /dev/null +++ b/plugins/telemetry-otel/licenses/opentelemetry-exporter-common-1.63.0.jar.sha1 @@ -0,0 +1 @@ +75e2658228eb885770b894f71e8d145d46095fdf \ No newline at end of file diff --git a/plugins/telemetry-otel/licenses/opentelemetry-exporter-logging-1.62.0.jar.sha1 b/plugins/telemetry-otel/licenses/opentelemetry-exporter-logging-1.62.0.jar.sha1 deleted file mode 100644 index 25565231ce2ff..0000000000000 --- a/plugins/telemetry-otel/licenses/opentelemetry-exporter-logging-1.62.0.jar.sha1 +++ /dev/null @@ -1 +0,0 @@ -79ec5f1f23e00da7a8c8a30136cfbfaf9aa38f93 \ No newline at end of file diff --git a/plugins/telemetry-otel/licenses/opentelemetry-exporter-logging-1.63.0.jar.sha1 b/plugins/telemetry-otel/licenses/opentelemetry-exporter-logging-1.63.0.jar.sha1 new file mode 100644 index 0000000000000..aaaef68da4186 --- /dev/null +++ b/plugins/telemetry-otel/licenses/opentelemetry-exporter-logging-1.63.0.jar.sha1 @@ -0,0 +1 @@ +0a5f203ebd35f1ff54171a01335d1a47c3600523 \ No newline at end of file diff --git a/plugins/telemetry-otel/licenses/opentelemetry-exporter-otlp-1.62.0.jar.sha1 b/plugins/telemetry-otel/licenses/opentelemetry-exporter-otlp-1.62.0.jar.sha1 deleted file mode 100644 index 1d0b19032d4ef..0000000000000 --- a/plugins/telemetry-otel/licenses/opentelemetry-exporter-otlp-1.62.0.jar.sha1 +++ /dev/null @@ -1 +0,0 @@ -8e4cb9199ac868332a1213ca27408a18905ba369 \ No newline at end of file diff --git a/plugins/telemetry-otel/licenses/opentelemetry-exporter-otlp-1.63.0.jar.sha1 b/plugins/telemetry-otel/licenses/opentelemetry-exporter-otlp-1.63.0.jar.sha1 new file mode 100644 index 0000000000000..6b5d64d4dfe72 --- /dev/null +++ b/plugins/telemetry-otel/licenses/opentelemetry-exporter-otlp-1.63.0.jar.sha1 @@ -0,0 +1 @@ +7f02be0f4f2c44476e957bce9d7f7d3e0a1a7ae4 \ No newline at end of file diff --git a/plugins/telemetry-otel/licenses/opentelemetry-exporter-otlp-common-1.62.0.jar.sha1 b/plugins/telemetry-otel/licenses/opentelemetry-exporter-otlp-common-1.62.0.jar.sha1 deleted file mode 100644 index 03a22b1f68946..0000000000000 --- a/plugins/telemetry-otel/licenses/opentelemetry-exporter-otlp-common-1.62.0.jar.sha1 +++ /dev/null @@ -1 +0,0 @@ -6f269df0e3f065fbd2e590458e7b2107cde2a106 \ No newline at end of file diff --git a/plugins/telemetry-otel/licenses/opentelemetry-exporter-otlp-common-1.63.0.jar.sha1 b/plugins/telemetry-otel/licenses/opentelemetry-exporter-otlp-common-1.63.0.jar.sha1 new file mode 100644 index 0000000000000..e1adbf694257b --- /dev/null +++ b/plugins/telemetry-otel/licenses/opentelemetry-exporter-otlp-common-1.63.0.jar.sha1 @@ -0,0 +1 @@ +d6a9a75d3e4b9f69a67c3a333162a3b0783178c0 \ No newline at end of file diff --git a/plugins/telemetry-otel/licenses/opentelemetry-exporter-sender-okhttp-1.62.0.jar.sha1 b/plugins/telemetry-otel/licenses/opentelemetry-exporter-sender-okhttp-1.62.0.jar.sha1 deleted file mode 100644 index 82a20ade44ef5..0000000000000 --- a/plugins/telemetry-otel/licenses/opentelemetry-exporter-sender-okhttp-1.62.0.jar.sha1 +++ /dev/null @@ -1 +0,0 @@ -19b5e023db9037a38fe2531afb6e44456e963fba \ No newline at end of file diff --git a/plugins/telemetry-otel/licenses/opentelemetry-exporter-sender-okhttp-1.63.0.jar.sha1 b/plugins/telemetry-otel/licenses/opentelemetry-exporter-sender-okhttp-1.63.0.jar.sha1 new file mode 100644 index 0000000000000..71c90cfe4405d --- /dev/null +++ b/plugins/telemetry-otel/licenses/opentelemetry-exporter-sender-okhttp-1.63.0.jar.sha1 @@ -0,0 +1 @@ +2045d90ed21954e0f61e900d168034e17711c7e9 \ No newline at end of file diff --git a/plugins/telemetry-otel/licenses/opentelemetry-sdk-1.62.0.jar.sha1 b/plugins/telemetry-otel/licenses/opentelemetry-sdk-1.62.0.jar.sha1 deleted file mode 100644 index 99ef9c2e1d260..0000000000000 --- a/plugins/telemetry-otel/licenses/opentelemetry-sdk-1.62.0.jar.sha1 +++ /dev/null @@ -1 +0,0 @@ -06fa52c4641322b14b8bd515eb048bb9b1365d0c \ No newline at end of file diff --git a/plugins/telemetry-otel/licenses/opentelemetry-sdk-1.63.0.jar.sha1 b/plugins/telemetry-otel/licenses/opentelemetry-sdk-1.63.0.jar.sha1 new file mode 100644 index 0000000000000..5cc2625880df0 --- /dev/null +++ b/plugins/telemetry-otel/licenses/opentelemetry-sdk-1.63.0.jar.sha1 @@ -0,0 +1 @@ +1d522dcf3fb4903c7f95835c3f6f631cf1c3824b \ No newline at end of file diff --git a/plugins/telemetry-otel/licenses/opentelemetry-sdk-common-1.62.0.jar.sha1 b/plugins/telemetry-otel/licenses/opentelemetry-sdk-common-1.62.0.jar.sha1 deleted file mode 100644 index 9828f6e0985cf..0000000000000 --- a/plugins/telemetry-otel/licenses/opentelemetry-sdk-common-1.62.0.jar.sha1 +++ /dev/null @@ -1 +0,0 @@ -b6742282daab8e13598b78a83ddfa54f10b5752b \ No newline at end of file diff --git a/plugins/telemetry-otel/licenses/opentelemetry-sdk-common-1.63.0.jar.sha1 b/plugins/telemetry-otel/licenses/opentelemetry-sdk-common-1.63.0.jar.sha1 new file mode 100644 index 0000000000000..c4a8a431030c5 --- /dev/null +++ b/plugins/telemetry-otel/licenses/opentelemetry-sdk-common-1.63.0.jar.sha1 @@ -0,0 +1 @@ +db2ec1f09c307adcf015bb8a87dd687eefa32e77 \ No newline at end of file diff --git a/plugins/telemetry-otel/licenses/opentelemetry-sdk-logs-1.62.0.jar.sha1 b/plugins/telemetry-otel/licenses/opentelemetry-sdk-logs-1.62.0.jar.sha1 deleted file mode 100644 index b008d4cb9a80e..0000000000000 --- a/plugins/telemetry-otel/licenses/opentelemetry-sdk-logs-1.62.0.jar.sha1 +++ /dev/null @@ -1 +0,0 @@ -f242422084100da0bd3a5f6f2bcf364aaf4d2c53 \ No newline at end of file diff --git a/plugins/telemetry-otel/licenses/opentelemetry-sdk-logs-1.63.0.jar.sha1 b/plugins/telemetry-otel/licenses/opentelemetry-sdk-logs-1.63.0.jar.sha1 new file mode 100644 index 0000000000000..dfdc2f6853b85 --- /dev/null +++ b/plugins/telemetry-otel/licenses/opentelemetry-sdk-logs-1.63.0.jar.sha1 @@ -0,0 +1 @@ +ef1e93dc6e05d2b9191d7fb9449fea8965086acb \ No newline at end of file diff --git a/plugins/telemetry-otel/licenses/opentelemetry-sdk-metrics-1.62.0.jar.sha1 b/plugins/telemetry-otel/licenses/opentelemetry-sdk-metrics-1.62.0.jar.sha1 deleted file mode 100644 index a845283b6a3d4..0000000000000 --- a/plugins/telemetry-otel/licenses/opentelemetry-sdk-metrics-1.62.0.jar.sha1 +++ /dev/null @@ -1 +0,0 @@ -5838371075930a4a15f7f61240b4b64cb3e924d8 \ No newline at end of file diff --git a/plugins/telemetry-otel/licenses/opentelemetry-sdk-metrics-1.63.0.jar.sha1 b/plugins/telemetry-otel/licenses/opentelemetry-sdk-metrics-1.63.0.jar.sha1 new file mode 100644 index 0000000000000..2db28d76c8a47 --- /dev/null +++ b/plugins/telemetry-otel/licenses/opentelemetry-sdk-metrics-1.63.0.jar.sha1 @@ -0,0 +1 @@ +decbaf0de6dbe8156aa426ddc85b5b38354a236b \ No newline at end of file diff --git a/plugins/telemetry-otel/licenses/opentelemetry-sdk-trace-1.62.0.jar.sha1 b/plugins/telemetry-otel/licenses/opentelemetry-sdk-trace-1.62.0.jar.sha1 deleted file mode 100644 index 9d3ce157565d3..0000000000000 --- a/plugins/telemetry-otel/licenses/opentelemetry-sdk-trace-1.62.0.jar.sha1 +++ /dev/null @@ -1 +0,0 @@ -7a337d2f887b151d27e734d1c221eb51b1c5b734 \ No newline at end of file diff --git a/plugins/telemetry-otel/licenses/opentelemetry-sdk-trace-1.63.0.jar.sha1 b/plugins/telemetry-otel/licenses/opentelemetry-sdk-trace-1.63.0.jar.sha1 new file mode 100644 index 0000000000000..0ee645084e84b --- /dev/null +++ b/plugins/telemetry-otel/licenses/opentelemetry-sdk-trace-1.63.0.jar.sha1 @@ -0,0 +1 @@ +5c948707623a0dbd79105eb29461f42bfe0db11b \ No newline at end of file From 055d09bd2ad37fb7617edcd51ad5f97b201f97ea Mon Sep 17 00:00:00 2001 From: Simeon Widdis Date: Thu, 11 Jun 2026 14:40:02 -0700 Subject: [PATCH 08/14] fix (ae): add scalars for `IS_{bool}` (#22091) * fix: add scalars for is with bool Signed-off-by: Simeon Widdis * add some test queries Signed-off-by: Simeon Widdis * also add is_true/etc as project functions Signed-off-by: Simeon Widdis --------- Signed-off-by: Simeon Widdis --- .../java/org/opensearch/analytics/spi/ScalarFunction.java | 4 ++++ .../be/datafusion/DataFusionAnalyticsBackendPlugin.java | 8 ++++++++ .../src/test/resources/datasets/functions/ppl/q19.ppl | 1 + .../src/test/resources/datasets/functions/ppl/q20.ppl | 1 + .../src/test/resources/datasets/functions/ppl/q21.ppl | 1 + .../src/test/resources/datasets/functions/ppl/q22.ppl | 1 + 6 files changed, 16 insertions(+) create mode 100644 sandbox/qa/analytics-engine-rest/src/test/resources/datasets/functions/ppl/q19.ppl create mode 100644 sandbox/qa/analytics-engine-rest/src/test/resources/datasets/functions/ppl/q20.ppl create mode 100644 sandbox/qa/analytics-engine-rest/src/test/resources/datasets/functions/ppl/q21.ppl create mode 100644 sandbox/qa/analytics-engine-rest/src/test/resources/datasets/functions/ppl/q22.ppl diff --git a/sandbox/libs/analytics-framework/src/main/java/org/opensearch/analytics/spi/ScalarFunction.java b/sandbox/libs/analytics-framework/src/main/java/org/opensearch/analytics/spi/ScalarFunction.java index 6b6f27440a470..6fbacfa86e8b9 100644 --- a/sandbox/libs/analytics-framework/src/main/java/org/opensearch/analytics/spi/ScalarFunction.java +++ b/sandbox/libs/analytics-framework/src/main/java/org/opensearch/analytics/spi/ScalarFunction.java @@ -38,6 +38,10 @@ public enum ScalarFunction { LESS_THAN_OR_EQUAL(Category.COMPARISON, SqlKind.LESS_THAN_OR_EQUAL), IS_NULL(Category.COMPARISON, SqlKind.IS_NULL), IS_NOT_NULL(Category.COMPARISON, SqlKind.IS_NOT_NULL), + IS_TRUE(Category.COMPARISON, SqlKind.IS_TRUE), + IS_FALSE(Category.COMPARISON, SqlKind.IS_FALSE), + IS_NOT_TRUE(Category.COMPARISON, SqlKind.IS_NOT_TRUE), + IS_NOT_FALSE(Category.COMPARISON, SqlKind.IS_NOT_FALSE), IN(Category.COMPARISON, SqlKind.IN), LIKE(Category.COMPARISON, SqlKind.LIKE), PREFIX(Category.COMPARISON, SqlKind.OTHER_FUNCTION), diff --git a/sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/DataFusionAnalyticsBackendPlugin.java b/sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/DataFusionAnalyticsBackendPlugin.java index b2df66c2526d8..8dff5eb1bb331 100644 --- a/sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/DataFusionAnalyticsBackendPlugin.java +++ b/sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/DataFusionAnalyticsBackendPlugin.java @@ -101,6 +101,10 @@ public class DataFusionAnalyticsBackendPlugin implements AnalyticsSearchBackendP ScalarFunction.LESS_THAN_OR_EQUAL, ScalarFunction.IS_NULL, ScalarFunction.IS_NOT_NULL, + ScalarFunction.IS_TRUE, + ScalarFunction.IS_FALSE, + ScalarFunction.IS_NOT_TRUE, + ScalarFunction.IS_NOT_FALSE, ScalarFunction.IN, ScalarFunction.LIKE, ScalarFunction.REGEXP_CONTAINS, @@ -146,6 +150,10 @@ public class DataFusionAnalyticsBackendPlugin implements AnalyticsSearchBackendP ScalarFunction.IS_NULL, ScalarFunction.IS_NOT_NULL, ScalarFunction.NULLIF, + ScalarFunction.IS_TRUE, + ScalarFunction.IS_FALSE, + ScalarFunction.IS_NOT_TRUE, + ScalarFunction.IS_NOT_FALSE, // ABS / SUBSTRING — PPL sort-pushdown moves these into the project tree; DataFusion has // both natively and isthmus's default catalog binds them, so no adapter needed. ScalarFunction.ABS, diff --git a/sandbox/qa/analytics-engine-rest/src/test/resources/datasets/functions/ppl/q19.ppl b/sandbox/qa/analytics-engine-rest/src/test/resources/datasets/functions/ppl/q19.ppl new file mode 100644 index 0000000000000..6ec9a7a8eb1f0 --- /dev/null +++ b/sandbox/qa/analytics-engine-rest/src/test/resources/datasets/functions/ppl/q19.ppl @@ -0,0 +1 @@ +source=function_tests | where not true = case(amount > 0, true, amount < 0, false else false) | fields user, amount, status | head 10 diff --git a/sandbox/qa/analytics-engine-rest/src/test/resources/datasets/functions/ppl/q20.ppl b/sandbox/qa/analytics-engine-rest/src/test/resources/datasets/functions/ppl/q20.ppl new file mode 100644 index 0000000000000..47089ef801a3e --- /dev/null +++ b/sandbox/qa/analytics-engine-rest/src/test/resources/datasets/functions/ppl/q20.ppl @@ -0,0 +1 @@ +source=function_tests | where true = case(amount > 0, true else false) | fields user, amount | head 10 diff --git a/sandbox/qa/analytics-engine-rest/src/test/resources/datasets/functions/ppl/q21.ppl b/sandbox/qa/analytics-engine-rest/src/test/resources/datasets/functions/ppl/q21.ppl new file mode 100644 index 0000000000000..c3f6843cd849d --- /dev/null +++ b/sandbox/qa/analytics-engine-rest/src/test/resources/datasets/functions/ppl/q21.ppl @@ -0,0 +1 @@ +source=function_tests | where false = case(amount > 0, true else false) | fields user, amount | head 10 diff --git a/sandbox/qa/analytics-engine-rest/src/test/resources/datasets/functions/ppl/q22.ppl b/sandbox/qa/analytics-engine-rest/src/test/resources/datasets/functions/ppl/q22.ppl new file mode 100644 index 0000000000000..dd18eaceb1247 --- /dev/null +++ b/sandbox/qa/analytics-engine-rest/src/test/resources/datasets/functions/ppl/q22.ppl @@ -0,0 +1 @@ +source=function_tests | eval is_positive = case(amount > 0, true else false) | where not false = is_positive | fields user, amount, is_positive | head 10 From b32e6fad12ee9bf38e469c4430859265946ed92c Mon Sep 17 00:00:00 2001 From: Andriy Redko Date: Thu, 11 Jun 2026 19:19:20 -0400 Subject: [PATCH 09/14] Support for HTTP/3 (client side) (#21718) Signed-off-by: Andriy Redko --- client/rest-http-client/build.gradle | 134 ++ .../licenses/bc-fips-2.1.2.jar.sha1 | 1 + .../licenses/bctls-fips-2.1.20.jar.sha1 | 1 + .../licenses/bcutil-fips-2.1.4.jar.sha1 | 1 + .../licenses/bouncycastle-LICENSE.txt | 14 + .../licenses/bouncycastle-NOTICE.txt | 1 + .../licenses/commons-codec-1.18.0.jar.sha1 | 1 + .../licenses/commons-codec-LICENSE.txt | 202 +++ .../licenses/commons-codec-NOTICE.txt | 17 + .../licenses/commons-logging-1.3.5.jar.sha1 | 1 + .../licenses/commons-logging-LICENSE.txt | 202 +++ .../licenses/commons-logging-NOTICE.txt | 6 + .../licenses/reactive-streams-1.0.4.jar.sha1 | 1 + .../licenses/reactive-streams-LICENSE.txt | 21 + .../licenses/reactive-streams-NOTICE.txt | 0 .../licenses/reactor-core-3.8.6.jar.sha1 | 1 + .../licenses/reactor-core-LICENSE.txt | 201 +++ .../licenses/reactor-core-NOTICE.txt | 0 .../licenses/slf4j-api-2.0.17.jar.sha1 | 1 + .../licenses/slf4j-api-LICENSE.txt | 21 + .../licenses/slf4j-api-NOTICE.txt | 0 .../httpclient/AsyncResponseProducer.java | 42 + .../internal/httpclient/BodyUtils.java | 268 ++++ .../internal/httpclient/Cancellable.java | 120 ++ .../httpclient/CompressedResponse.java | 41 + .../internal/httpclient/DeadHostState.java | 125 ++ .../internal/httpclient/HttpHost.java | 52 + .../opensearch/internal/httpclient/Node.java | 278 +++++ .../internal/httpclient/NodeSelector.java | 105 ++ .../httpclient/NonCompressedResponse.java | 32 + .../internal/httpclient/Request.java | 160 +++ .../internal/httpclient/RequestLine.java | 68 + .../internal/httpclient/RequestLogger.java | 190 +++ .../internal/httpclient/RequestOptions.java | 289 +++++ .../internal/httpclient/Response.java | 117 ++ .../httpclient/ResponseException.java | 61 + .../internal/httpclient/ResponseListener.java | 61 + .../httpclient/ResponseWarningsExtractor.java | 96 ++ .../internal/httpclient/RestHttpClient.java | 1096 +++++++++++++++++ .../httpclient/RestHttpClientBuilder.java | 248 ++++ .../internal/httpclient/StatusLine.java | 38 + .../internal/httpclient/StreamingRequest.java | 212 ++++ .../httpclient/StreamingResponse.java | 121 ++ .../httpclient/WarningFailureException.java | 79 ++ .../internal/httpclient/WarningsHandler.java | 80 ++ .../httpclient/BouncyCastleThreadFilter.java | 25 + .../HostsTrackingFailureListener.java | 71 ++ .../HttpClientThreadLeakFilter.java | 24 + .../httpclient/RestClientTestUtil.java | 120 ++ .../RestHttpClientCompressionTests.java | 142 +++ .../RestHttpClientGzipCompressionTests.java | 191 +++ ...RestHttpClientMultipleHostsIntegTests.java | 346 ++++++ .../RestHttpClientMultipleHostsTests.java | 346 ++++++ .../RestHttpClientSingleHostIntegTests.java | 501 ++++++++ ...tpClientSingleHostStreamingIntegTests.java | 180 +++ .../RestHttpClientSingleHostTests.java | 728 +++++++++++ .../httpclient/RestHttpClientTestCase.java | 99 ++ .../RestClientSingleHostIntegTests.java | 12 + settings.gradle | 1 + 59 files changed, 7592 insertions(+) create mode 100644 client/rest-http-client/build.gradle create mode 100644 client/rest-http-client/licenses/bc-fips-2.1.2.jar.sha1 create mode 100644 client/rest-http-client/licenses/bctls-fips-2.1.20.jar.sha1 create mode 100644 client/rest-http-client/licenses/bcutil-fips-2.1.4.jar.sha1 create mode 100644 client/rest-http-client/licenses/bouncycastle-LICENSE.txt create mode 100644 client/rest-http-client/licenses/bouncycastle-NOTICE.txt create mode 100644 client/rest-http-client/licenses/commons-codec-1.18.0.jar.sha1 create mode 100644 client/rest-http-client/licenses/commons-codec-LICENSE.txt create mode 100644 client/rest-http-client/licenses/commons-codec-NOTICE.txt create mode 100644 client/rest-http-client/licenses/commons-logging-1.3.5.jar.sha1 create mode 100644 client/rest-http-client/licenses/commons-logging-LICENSE.txt create mode 100644 client/rest-http-client/licenses/commons-logging-NOTICE.txt create mode 100644 client/rest-http-client/licenses/reactive-streams-1.0.4.jar.sha1 create mode 100644 client/rest-http-client/licenses/reactive-streams-LICENSE.txt create mode 100644 client/rest-http-client/licenses/reactive-streams-NOTICE.txt create mode 100644 client/rest-http-client/licenses/reactor-core-3.8.6.jar.sha1 create mode 100644 client/rest-http-client/licenses/reactor-core-LICENSE.txt create mode 100644 client/rest-http-client/licenses/reactor-core-NOTICE.txt create mode 100644 client/rest-http-client/licenses/slf4j-api-2.0.17.jar.sha1 create mode 100644 client/rest-http-client/licenses/slf4j-api-LICENSE.txt create mode 100644 client/rest-http-client/licenses/slf4j-api-NOTICE.txt create mode 100644 client/rest-http-client/src/main/java/org/opensearch/internal/httpclient/AsyncResponseProducer.java create mode 100644 client/rest-http-client/src/main/java/org/opensearch/internal/httpclient/BodyUtils.java create mode 100644 client/rest-http-client/src/main/java/org/opensearch/internal/httpclient/Cancellable.java create mode 100644 client/rest-http-client/src/main/java/org/opensearch/internal/httpclient/CompressedResponse.java create mode 100644 client/rest-http-client/src/main/java/org/opensearch/internal/httpclient/DeadHostState.java create mode 100644 client/rest-http-client/src/main/java/org/opensearch/internal/httpclient/HttpHost.java create mode 100644 client/rest-http-client/src/main/java/org/opensearch/internal/httpclient/Node.java create mode 100644 client/rest-http-client/src/main/java/org/opensearch/internal/httpclient/NodeSelector.java create mode 100644 client/rest-http-client/src/main/java/org/opensearch/internal/httpclient/NonCompressedResponse.java create mode 100644 client/rest-http-client/src/main/java/org/opensearch/internal/httpclient/Request.java create mode 100644 client/rest-http-client/src/main/java/org/opensearch/internal/httpclient/RequestLine.java create mode 100644 client/rest-http-client/src/main/java/org/opensearch/internal/httpclient/RequestLogger.java create mode 100644 client/rest-http-client/src/main/java/org/opensearch/internal/httpclient/RequestOptions.java create mode 100644 client/rest-http-client/src/main/java/org/opensearch/internal/httpclient/Response.java create mode 100644 client/rest-http-client/src/main/java/org/opensearch/internal/httpclient/ResponseException.java create mode 100644 client/rest-http-client/src/main/java/org/opensearch/internal/httpclient/ResponseListener.java create mode 100644 client/rest-http-client/src/main/java/org/opensearch/internal/httpclient/ResponseWarningsExtractor.java create mode 100644 client/rest-http-client/src/main/java/org/opensearch/internal/httpclient/RestHttpClient.java create mode 100644 client/rest-http-client/src/main/java/org/opensearch/internal/httpclient/RestHttpClientBuilder.java create mode 100644 client/rest-http-client/src/main/java/org/opensearch/internal/httpclient/StatusLine.java create mode 100644 client/rest-http-client/src/main/java/org/opensearch/internal/httpclient/StreamingRequest.java create mode 100644 client/rest-http-client/src/main/java/org/opensearch/internal/httpclient/StreamingResponse.java create mode 100644 client/rest-http-client/src/main/java/org/opensearch/internal/httpclient/WarningFailureException.java create mode 100644 client/rest-http-client/src/main/java/org/opensearch/internal/httpclient/WarningsHandler.java create mode 100644 client/rest-http-client/src/test/java/org/opensearch/internal/httpclient/BouncyCastleThreadFilter.java create mode 100644 client/rest-http-client/src/test/java/org/opensearch/internal/httpclient/HostsTrackingFailureListener.java create mode 100644 client/rest-http-client/src/test/java/org/opensearch/internal/httpclient/HttpClientThreadLeakFilter.java create mode 100644 client/rest-http-client/src/test/java/org/opensearch/internal/httpclient/RestClientTestUtil.java create mode 100644 client/rest-http-client/src/test/java/org/opensearch/internal/httpclient/RestHttpClientCompressionTests.java create mode 100644 client/rest-http-client/src/test/java/org/opensearch/internal/httpclient/RestHttpClientGzipCompressionTests.java create mode 100644 client/rest-http-client/src/test/java/org/opensearch/internal/httpclient/RestHttpClientMultipleHostsIntegTests.java create mode 100644 client/rest-http-client/src/test/java/org/opensearch/internal/httpclient/RestHttpClientMultipleHostsTests.java create mode 100644 client/rest-http-client/src/test/java/org/opensearch/internal/httpclient/RestHttpClientSingleHostIntegTests.java create mode 100644 client/rest-http-client/src/test/java/org/opensearch/internal/httpclient/RestHttpClientSingleHostStreamingIntegTests.java create mode 100644 client/rest-http-client/src/test/java/org/opensearch/internal/httpclient/RestHttpClientSingleHostTests.java create mode 100644 client/rest-http-client/src/test/java/org/opensearch/internal/httpclient/RestHttpClientTestCase.java diff --git a/client/rest-http-client/build.gradle b/client/rest-http-client/build.gradle new file mode 100644 index 0000000000000..69a6813f0329e --- /dev/null +++ b/client/rest-http-client/build.gradle @@ -0,0 +1,134 @@ +/* + * SPDX-License-Identifier: Apache-2.0 + * + * The OpenSearch Contributors require contributions made to + * this file be licensed under the Apache-2.0 license or a + * compatible open source license. + * + * Modifications Copyright OpenSearch Contributors. See + * GitHub history for details. + */ + +import de.thetaphi.forbiddenapis.gradle.CheckForbiddenApis + +apply plugin: 'opensearch.build' +apply plugin: 'opensearch.publish' +apply from: "$rootDir/gradle/fips.gradle" + +java { + targetCompatibility = JavaVersion.VERSION_21 + sourceCompatibility = JavaVersion.VERSION_21 +} + +base { + group = 'org.opensearch.client' + archivesName = 'opensearch-rest-http-client' +} + +dependencies { + api "commons-codec:commons-codec:${versions.commonscodec}" + api "commons-logging:commons-logging:${versions.commonslogging}" + api "org.slf4j:slf4j-api:${versions.slf4j}" + fipsRuntimeOnly "org.bouncycastle:bc-fips:${versions.bouncycastle_jce}" + fipsRuntimeOnly "org.bouncycastle:bctls-fips:${versions.bouncycastle_tls}" + fipsRuntimeOnly "org.bouncycastle:bcutil-fips:${versions.bouncycastle_util}" + + // reactor + api "io.projectreactor:reactor-core:${versions.reactor}" + api "org.reactivestreams:reactive-streams:${versions.reactivestreams}" + + testImplementation "com.carrotsearch.randomizedtesting:randomizedtesting-runner:${versions.randomizedrunner}" + testImplementation "junit:junit:${versions.junit}" + testImplementation "org.hamcrest:hamcrest:${versions.hamcrest}" + testImplementation "org.mockito:mockito-core:${versions.mockito}" + testImplementation "org.objenesis:objenesis:${versions.objenesis}" + testImplementation "net.bytebuddy:byte-buddy:${versions.bytebuddy}" + testImplementation "net.bytebuddy:byte-buddy-agent:${versions.bytebuddy}" + testImplementation "org.apache.logging.log4j:log4j-api:${versions.log4j}" + testImplementation "org.apache.logging.log4j:log4j-core:${versions.log4j}" + testImplementation "org.apache.logging.log4j:log4j-jul:${versions.log4j}" + testImplementation "org.apache.logging.log4j:log4j-slf4j2-impl:${versions.log4j}" + testImplementation "io.projectreactor:reactor-test:${versions.reactor}" +} + +tasks.named("dependencyLicenses").configure { + mapping from: /bc.*/, to: 'bouncycastle' +} + +tasks.withType(CheckForbiddenApis).configureEach { + //client does not depend on server, so only jdk and http signatures should be checked + replaceSignatureFiles('jdk-signatures') +} + +tasks.named('forbiddenApisTest').configure { + //we are using jdk-internal instead of jdk-non-portable to allow for com.sun.net.httpserver.* usage + bundledSignatures -= 'jdk-non-portable' + bundledSignatures += 'jdk-internal' +} + +testingConventions { + naming.clear() + naming { + Tests { + baseClass 'org.opensearch.internal.httpclient.RestHttpClientTestCase' + } + } +} + +thirdPartyAudit { + ignoreMissingClasses( + //commons-logging optional dependencies + 'org.apache.avalon.framework.logger.Logger', + 'org.apache.log.Hierarchy', + 'org.apache.log.Logger', + 'org.apache.log4j.Level', + 'org.apache.log4j.Logger', + 'org.apache.log4j.Priority', + 'org.apache.logging.log4j.Level', + 'org.apache.logging.log4j.LogManager', + 'org.apache.logging.log4j.Marker', + 'org.apache.logging.log4j.MarkerManager', + 'org.apache.logging.log4j.spi.AbstractLoggerAdapter', + 'org.apache.logging.log4j.spi.ExtendedLogger', + 'org.apache.logging.log4j.spi.LoggerAdapter', + 'org.apache.logging.log4j.spi.LoggerContext', + 'org.apache.logging.log4j.spi.LoggerContextFactory', + 'org.apache.logging.log4j.util.StackLocatorUtil', + //commons-logging provided dependencies + 'javax.servlet.ServletContextEvent', + 'javax.servlet.ServletContextListener', + 'io.micrometer.context.ContextAccessor', + 'io.micrometer.context.ContextRegistry', + 'io.micrometer.context.ContextSnapshot', + 'io.micrometer.context.ContextSnapshot$Scope', + 'io.micrometer.context.ContextSnapshotFactory', + 'io.micrometer.context.ContextSnapshotFactory$Builder', + 'io.micrometer.context.ThreadLocalAccessor', + 'io.micrometer.core.instrument.Clock', + 'io.micrometer.core.instrument.Counter', + 'io.micrometer.core.instrument.Counter$Builder', + 'io.micrometer.core.instrument.DistributionSummary', + 'io.micrometer.core.instrument.DistributionSummary$Builder', + 'io.micrometer.core.instrument.Meter', + 'io.micrometer.core.instrument.MeterRegistry', + 'io.micrometer.core.instrument.Metrics', + 'io.micrometer.core.instrument.Tag', + 'io.micrometer.core.instrument.Tags', + 'io.micrometer.core.instrument.Timer', + 'io.micrometer.core.instrument.Timer$Builder', + 'io.micrometer.core.instrument.Timer$Sample', + 'io.micrometer.core.instrument.binder.jvm.ExecutorServiceMetrics', + 'io.micrometer.core.instrument.composite.CompositeMeterRegistry', + 'io.micrometer.core.instrument.search.Search', + 'reactor.blockhound.BlockHound$Builder', + 'reactor.blockhound.integration.BlockHoundIntegration' + ) + ignoreViolations( + 'reactor.core.publisher.CallSiteSupplierFactory$SharedSecretsCallSiteSupplierFactory', + 'reactor.core.publisher.CallSiteSupplierFactory$SharedSecretsCallSiteSupplierFactory$TracingException' + ) +} + +tasks.named("missingJavadoc").configure { + it.enabled = false +} diff --git a/client/rest-http-client/licenses/bc-fips-2.1.2.jar.sha1 b/client/rest-http-client/licenses/bc-fips-2.1.2.jar.sha1 new file mode 100644 index 0000000000000..bd2f333cb12d0 --- /dev/null +++ b/client/rest-http-client/licenses/bc-fips-2.1.2.jar.sha1 @@ -0,0 +1 @@ +061fbe8383f70489dda95a11a2a4739eb818ff2c \ No newline at end of file diff --git a/client/rest-http-client/licenses/bctls-fips-2.1.20.jar.sha1 b/client/rest-http-client/licenses/bctls-fips-2.1.20.jar.sha1 new file mode 100644 index 0000000000000..7266ec5abf10a --- /dev/null +++ b/client/rest-http-client/licenses/bctls-fips-2.1.20.jar.sha1 @@ -0,0 +1 @@ +9c0632a6c5ca09a86434cf5e02e72c221e1c930f \ No newline at end of file diff --git a/client/rest-http-client/licenses/bcutil-fips-2.1.4.jar.sha1 b/client/rest-http-client/licenses/bcutil-fips-2.1.4.jar.sha1 new file mode 100644 index 0000000000000..73b19722430fb --- /dev/null +++ b/client/rest-http-client/licenses/bcutil-fips-2.1.4.jar.sha1 @@ -0,0 +1 @@ +1d37b7a28560684f5b8e4fd65478c9130d4015d0 \ No newline at end of file diff --git a/client/rest-http-client/licenses/bouncycastle-LICENSE.txt b/client/rest-http-client/licenses/bouncycastle-LICENSE.txt new file mode 100644 index 0000000000000..5c7c14696849d --- /dev/null +++ b/client/rest-http-client/licenses/bouncycastle-LICENSE.txt @@ -0,0 +1,14 @@ +Copyright (c) 2000 - 2023 The Legion of the Bouncy Castle Inc. (https://www.bouncycastle.org) + +Permission is hereby granted, free of charge, to any person obtaining a copy of this software and associated +documentation files (the "Software"), to deal in the Software without restriction, including without limitation +the rights to use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies of the Software, +and to permit persons to whom the Software is furnished to do so, subject to the following conditions: + +The above copyright notice and this permission notice shall be included in all copies or substantial portions of the +Software. + +THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE +WARRANTIES OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR +COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR +OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE. diff --git a/client/rest-http-client/licenses/bouncycastle-NOTICE.txt b/client/rest-http-client/licenses/bouncycastle-NOTICE.txt new file mode 100644 index 0000000000000..8b137891791fe --- /dev/null +++ b/client/rest-http-client/licenses/bouncycastle-NOTICE.txt @@ -0,0 +1 @@ + diff --git a/client/rest-http-client/licenses/commons-codec-1.18.0.jar.sha1 b/client/rest-http-client/licenses/commons-codec-1.18.0.jar.sha1 new file mode 100644 index 0000000000000..01a6a8f446302 --- /dev/null +++ b/client/rest-http-client/licenses/commons-codec-1.18.0.jar.sha1 @@ -0,0 +1 @@ +ee45d1cf6ec2cc2b809ff04b4dc7aec858e0df8f \ No newline at end of file diff --git a/client/rest-http-client/licenses/commons-codec-LICENSE.txt b/client/rest-http-client/licenses/commons-codec-LICENSE.txt new file mode 100644 index 0000000000000..d645695673349 --- /dev/null +++ b/client/rest-http-client/licenses/commons-codec-LICENSE.txt @@ -0,0 +1,202 @@ + + Apache License + Version 2.0, January 2004 + http://www.apache.org/licenses/ + + TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION + + 1. Definitions. + + "License" shall mean the terms and conditions for use, reproduction, + and distribution as defined by Sections 1 through 9 of this document. + + "Licensor" shall mean the copyright owner or entity authorized by + the copyright owner that is granting the License. + + "Legal Entity" shall mean the union of the acting entity and all + other entities that control, are controlled by, or are under common + control with that entity. For the purposes of this definition, + "control" means (i) the power, direct or indirect, to cause the + direction or management of such entity, whether by contract or + otherwise, or (ii) ownership of fifty percent (50%) or more of the + outstanding shares, or (iii) beneficial ownership of such entity. + + "You" (or "Your") shall mean an individual or Legal Entity + exercising permissions granted by this License. + + "Source" form shall mean the preferred form for making modifications, + including but not limited to software source code, documentation + source, and configuration files. + + "Object" form shall mean any form resulting from mechanical + transformation or translation of a Source form, including but + not limited to compiled object code, generated documentation, + and conversions to other media types. + + "Work" shall mean the work of authorship, whether in Source or + Object form, made available under the License, as indicated by a + copyright notice that is included in or attached to the work + (an example is provided in the Appendix below). + + "Derivative Works" shall mean any work, whether in Source or Object + form, that is based on (or derived from) the Work and for which the + editorial revisions, annotations, elaborations, or other modifications + represent, as a whole, an original work of authorship. For the purposes + of this License, Derivative Works shall not include works that remain + separable from, or merely link (or bind by name) to the interfaces of, + the Work and Derivative Works thereof. + + "Contribution" shall mean any work of authorship, including + the original version of the Work and any modifications or additions + to that Work or Derivative Works thereof, that is intentionally + submitted to Licensor for inclusion in the Work by the copyright owner + or by an individual or Legal Entity authorized to submit on behalf of + the copyright owner. For the purposes of this definition, "submitted" + means any form of electronic, verbal, or written communication sent + to the Licensor or its representatives, including but not limited to + communication on electronic mailing lists, source code control systems, + and issue tracking systems that are managed by, or on behalf of, the + Licensor for the purpose of discussing and improving the Work, but + excluding communication that is conspicuously marked or otherwise + designated in writing by the copyright owner as "Not a Contribution." + + "Contributor" shall mean Licensor and any individual or Legal Entity + on behalf of whom a Contribution has been received by Licensor and + subsequently incorporated within the Work. + + 2. Grant of Copyright License. Subject to the terms and conditions of + this License, each Contributor hereby grants to You a perpetual, + worldwide, non-exclusive, no-charge, royalty-free, irrevocable + copyright license to reproduce, prepare Derivative Works of, + publicly display, publicly perform, sublicense, and distribute the + Work and such Derivative Works in Source or Object form. + + 3. Grant of Patent License. Subject to the terms and conditions of + this License, each Contributor hereby grants to You a perpetual, + worldwide, non-exclusive, no-charge, royalty-free, irrevocable + (except as stated in this section) patent license to make, have made, + use, offer to sell, sell, import, and otherwise transfer the Work, + where such license applies only to those patent claims licensable + by such Contributor that are necessarily infringed by their + Contribution(s) alone or by combination of their Contribution(s) + with the Work to which such Contribution(s) was submitted. If You + institute patent litigation against any entity (including a + cross-claim or counterclaim in a lawsuit) alleging that the Work + or a Contribution incorporated within the Work constitutes direct + or contributory patent infringement, then any patent licenses + granted to You under this License for that Work shall terminate + as of the date such litigation is filed. + + 4. Redistribution. You may reproduce and distribute copies of the + Work or Derivative Works thereof in any medium, with or without + modifications, and in Source or Object form, provided that You + meet the following conditions: + + (a) You must give any other recipients of the Work or + Derivative Works a copy of this License; and + + (b) You must cause any modified files to carry prominent notices + stating that You changed the files; and + + (c) You must retain, in the Source form of any Derivative Works + that You distribute, all copyright, patent, trademark, and + attribution notices from the Source form of the Work, + excluding those notices that do not pertain to any part of + the Derivative Works; and + + (d) If the Work includes a "NOTICE" text file as part of its + distribution, then any Derivative Works that You distribute must + include a readable copy of the attribution notices contained + within such NOTICE file, excluding those notices that do not + pertain to any part of the Derivative Works, in at least one + of the following places: within a NOTICE text file distributed + as part of the Derivative Works; within the Source form or + documentation, if provided along with the Derivative Works; or, + within a display generated by the Derivative Works, if and + wherever such third-party notices normally appear. The contents + of the NOTICE file are for informational purposes only and + do not modify the License. You may add Your own attribution + notices within Derivative Works that You distribute, alongside + or as an addendum to the NOTICE text from the Work, provided + that such additional attribution notices cannot be construed + as modifying the License. + + You may add Your own copyright statement to Your modifications and + may provide additional or different license terms and conditions + for use, reproduction, or distribution of Your modifications, or + for any such Derivative Works as a whole, provided Your use, + reproduction, and distribution of the Work otherwise complies with + the conditions stated in this License. + + 5. Submission of Contributions. Unless You explicitly state otherwise, + any Contribution intentionally submitted for inclusion in the Work + by You to the Licensor shall be under the terms and conditions of + this License, without any additional terms or conditions. + Notwithstanding the above, nothing herein shall supersede or modify + the terms of any separate license agreement you may have executed + with Licensor regarding such Contributions. + + 6. Trademarks. This License does not grant permission to use the trade + names, trademarks, service marks, or product names of the Licensor, + except as required for reasonable and customary use in describing the + origin of the Work and reproducing the content of the NOTICE file. + + 7. Disclaimer of Warranty. Unless required by applicable law or + agreed to in writing, Licensor provides the Work (and each + Contributor provides its Contributions) on an "AS IS" BASIS, + WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or + implied, including, without limitation, any warranties or conditions + of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A + PARTICULAR PURPOSE. You are solely responsible for determining the + appropriateness of using or redistributing the Work and assume any + risks associated with Your exercise of permissions under this License. + + 8. Limitation of Liability. In no event and under no legal theory, + whether in tort (including negligence), contract, or otherwise, + unless required by applicable law (such as deliberate and grossly + negligent acts) or agreed to in writing, shall any Contributor be + liable to You for damages, including any direct, indirect, special, + incidental, or consequential damages of any character arising as a + result of this License or out of the use or inability to use the + Work (including but not limited to damages for loss of goodwill, + work stoppage, computer failure or malfunction, or any and all + other commercial damages or losses), even if such Contributor + has been advised of the possibility of such damages. + + 9. Accepting Warranty or Additional Liability. While redistributing + the Work or Derivative Works thereof, You may choose to offer, + and charge a fee for, acceptance of support, warranty, indemnity, + or other liability obligations and/or rights consistent with this + License. However, in accepting such obligations, You may act only + on Your own behalf and on Your sole responsibility, not on behalf + of any other Contributor, and only if You agree to indemnify, + defend, and hold each Contributor harmless for any liability + incurred by, or claims asserted against, such Contributor by reason + of your accepting any such warranty or additional liability. + + END OF TERMS AND CONDITIONS + + APPENDIX: How to apply the Apache License to your work. + + To apply the Apache License to your work, attach the following + boilerplate notice, with the fields enclosed by brackets "[]" + replaced with your own identifying information. (Don't include + the brackets!) The text should be enclosed in the appropriate + comment syntax for the file format. We also recommend that a + file or class name and description of purpose be included on the + same "printed page" as the copyright notice for easier + identification within third-party archives. + + Copyright [yyyy] [name of copyright owner] + + Licensed 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. diff --git a/client/rest-http-client/licenses/commons-codec-NOTICE.txt b/client/rest-http-client/licenses/commons-codec-NOTICE.txt new file mode 100644 index 0000000000000..1da9af50f6008 --- /dev/null +++ b/client/rest-http-client/licenses/commons-codec-NOTICE.txt @@ -0,0 +1,17 @@ +Apache Commons Codec +Copyright 2002-2014 The Apache Software Foundation + +This product includes software developed at +The Apache Software Foundation (http://www.apache.org/). + +src/test/org/apache/commons/codec/language/DoubleMetaphoneTest.java +contains test data from http://aspell.net/test/orig/batch0.tab. +Copyright (C) 2002 Kevin Atkinson (kevina@gnu.org) + +=============================================================================== + +The content of package org.apache.commons.codec.language.bm has been translated +from the original php source code available at http://stevemorse.org/phoneticinfo.htm +with permission from the original authors. +Original source copyright: +Copyright (c) 2008 Alexander Beider & Stephen P. Morse. diff --git a/client/rest-http-client/licenses/commons-logging-1.3.5.jar.sha1 b/client/rest-http-client/licenses/commons-logging-1.3.5.jar.sha1 new file mode 100644 index 0000000000000..f7ddad61aaeaa --- /dev/null +++ b/client/rest-http-client/licenses/commons-logging-1.3.5.jar.sha1 @@ -0,0 +1 @@ +a3fcc5d3c29b2b03433aa2d2f2d2c1b1638924a1 \ No newline at end of file diff --git a/client/rest-http-client/licenses/commons-logging-LICENSE.txt b/client/rest-http-client/licenses/commons-logging-LICENSE.txt new file mode 100644 index 0000000000000..d645695673349 --- /dev/null +++ b/client/rest-http-client/licenses/commons-logging-LICENSE.txt @@ -0,0 +1,202 @@ + + Apache License + Version 2.0, January 2004 + http://www.apache.org/licenses/ + + TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION + + 1. Definitions. + + "License" shall mean the terms and conditions for use, reproduction, + and distribution as defined by Sections 1 through 9 of this document. + + "Licensor" shall mean the copyright owner or entity authorized by + the copyright owner that is granting the License. + + "Legal Entity" shall mean the union of the acting entity and all + other entities that control, are controlled by, or are under common + control with that entity. For the purposes of this definition, + "control" means (i) the power, direct or indirect, to cause the + direction or management of such entity, whether by contract or + otherwise, or (ii) ownership of fifty percent (50%) or more of the + outstanding shares, or (iii) beneficial ownership of such entity. + + "You" (or "Your") shall mean an individual or Legal Entity + exercising permissions granted by this License. + + "Source" form shall mean the preferred form for making modifications, + including but not limited to software source code, documentation + source, and configuration files. + + "Object" form shall mean any form resulting from mechanical + transformation or translation of a Source form, including but + not limited to compiled object code, generated documentation, + and conversions to other media types. + + "Work" shall mean the work of authorship, whether in Source or + Object form, made available under the License, as indicated by a + copyright notice that is included in or attached to the work + (an example is provided in the Appendix below). + + "Derivative Works" shall mean any work, whether in Source or Object + form, that is based on (or derived from) the Work and for which the + editorial revisions, annotations, elaborations, or other modifications + represent, as a whole, an original work of authorship. For the purposes + of this License, Derivative Works shall not include works that remain + separable from, or merely link (or bind by name) to the interfaces of, + the Work and Derivative Works thereof. + + "Contribution" shall mean any work of authorship, including + the original version of the Work and any modifications or additions + to that Work or Derivative Works thereof, that is intentionally + submitted to Licensor for inclusion in the Work by the copyright owner + or by an individual or Legal Entity authorized to submit on behalf of + the copyright owner. For the purposes of this definition, "submitted" + means any form of electronic, verbal, or written communication sent + to the Licensor or its representatives, including but not limited to + communication on electronic mailing lists, source code control systems, + and issue tracking systems that are managed by, or on behalf of, the + Licensor for the purpose of discussing and improving the Work, but + excluding communication that is conspicuously marked or otherwise + designated in writing by the copyright owner as "Not a Contribution." + + "Contributor" shall mean Licensor and any individual or Legal Entity + on behalf of whom a Contribution has been received by Licensor and + subsequently incorporated within the Work. + + 2. Grant of Copyright License. Subject to the terms and conditions of + this License, each Contributor hereby grants to You a perpetual, + worldwide, non-exclusive, no-charge, royalty-free, irrevocable + copyright license to reproduce, prepare Derivative Works of, + publicly display, publicly perform, sublicense, and distribute the + Work and such Derivative Works in Source or Object form. + + 3. Grant of Patent License. Subject to the terms and conditions of + this License, each Contributor hereby grants to You a perpetual, + worldwide, non-exclusive, no-charge, royalty-free, irrevocable + (except as stated in this section) patent license to make, have made, + use, offer to sell, sell, import, and otherwise transfer the Work, + where such license applies only to those patent claims licensable + by such Contributor that are necessarily infringed by their + Contribution(s) alone or by combination of their Contribution(s) + with the Work to which such Contribution(s) was submitted. If You + institute patent litigation against any entity (including a + cross-claim or counterclaim in a lawsuit) alleging that the Work + or a Contribution incorporated within the Work constitutes direct + or contributory patent infringement, then any patent licenses + granted to You under this License for that Work shall terminate + as of the date such litigation is filed. + + 4. Redistribution. You may reproduce and distribute copies of the + Work or Derivative Works thereof in any medium, with or without + modifications, and in Source or Object form, provided that You + meet the following conditions: + + (a) You must give any other recipients of the Work or + Derivative Works a copy of this License; and + + (b) You must cause any modified files to carry prominent notices + stating that You changed the files; and + + (c) You must retain, in the Source form of any Derivative Works + that You distribute, all copyright, patent, trademark, and + attribution notices from the Source form of the Work, + excluding those notices that do not pertain to any part of + the Derivative Works; and + + (d) If the Work includes a "NOTICE" text file as part of its + distribution, then any Derivative Works that You distribute must + include a readable copy of the attribution notices contained + within such NOTICE file, excluding those notices that do not + pertain to any part of the Derivative Works, in at least one + of the following places: within a NOTICE text file distributed + as part of the Derivative Works; within the Source form or + documentation, if provided along with the Derivative Works; or, + within a display generated by the Derivative Works, if and + wherever such third-party notices normally appear. The contents + of the NOTICE file are for informational purposes only and + do not modify the License. You may add Your own attribution + notices within Derivative Works that You distribute, alongside + or as an addendum to the NOTICE text from the Work, provided + that such additional attribution notices cannot be construed + as modifying the License. + + You may add Your own copyright statement to Your modifications and + may provide additional or different license terms and conditions + for use, reproduction, or distribution of Your modifications, or + for any such Derivative Works as a whole, provided Your use, + reproduction, and distribution of the Work otherwise complies with + the conditions stated in this License. + + 5. Submission of Contributions. Unless You explicitly state otherwise, + any Contribution intentionally submitted for inclusion in the Work + by You to the Licensor shall be under the terms and conditions of + this License, without any additional terms or conditions. + Notwithstanding the above, nothing herein shall supersede or modify + the terms of any separate license agreement you may have executed + with Licensor regarding such Contributions. + + 6. Trademarks. This License does not grant permission to use the trade + names, trademarks, service marks, or product names of the Licensor, + except as required for reasonable and customary use in describing the + origin of the Work and reproducing the content of the NOTICE file. + + 7. Disclaimer of Warranty. Unless required by applicable law or + agreed to in writing, Licensor provides the Work (and each + Contributor provides its Contributions) on an "AS IS" BASIS, + WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or + implied, including, without limitation, any warranties or conditions + of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A + PARTICULAR PURPOSE. You are solely responsible for determining the + appropriateness of using or redistributing the Work and assume any + risks associated with Your exercise of permissions under this License. + + 8. Limitation of Liability. In no event and under no legal theory, + whether in tort (including negligence), contract, or otherwise, + unless required by applicable law (such as deliberate and grossly + negligent acts) or agreed to in writing, shall any Contributor be + liable to You for damages, including any direct, indirect, special, + incidental, or consequential damages of any character arising as a + result of this License or out of the use or inability to use the + Work (including but not limited to damages for loss of goodwill, + work stoppage, computer failure or malfunction, or any and all + other commercial damages or losses), even if such Contributor + has been advised of the possibility of such damages. + + 9. Accepting Warranty or Additional Liability. While redistributing + the Work or Derivative Works thereof, You may choose to offer, + and charge a fee for, acceptance of support, warranty, indemnity, + or other liability obligations and/or rights consistent with this + License. However, in accepting such obligations, You may act only + on Your own behalf and on Your sole responsibility, not on behalf + of any other Contributor, and only if You agree to indemnify, + defend, and hold each Contributor harmless for any liability + incurred by, or claims asserted against, such Contributor by reason + of your accepting any such warranty or additional liability. + + END OF TERMS AND CONDITIONS + + APPENDIX: How to apply the Apache License to your work. + + To apply the Apache License to your work, attach the following + boilerplate notice, with the fields enclosed by brackets "[]" + replaced with your own identifying information. (Don't include + the brackets!) The text should be enclosed in the appropriate + comment syntax for the file format. We also recommend that a + file or class name and description of purpose be included on the + same "printed page" as the copyright notice for easier + identification within third-party archives. + + Copyright [yyyy] [name of copyright owner] + + Licensed 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. diff --git a/client/rest-http-client/licenses/commons-logging-NOTICE.txt b/client/rest-http-client/licenses/commons-logging-NOTICE.txt new file mode 100644 index 0000000000000..a37977d45a168 --- /dev/null +++ b/client/rest-http-client/licenses/commons-logging-NOTICE.txt @@ -0,0 +1,6 @@ +Apache Commons Logging +Copyright 2003-2013 The Apache Software Foundation + +This product includes software developed at +The Apache Software Foundation (http://www.apache.org/). + diff --git a/client/rest-http-client/licenses/reactive-streams-1.0.4.jar.sha1 b/client/rest-http-client/licenses/reactive-streams-1.0.4.jar.sha1 new file mode 100644 index 0000000000000..45a80e3f7e361 --- /dev/null +++ b/client/rest-http-client/licenses/reactive-streams-1.0.4.jar.sha1 @@ -0,0 +1 @@ +3864a1320d97d7b045f729a326e1e077661f31b7 \ No newline at end of file diff --git a/client/rest-http-client/licenses/reactive-streams-LICENSE.txt b/client/rest-http-client/licenses/reactive-streams-LICENSE.txt new file mode 100644 index 0000000000000..1e3c7e7c77495 --- /dev/null +++ b/client/rest-http-client/licenses/reactive-streams-LICENSE.txt @@ -0,0 +1,21 @@ +MIT No Attribution + +Copyright 2014 Reactive Streams + +Permission is hereby granted, free of charge, to any person obtaining a copy +of this software and associated documentation files (the "Software"), to deal +in the Software without restriction, including without limitation the rights +to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +copies of the Software, and to permit persons to whom the Software is +furnished to do so, subject to the following conditions: + +The above copyright notice and this permission notice shall be included in all +copies or substantial portions of the Software. + +THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +SOFTWARE. \ No newline at end of file diff --git a/client/rest-http-client/licenses/reactive-streams-NOTICE.txt b/client/rest-http-client/licenses/reactive-streams-NOTICE.txt new file mode 100644 index 0000000000000..e69de29bb2d1d diff --git a/client/rest-http-client/licenses/reactor-core-3.8.6.jar.sha1 b/client/rest-http-client/licenses/reactor-core-3.8.6.jar.sha1 new file mode 100644 index 0000000000000..6cdefa9580bdb --- /dev/null +++ b/client/rest-http-client/licenses/reactor-core-3.8.6.jar.sha1 @@ -0,0 +1 @@ +76285d63d5da4ed8679357628dd5309b0feb77ff \ No newline at end of file diff --git a/client/rest-http-client/licenses/reactor-core-LICENSE.txt b/client/rest-http-client/licenses/reactor-core-LICENSE.txt new file mode 100644 index 0000000000000..e5583c184e67a --- /dev/null +++ b/client/rest-http-client/licenses/reactor-core-LICENSE.txt @@ -0,0 +1,201 @@ +Apache License + Version 2.0, January 2004 + https://www.apache.org/licenses/ + + TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION + + 1. Definitions. + + "License" shall mean the terms and conditions for use, reproduction, + and distribution as defined by Sections 1 through 9 of this document. + + "Licensor" shall mean the copyright owner or entity authorized by + the copyright owner that is granting the License. + + "Legal Entity" shall mean the union of the acting entity and all + other entities that control, are controlled by, or are under common + control with that entity. For the purposes of this definition, + "control" means (i) the power, direct or indirect, to cause the + direction or management of such entity, whether by contract or + otherwise, or (ii) ownership of fifty percent (50%) or more of the + outstanding shares, or (iii) beneficial ownership of such entity. + + "You" (or "Your") shall mean an individual or Legal Entity + exercising permissions granted by this License. + + "Source" form shall mean the preferred form for making modifications, + including but not limited to software source code, documentation + source, and configuration files. + + "Object" form shall mean any form resulting from mechanical + transformation or translation of a Source form, including but + not limited to compiled object code, generated documentation, + and conversions to other media types. + + "Work" shall mean the work of authorship, whether in Source or + Object form, made available under the License, as indicated by a + copyright notice that is included in or attached to the work + (an example is provided in the Appendix below). + + "Derivative Works" shall mean any work, whether in Source or Object + form, that is based on (or derived from) the Work and for which the + editorial revisions, annotations, elaborations, or other modifications + represent, as a whole, an original work of authorship. For the purposes + of this License, Derivative Works shall not include works that remain + separable from, or merely link (or bind by name) to the interfaces of, + the Work and Derivative Works thereof. + + "Contribution" shall mean any work of authorship, including + the original version of the Work and any modifications or additions + to that Work or Derivative Works thereof, that is intentionally + submitted to Licensor for inclusion in the Work by the copyright owner + or by an individual or Legal Entity authorized to submit on behalf of + the copyright owner. For the purposes of this definition, "submitted" + means any form of electronic, verbal, or written communication sent + to the Licensor or its representatives, including but not limited to + communication on electronic mailing lists, source code control systems, + and issue tracking systems that are managed by, or on behalf of, the + Licensor for the purpose of discussing and improving the Work, but + excluding communication that is conspicuously marked or otherwise + designated in writing by the copyright owner as "Not a Contribution." + + "Contributor" shall mean Licensor and any individual or Legal Entity + on behalf of whom a Contribution has been received by Licensor and + subsequently incorporated within the Work. + + 2. Grant of Copyright License. Subject to the terms and conditions of + this License, each Contributor hereby grants to You a perpetual, + worldwide, non-exclusive, no-charge, royalty-free, irrevocable + copyright license to reproduce, prepare Derivative Works of, + publicly display, publicly perform, sublicense, and distribute the + Work and such Derivative Works in Source or Object form. + + 3. Grant of Patent License. Subject to the terms and conditions of + this License, each Contributor hereby grants to You a perpetual, + worldwide, non-exclusive, no-charge, royalty-free, irrevocable + (except as stated in this section) patent license to make, have made, + use, offer to sell, sell, import, and otherwise transfer the Work, + where such license applies only to those patent claims licensable + by such Contributor that are necessarily infringed by their + Contribution(s) alone or by combination of their Contribution(s) + with the Work to which such Contribution(s) was submitted. If You + institute patent litigation against any entity (including a + cross-claim or counterclaim in a lawsuit) alleging that the Work + or a Contribution incorporated within the Work constitutes direct + or contributory patent infringement, then any patent licenses + granted to You under this License for that Work shall terminate + as of the date such litigation is filed. + + 4. Redistribution. You may reproduce and distribute copies of the + Work or Derivative Works thereof in any medium, with or without + modifications, and in Source or Object form, provided that You + meet the following conditions: + + (a) You must give any other recipients of the Work or + Derivative Works a copy of this License; and + + (b) You must cause any modified files to carry prominent notices + stating that You changed the files; and + + (c) You must retain, in the Source form of any Derivative Works + that You distribute, all copyright, patent, trademark, and + attribution notices from the Source form of the Work, + excluding those notices that do not pertain to any part of + the Derivative Works; and + + (d) If the Work includes a "NOTICE" text file as part of its + distribution, then any Derivative Works that You distribute must + include a readable copy of the attribution notices contained + within such NOTICE file, excluding those notices that do not + pertain to any part of the Derivative Works, in at least one + of the following places: within a NOTICE text file distributed + as part of the Derivative Works; within the Source form or + documentation, if provided along with the Derivative Works; or, + within a display generated by the Derivative Works, if and + wherever such third-party notices normally appear. The contents + of the NOTICE file are for informational purposes only and + do not modify the License. You may add Your own attribution + notices within Derivative Works that You distribute, alongside + or as an addendum to the NOTICE text from the Work, provided + that such additional attribution notices cannot be construed + as modifying the License. + + You may add Your own copyright statement to Your modifications and + may provide additional or different license terms and conditions + for use, reproduction, or distribution of Your modifications, or + for any such Derivative Works as a whole, provided Your use, + reproduction, and distribution of the Work otherwise complies with + the conditions stated in this License. + + 5. Submission of Contributions. Unless You explicitly state otherwise, + any Contribution intentionally submitted for inclusion in the Work + by You to the Licensor shall be under the terms and conditions of + this License, without any additional terms or conditions. + Notwithstanding the above, nothing herein shall supersede or modify + the terms of any separate license agreement you may have executed + with Licensor regarding such Contributions. + + 6. Trademarks. This License does not grant permission to use the trade + names, trademarks, service marks, or product names of the Licensor, + except as required for reasonable and customary use in describing the + origin of the Work and reproducing the content of the NOTICE file. + + 7. Disclaimer of Warranty. Unless required by applicable law or + agreed to in writing, Licensor provides the Work (and each + Contributor provides its Contributions) on an "AS IS" BASIS, + WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or + implied, including, without limitation, any warranties or conditions + of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A + PARTICULAR PURPOSE. You are solely responsible for determining the + appropriateness of using or redistributing the Work and assume any + risks associated with Your exercise of permissions under this License. + + 8. Limitation of Liability. In no event and under no legal theory, + whether in tort (including negligence), contract, or otherwise, + unless required by applicable law (such as deliberate and grossly + negligent acts) or agreed to in writing, shall any Contributor be + liable to You for damages, including any direct, indirect, special, + incidental, or consequential damages of any character arising as a + result of this License or out of the use or inability to use the + Work (including but not limited to damages for loss of goodwill, + work stoppage, computer failure or malfunction, or any and all + other commercial damages or losses), even if such Contributor + has been advised of the possibility of such damages. + + 9. Accepting Warranty or Additional Liability. While redistributing + the Work or Derivative Works thereof, You may choose to offer, + and charge a fee for, acceptance of support, warranty, indemnity, + or other liability obligations and/or rights consistent with this + License. However, in accepting such obligations, You may act only + on Your own behalf and on Your sole responsibility, not on behalf + of any other Contributor, and only if You agree to indemnify, + defend, and hold each Contributor harmless for any liability + incurred by, or claims asserted against, such Contributor by reason + of your accepting any such warranty or additional liability. + + END OF TERMS AND CONDITIONS + + APPENDIX: How to apply the Apache License to your work. + + To apply the Apache License to your work, attach the following + boilerplate notice, with the fields enclosed by brackets "{}" + replaced with your own identifying information. (Don't include + the brackets!) The text should be enclosed in the appropriate + comment syntax for the file format. We also recommend that a + file or class name and description of purpose be included on the + same "printed page" as the copyright notice for easier + identification within third-party archives. + + Copyright {yyyy} {name of copyright owner} + + Licensed 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 + + https://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. diff --git a/client/rest-http-client/licenses/reactor-core-NOTICE.txt b/client/rest-http-client/licenses/reactor-core-NOTICE.txt new file mode 100644 index 0000000000000..e69de29bb2d1d diff --git a/client/rest-http-client/licenses/slf4j-api-2.0.17.jar.sha1 b/client/rest-http-client/licenses/slf4j-api-2.0.17.jar.sha1 new file mode 100644 index 0000000000000..435f6c13a28b6 --- /dev/null +++ b/client/rest-http-client/licenses/slf4j-api-2.0.17.jar.sha1 @@ -0,0 +1 @@ +d9e58ac9c7779ba3bf8142aff6c830617a7fe60f \ No newline at end of file diff --git a/client/rest-http-client/licenses/slf4j-api-LICENSE.txt b/client/rest-http-client/licenses/slf4j-api-LICENSE.txt new file mode 100644 index 0000000000000..8fda22f4d72f6 --- /dev/null +++ b/client/rest-http-client/licenses/slf4j-api-LICENSE.txt @@ -0,0 +1,21 @@ +Copyright (c) 2004-2014 QOS.ch +All rights reserved. + +Permission is hereby granted, free of charge, to any person obtaining +a copy of this software and associated documentation files (the +"Software"), to deal in the Software without restriction, including +without limitation the rights to use, copy, modify, merge, publish, +distribute, sublicense, and/or sell copies of the Software, and to +permit persons to whom the Software is furnished to do so, subject to +the following conditions: + +The above copyright notice and this permission notice shall be +included in all copies or substantial portions of the Software. + +THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, +EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF +MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND +NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE +LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION +OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION +WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE. diff --git a/client/rest-http-client/licenses/slf4j-api-NOTICE.txt b/client/rest-http-client/licenses/slf4j-api-NOTICE.txt new file mode 100644 index 0000000000000..e69de29bb2d1d diff --git a/client/rest-http-client/src/main/java/org/opensearch/internal/httpclient/AsyncResponseProducer.java b/client/rest-http-client/src/main/java/org/opensearch/internal/httpclient/AsyncResponseProducer.java new file mode 100644 index 0000000000000..aecc1345100c3 --- /dev/null +++ b/client/rest-http-client/src/main/java/org/opensearch/internal/httpclient/AsyncResponseProducer.java @@ -0,0 +1,42 @@ +/* + * SPDX-License-Identifier: Apache-2.0 + * + * The OpenSearch Contributors require contributions made to + * this file be licensed under the Apache-2.0 license or a + * compatible open source license. + */ + +package org.opensearch.internal.httpclient; + +import java.nio.ByteBuffer; +import java.util.List; +import java.util.Vector; +import java.util.concurrent.Flow.Subscriber; +import java.util.concurrent.Flow.Subscription; + +class AsyncResponseProducer implements Subscriber> { + private Subscription subscription; + private final List buffers = new Vector<>(); + + @Override + public void onSubscribe(Subscription subscription) { + this.subscription = subscription; + subscription.request(1); + } + + @Override + public void onNext(List item) { + buffers.addAll(item); + subscription.request(1); + } + + @Override + public void onError(Throwable throwable) {} + + @Override + public void onComplete() {} + + public List getResult() { + return buffers; + } +} diff --git a/client/rest-http-client/src/main/java/org/opensearch/internal/httpclient/BodyUtils.java b/client/rest-http-client/src/main/java/org/opensearch/internal/httpclient/BodyUtils.java new file mode 100644 index 0000000000000..4b1ef84f8a702 --- /dev/null +++ b/client/rest-http-client/src/main/java/org/opensearch/internal/httpclient/BodyUtils.java @@ -0,0 +1,268 @@ +/* + * SPDX-License-Identifier: Apache-2.0 + * + * The OpenSearch Contributors require contributions made to + * this file be licensed under the Apache-2.0 license or a + * compatible open source license. + */ + +package org.opensearch.internal.httpclient; + +import java.io.EOFException; +import java.io.IOException; +import java.io.InputStream; +import java.io.OutputStream; +import java.io.UncheckedIOException; +import java.net.http.HttpRequest; +import java.net.http.HttpResponse; +import java.nio.ByteBuffer; +import java.nio.channels.Channels; +import java.nio.channels.WritableByteChannel; +import java.nio.charset.StandardCharsets; +import java.util.ArrayList; +import java.util.List; +import java.util.concurrent.Flow; +import java.util.zip.GZIPInputStream; +import java.util.zip.GZIPOutputStream; + +import reactor.adapter.JdkFlowAdapter; +import reactor.core.publisher.Mono; + +/** + * Helper class to deal with request / response bodies. + */ +final class BodyUtils { + static Mono getBody(HttpRequest httpRequest) { + return httpRequest.bodyPublisher().map(JdkFlowAdapter::flowPublisherToFlux).map(Mono::from).orElseGet(Mono::empty); + } + + static String getBodyAsString(Response response) { + return getBodyAsString(response.entity()); + } + + static Mono getBodyAsString(HttpRequest httpRequest) { + return httpRequest.bodyPublisher().map(p -> { + var bodySubscriber = HttpResponse.BodySubscribers.ofString(StandardCharsets.UTF_8); + var flowSubscriber = new StringSubscriber(bodySubscriber); + p.subscribe(flowSubscriber); + return Mono.fromCompletionStage(bodySubscriber.getBody()); + }).orElseGet(Mono::empty); + } + + static String getBodyAsString(List body) { + final StringBuilder builder = new StringBuilder(); + if (body != null && body.isEmpty() == false) { + for (ByteBuffer chunk : body) { + chunk.mark(); + builder.append(StandardCharsets.UTF_8.decode(chunk).toString()); + chunk.reset(); + } + } + return builder.toString(); + } + + static List compress(List body) { + if (body == null || body.isEmpty()) { + return body; + } else { + body.stream().forEach(ByteBuffer::mark); + try (ByteBufferOutputStream bbout = new ByteBufferOutputStream(); OutputStream out = new GZIPOutputStream(bbout)) { + try (WritableByteChannel channel = Channels.newChannel(out)) { + for (ByteBuffer buffer : body) { + channel.write(buffer); + } + out.flush(); + } + return bbout.getBufferList(); + } catch (final IOException ex) { + throw new UncheckedIOException(ex); + } finally { + body.stream().forEach(ByteBuffer::reset); + } + } + } + + static ByteBuffer compress(ByteBuffer body) { + if (body == null) { + return body; + } else { + body.mark(); + try (ByteBufferOutputStream bbout = new ByteBufferOutputStream(); OutputStream out = new GZIPOutputStream(bbout)) { + try (WritableByteChannel channel = Channels.newChannel(out)) { + channel.write(body); + out.flush(); + } + + final List bufferList = bbout.getBufferList(); + if (bufferList.isEmpty() == false) { + return bufferList.get(0); + } else { + // We should never end up here + return ByteBuffer.allocate(0); /* empty byte buffer */ + } + } catch (final IOException ex) { + throw new UncheckedIOException(ex); + } finally { + body.reset(); + } + } + } + + static ByteBuffer decompress(ByteBuffer body) { + if (body == null) { + return body; + } else { + body.mark(); + + try (ByteBufferInputStream bbin = new ByteBufferInputStream(List.of(body)); InputStream in = new GZIPInputStream(bbin)) { + return ByteBuffer.wrap(in.readAllBytes()); + } catch (final IOException ex) { + throw new UncheckedIOException(ex); + } finally { + body.reset(); + } + } + } + + static List decompress(List body) { + if (body == null || body.isEmpty()) { + return body; + } else { + body.stream().forEach(ByteBuffer::mark); + try (ByteBufferInputStream bbin = new ByteBufferInputStream(body); InputStream in = new GZIPInputStream(bbin)) { + return List.of(ByteBuffer.wrap(in.readAllBytes())); + } catch (final IOException ex) { + throw new UncheckedIOException(ex); + } finally { + body.stream().forEach(ByteBuffer::reset); + } + } + } + + /** + * See please https://github.com/justinsb/avro/blob/master/src/java/org/apache/avro/ipc/ByteBufferOutputStream.java + */ + private final static class ByteBufferOutputStream extends OutputStream { + private static final int BUFFER_SIZE = 8192; + private List buffers; + + ByteBufferOutputStream() { + reset(); + } + + /** Returns all data written and resets the stream to be empty. */ + List getBufferList() { + List result = buffers; + reset(); + for (ByteBuffer buffer : result) { + buffer.flip(); + } + return result; + } + + private void reset() { + buffers = new ArrayList(1); + buffers.add(ByteBuffer.allocate(BUFFER_SIZE)); + } + + @Override + public void write(int b) { + ByteBuffer buffer = buffers.get(buffers.size() - 1); + if (buffer.remaining() < 1) { + buffer = ByteBuffer.allocate(BUFFER_SIZE); + buffers.add(buffer); + } + buffer.put((byte) b); + } + + @Override + public void write(byte[] b, int off, int len) { + ByteBuffer buffer = buffers.get(buffers.size() - 1); + int remaining = buffer.remaining(); + while (len > remaining) { + buffer.put(b, off, remaining); + len -= remaining; + off += remaining; + buffer = ByteBuffer.allocate(BUFFER_SIZE); + buffers.add(buffer); + remaining = buffer.remaining(); + } + buffer.put(b, off, len); + } + } + + /** + * See please https://github.com/justinsb/avro/blob/master/src/java/org/apache/avro/ipc/ByteBufferInputStream.java + */ + private static final class ByteBufferInputStream extends InputStream { + private List buffers; + private int current; + + ByteBufferInputStream(List buffers) { + this.buffers = buffers; + } + + /** @see InputStream#read() + * @throws EOFException if EOF is reached. */ + @Override + public int read() throws IOException { + return getBuffer().get() & 0xff; + } + + /** @see InputStream#read(byte[], int, int) + * @throws EOFException if EOF is reached before reading all the bytes. */ + @Override + public int read(byte[] b, int off, int len) throws IOException { + if (len == 0) return 0; + ByteBuffer buffer = getBuffer(); + int remaining = buffer.remaining(); + if (len > remaining) { + buffer.get(b, off, remaining); + return remaining; + } else { + buffer.get(b, off, len); + return len; + } + } + + /** Returns the next non-empty buffer. + * @throws EOFException if EOF is reached before reading all the bytes. + */ + private ByteBuffer getBuffer() throws IOException { + while (current < buffers.size()) { + ByteBuffer buffer = buffers.get(current); + if (buffer.hasRemaining()) return buffer; + current++; + } + throw new EOFException(); + } + } + + private static final class StringSubscriber implements Flow.Subscriber { + final HttpResponse.BodySubscriber delegate; + + StringSubscriber(HttpResponse.BodySubscriber delegate) { + this.delegate = delegate; + } + + @Override + public void onSubscribe(Flow.Subscription subscription) { + delegate.onSubscribe(subscription); + } + + @Override + public void onNext(ByteBuffer item) { + delegate.onNext(List.of(item)); + } + + @Override + public void onError(Throwable throwable) { + delegate.onError(throwable); + } + + @Override + public void onComplete() { + delegate.onComplete(); + } + } +} diff --git a/client/rest-http-client/src/main/java/org/opensearch/internal/httpclient/Cancellable.java b/client/rest-http-client/src/main/java/org/opensearch/internal/httpclient/Cancellable.java new file mode 100644 index 0000000000000..6357ac163e410 --- /dev/null +++ b/client/rest-http-client/src/main/java/org/opensearch/internal/httpclient/Cancellable.java @@ -0,0 +1,120 @@ +/* + * SPDX-License-Identifier: Apache-2.0 + * + * The OpenSearch Contributors require contributions made to + * this file be licensed under the Apache-2.0 license or a + * compatible open source license. + */ + +/* + * Licensed to Elasticsearch under one or more contributor + * license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright + * ownership. Elasticsearch 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. + */ +/* + * Modifications Copyright OpenSearch Contributors. See + * GitHub history for details. + */ + +package org.opensearch.internal.httpclient; + +import java.io.IOException; +import java.util.concurrent.Callable; +import java.util.concurrent.CancellationException; +import java.util.concurrent.Future; + +/** + * Represents an operation that can be cancelled. + * Returned when executing async requests through {@link RestHttpClient#performRequestAsync(Request, ResponseListener)}, so that the request + * can be cancelled if needed. Cancelling a request will result in calling {@link Future#cancel(boolean mayInterruptIfRunning)} on the + * underlying request future object. Note that cancelling a request does not automatically translate to aborting its execution on the + * server side, which needs to be specifically implemented in each API. + */ +public class Cancellable { + static final Cancellable NO_OP = new Cancellable(null) { + @Override + public boolean cancel() { + throw new UnsupportedOperationException(); + } + + @Override + void runIfNotCancelled(Runnable runnable) { + throw new UnsupportedOperationException(); + } + }; + + static Cancellable fromFuture(Future f) { + return new Cancellable(f); + } + + private final Future future; + + private Cancellable(Future f) { + this.future = f; + } + + /** + * Cancels the on-going request that is associated with the current instance of {@link Cancellable}. + * + */ + public synchronized boolean cancel() { + return this.future.cancel(true); + } + + /** + * Executes some arbitrary code if the on-going request has not been cancelled, otherwise throws {@link CancellationException}. + * This is needed to guarantee that cancelling a request works correctly even in case {@link #cancel()} is called between different + * attempts of the same request. + * If the request has already been cancelled we don't go ahead with the next attempt, and artificially raise the + * {@link CancellationException}, otherwise we run the provided {@link Runnable} which will reset the request and send the next attempt. + * Note that this method must be synchronized as well as the {@link #cancel()} method, to prevent a request from being cancelled + * when there is no future to cancel, which would make cancelling the request a no-op. + */ + synchronized void runIfNotCancelled(Runnable runnable) { + if (this.future.isCancelled()) { + throw newCancellationException(); + } + runnable.run(); + } + + /** + * Executes some arbitrary code if the on-going request has not been cancelled, otherwise throws {@link CancellationException}. + * This is needed to guarantee that cancelling a request works correctly even in case {@link #cancel()} is called between different + * attempts of the same request. The {@link #cancel()} method can be called at anytime, + * and we need to handle the case where it gets called while there is no request being executed as one attempt may have failed and + * the subsequent attempt has not been started yet. + * If the request has already been cancelled we don't go ahead with the next attempt, and artificially raise the + * {@link CancellationException}, otherwise we run the provided {@link Runnable} which will reset the request and send the next attempt. + * Note that this method must be synchronized as well as the {@link #cancel()} method, to prevent a request from being cancelled + * when there is no future to cancel, which would make cancelling the request a no-op. + */ + synchronized T callIfNotCancelled(Callable callable) throws IOException { + if (this.future.isCancelled()) { + throw newCancellationException(); + } + try { + return callable.call(); + } catch (final IOException ex) { + throw ex; + } catch (final Exception ex) { + throw new IOException(ex); + } + } + + static CancellationException newCancellationException() { + return new CancellationException("request was cancelled"); + } +} diff --git a/client/rest-http-client/src/main/java/org/opensearch/internal/httpclient/CompressedResponse.java b/client/rest-http-client/src/main/java/org/opensearch/internal/httpclient/CompressedResponse.java new file mode 100644 index 0000000000000..4060891dd48dd --- /dev/null +++ b/client/rest-http-client/src/main/java/org/opensearch/internal/httpclient/CompressedResponse.java @@ -0,0 +1,41 @@ +/* + * SPDX-License-Identifier: Apache-2.0 + * + * The OpenSearch Contributors require contributions made to + * this file be licensed under the Apache-2.0 license or a + * compatible open source license. + */ + +package org.opensearch.internal.httpclient; + +import java.io.InputStream; +import java.net.http.HttpResponse; +import java.nio.ByteBuffer; +import java.util.List; +import java.util.Objects; + +final record CompressedResponse(RequestLine requestLine, HttpHost host, HttpResponse httpResponse, List entity) + implements + Response { + CompressedResponse { + requestLine = Objects.requireNonNull(requestLine, "requestLine cannot be null"); + host = Objects.requireNonNull(host, "host cannot be null"); + httpResponse = Objects.requireNonNull(httpResponse, "response cannot be null"); + } + + /** + * Returns the response body available, null otherwise + * @see InputStream + */ + public List entity() { + return BodyUtils.decompress(entity); + } + + /** + * Convert response to string representation + */ + @Override + public String toString() { + return "Response{" + "requestLine=" + requestLine() + ", host=" + host() + ", response=" + statusLine() + '}'; + } +} diff --git a/client/rest-http-client/src/main/java/org/opensearch/internal/httpclient/DeadHostState.java b/client/rest-http-client/src/main/java/org/opensearch/internal/httpclient/DeadHostState.java new file mode 100644 index 0000000000000..de2e5161ea9c9 --- /dev/null +++ b/client/rest-http-client/src/main/java/org/opensearch/internal/httpclient/DeadHostState.java @@ -0,0 +1,125 @@ +/* + * SPDX-License-Identifier: Apache-2.0 + * + * The OpenSearch Contributors require contributions made to + * this file be licensed under the Apache-2.0 license or a + * compatible open source license. + */ + +/* + * Licensed to Elasticsearch under one or more contributor + * license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright + * ownership. Elasticsearch 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. + */ + +/* + * Modifications Copyright OpenSearch Contributors. See + * GitHub history for details. + */ + +package org.opensearch.internal.httpclient; + +import java.util.concurrent.TimeUnit; +import java.util.function.Supplier; + +/** + * Holds the state of a dead connection to a host. Keeps track of how many failed attempts were performed and + * when the host should be retried (based on number of previous failed attempts). + * Class is immutable, a new copy of it should be created each time the state has to be changed. + */ +final class DeadHostState implements Comparable { + + private static final long MIN_CONNECTION_TIMEOUT_NANOS = TimeUnit.MINUTES.toNanos(1); + static final long MAX_CONNECTION_TIMEOUT_NANOS = TimeUnit.MINUTES.toNanos(30); + static final Supplier DEFAULT_TIME_SUPPLIER = System::nanoTime; + + private final int failedAttempts; + private final long deadUntilNanos; + private final Supplier timeSupplier; + + /** + * Build the initial dead state of a host. Useful when a working host stops functioning + * and needs to be marked dead after its first failure. In such case the host will be retried after a minute or so. + * + * @param timeSupplier a way to supply the current time and allow for unit testing + */ + DeadHostState(Supplier timeSupplier) { + this.failedAttempts = 1; + this.deadUntilNanos = timeSupplier.get() + MIN_CONNECTION_TIMEOUT_NANOS; + this.timeSupplier = timeSupplier; + } + + /** + * Build the dead state of a host given its previous dead state. Useful when a host has been failing before, hence + * it already failed for one or more consecutive times. The more failed attempts we register the longer we wait + * to retry that same host again. Minimum is 1 minute (for a node the only failed once created + * through {@link #DeadHostState(Supplier)}), maximum is 30 minutes (for a node that failed more than 10 consecutive times) + * + * @param previousDeadHostState the previous state of the host which allows us to increase the wait till the next retry attempt + */ + DeadHostState(DeadHostState previousDeadHostState) { + long timeoutNanos = (long) Math.min( + MIN_CONNECTION_TIMEOUT_NANOS * 2 * Math.pow(2, previousDeadHostState.failedAttempts * 0.5 - 1), + MAX_CONNECTION_TIMEOUT_NANOS + ); + this.deadUntilNanos = previousDeadHostState.timeSupplier.get() + timeoutNanos; + this.failedAttempts = previousDeadHostState.failedAttempts + 1; + this.timeSupplier = previousDeadHostState.timeSupplier; + } + + /** + * Indicates whether it's time to retry to failed host or not. + * + * @return true if the host should be retried, false otherwise + */ + boolean shallBeRetried() { + return timeSupplier.get() - deadUntilNanos > 0; + } + + /** + * Returns the timestamp (nanos) till the host is supposed to stay dead without being retried. + * After that the host should be retried. + */ + long getDeadUntilNanos() { + return deadUntilNanos; + } + + int getFailedAttempts() { + return failedAttempts; + } + + @Override + public int compareTo(DeadHostState other) { + if (timeSupplier != other.timeSupplier) { + throw new IllegalArgumentException( + "can't compare DeadHostStates holding different time suppliers as they may " + "be based on different clocks" + ); + } + return Long.compare(deadUntilNanos, other.deadUntilNanos); + } + + @Override + public String toString() { + return "DeadHostState{" + + "failedAttempts=" + + failedAttempts + + ", deadUntilNanos=" + + deadUntilNanos + + ", timeSupplier=" + + timeSupplier + + '}'; + } +} diff --git a/client/rest-http-client/src/main/java/org/opensearch/internal/httpclient/HttpHost.java b/client/rest-http-client/src/main/java/org/opensearch/internal/httpclient/HttpHost.java new file mode 100644 index 0000000000000..cb7d7d6106a20 --- /dev/null +++ b/client/rest-http-client/src/main/java/org/opensearch/internal/httpclient/HttpHost.java @@ -0,0 +1,52 @@ +/* + * SPDX-License-Identifier: Apache-2.0 + * + * The OpenSearch Contributors require contributions made to + * this file be licensed under the Apache-2.0 license or a + * compatible open source license. + */ + +package org.opensearch.internal.httpclient; + +import java.net.URI; +import java.net.URISyntaxException; +import java.util.Objects; + +/** + * Represents HTTP host (scheme, hostname and port) + * Note: This is an experimental API. + */ +public record HttpHost(String scheme, String hostname, int port) { + /** + * Converts the host to string + */ + @Override + public String toString() { + final StringBuilder buffer = new StringBuilder(); + buffer.append(this.scheme); + buffer.append("://"); + buffer.append(this.hostname); + if (this.port != -1) { + buffer.append(':'); + buffer.append(Integer.toString(this.port)); + } + return buffer.toString(); + } + + public static HttpHost create(String uriStr) throws URISyntaxException { + Objects.requireNonNull(uriStr); + + String text = uriStr; + String scheme = null; + final int schemeIdx = text.indexOf("://"); + if (schemeIdx > 0) { + scheme = text.substring(0, schemeIdx); + if (scheme.isBlank()) { + throw new URISyntaxException(uriStr, "scheme contains blanks"); + } + } + + final URI uri = new URI(uriStr); + return new HttpHost(scheme, uri.getHost(), uri.getPort()); + } +} diff --git a/client/rest-http-client/src/main/java/org/opensearch/internal/httpclient/Node.java b/client/rest-http-client/src/main/java/org/opensearch/internal/httpclient/Node.java new file mode 100644 index 0000000000000..c892b6fae2b64 --- /dev/null +++ b/client/rest-http-client/src/main/java/org/opensearch/internal/httpclient/Node.java @@ -0,0 +1,278 @@ +/* + * SPDX-License-Identifier: Apache-2.0 + * + * The OpenSearch Contributors require contributions made to + * this file be licensed under the Apache-2.0 license or a + * compatible open source license. + */ + +/* + * Licensed to Elasticsearch under one or more contributor + * license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright + * ownership. Elasticsearch 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. + */ + +/* + * Modifications Copyright OpenSearch Contributors. See + * GitHub history for details. + */ + +package org.opensearch.internal.httpclient; + +import java.util.List; +import java.util.Map; +import java.util.Objects; +import java.util.Set; +import java.util.TreeSet; + +/** + * Metadata about an {@link HttpHost} running OpenSearch. + */ +public class Node { + /** + * Address that this host claims is its primary contact point. + */ + private final HttpHost host; + /** + * Addresses on which the host is listening. These are useful to have + * around because they allow you to find a host based on any address it + * is listening on. + */ + private final Set boundHosts; + /** + * Name of the node as configured by the {@code node.name} attribute. + */ + private final String name; + /** + * Version of OpenSearch that the node is running or {@code null} + * if we don't know the version. + */ + private final String version; + /** + * Roles that the OpenSearch process on the host has or {@code null} + * if we don't know what roles the node has. + */ + private final Roles roles; + /** + * Attributes declared on the node. + */ + private final Map> attributes; + + /** + * Create a {@linkplain Node} with metadata. All parameters except + * {@code host} are nullable and implementations of {@link NodeSelector} + * need to decide what to do in their absence. + * + * @param host primary host address + * @param boundHosts addresses on which the host is listening + * @param name name of the node + * @param version version of OpenSearch + * @param roles roles that the OpenSearch process has on the host + * @param attributes attributes declared on the node + */ + public Node(HttpHost host, Set boundHosts, String name, String version, Roles roles, Map> attributes) { + if (host == null) { + throw new IllegalArgumentException("host cannot be null"); + } + this.host = host; + this.boundHosts = boundHosts; + this.name = name; + this.version = version; + this.roles = roles; + this.attributes = attributes; + } + + /** + * Create a {@linkplain Node} without any metadata. + * + * @param host primary host address + */ + public Node(HttpHost host) { + this(host, null, null, null, null, null); + } + + /** + * Contact information for the host. + */ + public HttpHost getHost() { + return host; + } + + /** + * Addresses on which the host is listening. These are useful to have + * around because they allow you to find a host based on any address it + * is listening on. + */ + public Set getBoundHosts() { + return boundHosts; + } + + /** + * The {@code node.name} of the node. + */ + public String getName() { + return name; + } + + /** + * Version of OpenSearch that the node is running or {@code null} + * if we don't know the version. + */ + public String getVersion() { + return version; + } + + /** + * Roles that the OpenSearch process on the host has or {@code null} + * if we don't know what roles the node has. + */ + public Roles getRoles() { + return roles; + } + + /** + * Attributes declared on the node. + */ + public Map> getAttributes() { + return attributes; + } + + /** + * Convert node to string representation + */ + @Override + public String toString() { + StringBuilder b = new StringBuilder(); + b.append("[host=").append(host); + if (boundHosts != null) { + b.append(", bound=").append(boundHosts); + } + if (name != null) { + b.append(", name=").append(name); + } + if (version != null) { + b.append(", version=").append(version); + } + if (roles != null) { + b.append(", roles=").append(roles); + } + if (attributes != null) { + b.append(", attributes=").append(attributes); + } + return b.append(']').toString(); + } + + /** + * Compare two nodes for equality + * @param obj node instance to compare with + */ + @Override + public boolean equals(Object obj) { + if (obj == null || obj.getClass() != getClass()) { + return false; + } + Node other = (Node) obj; + return host.equals(other.host) + && Objects.equals(boundHosts, other.boundHosts) + && Objects.equals(name, other.name) + && Objects.equals(version, other.version) + && Objects.equals(roles, other.roles) + && Objects.equals(attributes, other.attributes); + } + + /** + * Calculate the hash code of the node + */ + @Override + public int hashCode() { + return Objects.hash(host, boundHosts, name, version, roles, attributes); + } + + /** + * Role information about an OpenSearch process. + */ + public static final class Roles { + + private final Set roles; + + /** + * Create a {@link Roles} instance of the given string set. + * + * @param roles set of role names. + */ + public Roles(final Set roles) { + this.roles = new TreeSet<>(roles); + } + + /** + * Returns whether or not the node could be elected cluster-manager. + */ + public boolean isClusterManagerEligible() { + return roles.contains("master") || roles.contains("cluster_manager"); + } + + /** + * Returns whether or not the node stores data. + */ + public boolean isData() { + return roles.contains("data"); + } + + /** + * Returns whether or not the node runs ingest pipelines. + */ + public boolean isIngest() { + return roles.contains("ingest"); + } + + /** + * Returns whether the node is dedicated to provide search capability. + */ + public boolean isSearch() { + return roles.contains("search"); + } + + /** + * Convert roles to string representation + */ + @Override + public String toString() { + return String.join(",", roles); + } + + /** + * Compare two roles for equality + * @param obj roles instance to compare with + */ + @Override + public boolean equals(Object obj) { + if (obj == null || obj.getClass() != getClass()) { + return false; + } + Roles other = (Roles) obj; + return roles.equals(other.roles); + } + + /** + * Calculate the hash code of the roles + */ + @Override + public int hashCode() { + return roles.hashCode(); + } + + } +} diff --git a/client/rest-http-client/src/main/java/org/opensearch/internal/httpclient/NodeSelector.java b/client/rest-http-client/src/main/java/org/opensearch/internal/httpclient/NodeSelector.java new file mode 100644 index 0000000000000..4a4856a920b48 --- /dev/null +++ b/client/rest-http-client/src/main/java/org/opensearch/internal/httpclient/NodeSelector.java @@ -0,0 +1,105 @@ +/* + * SPDX-License-Identifier: Apache-2.0 + * + * The OpenSearch Contributors require contributions made to + * this file be licensed under the Apache-2.0 license or a + * compatible open source license. + */ + +/* + * Licensed to Elasticsearch under one or more contributor + * license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright + * ownership. Elasticsearch 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. + */ + +/* + * Modifications Copyright OpenSearch Contributors. See + * GitHub history for details. + */ + +package org.opensearch.internal.httpclient; + +import java.util.Iterator; + +/** + * Selects nodes that can receive requests. Used to keep requests away + * from cluster-manager nodes or to send them to nodes with a particular attribute. + * Use with {@link RestHttpClientBuilder#setNodeSelector(NodeSelector)}. + */ +public interface NodeSelector { + /** + * Select the {@link Node}s to which to send requests. This is called with + * a mutable {@link Iterable} of {@linkplain Node}s in the order that the + * rest client would prefer to use them and implementers should remove + * nodes from the that should not receive the request. Implementers may + * iterate the nodes as many times as they need. + *

    + * This may be called twice per request: first for "living" nodes that + * have not been denylisted by previous errors. If the selector removes + * all nodes from the list or if there aren't any living nodes then the + * {@link RestHttpClient} will call this method with a list of "dead" nodes. + *

    + * Implementers should not rely on the ordering of the nodes. + * + * @param nodes the {@link Node}s targeted for the sending requests + */ + void select(Iterable nodes); + /* + * We were fairly careful with our choice of Iterable here. The caller has + * a List but reordering the list is likely to break round robin. Luckily + * Iterable doesn't allow any reordering. + */ + + /** + * Selector that matches any node. + */ + NodeSelector ANY = new NodeSelector() { + @Override + public void select(Iterable nodes) { + // Intentionally does nothing + } + + @Override + public String toString() { + return "ANY"; + } + }; + + /** + * Selector that matches any node that has metadata and doesn't + * have the {@code cluster_manager} role OR it has the data {@code data} + * role. + */ + NodeSelector SKIP_DEDICATED_CLUSTER_MANAGERS = new NodeSelector() { + @Override + public void select(Iterable nodes) { + for (Iterator itr = nodes.iterator(); itr.hasNext();) { + Node node = itr.next(); + if (node.getRoles() == null) continue; + if (node.getRoles().isClusterManagerEligible() + && false == node.getRoles().isData() + && false == node.getRoles().isIngest()) { + itr.remove(); + } + } + } + + @Override + public String toString() { + return "SKIP_DEDICATED_CLUSTER_MANAGERS"; + } + }; +} diff --git a/client/rest-http-client/src/main/java/org/opensearch/internal/httpclient/NonCompressedResponse.java b/client/rest-http-client/src/main/java/org/opensearch/internal/httpclient/NonCompressedResponse.java new file mode 100644 index 0000000000000..67137c468d2d6 --- /dev/null +++ b/client/rest-http-client/src/main/java/org/opensearch/internal/httpclient/NonCompressedResponse.java @@ -0,0 +1,32 @@ +/* + * SPDX-License-Identifier: Apache-2.0 + * + * The OpenSearch Contributors require contributions made to + * this file be licensed under the Apache-2.0 license or a + * compatible open source license. + */ + +package org.opensearch.internal.httpclient; + +import java.net.http.HttpResponse; +import java.nio.ByteBuffer; +import java.util.List; +import java.util.Objects; + +final record NonCompressedResponse(RequestLine requestLine, HttpHost host, HttpResponse httpResponse, List entity) + implements + Response { + NonCompressedResponse { + requestLine = Objects.requireNonNull(requestLine, "requestLine cannot be null"); + host = Objects.requireNonNull(host, "host cannot be null"); + httpResponse = Objects.requireNonNull(httpResponse, "response cannot be null"); + } + + /** + * Convert response to string representation + */ + @Override + public String toString() { + return "Response{" + "requestLine=" + requestLine() + ", host=" + host() + ", response=" + statusLine() + '}'; + } +} diff --git a/client/rest-http-client/src/main/java/org/opensearch/internal/httpclient/Request.java b/client/rest-http-client/src/main/java/org/opensearch/internal/httpclient/Request.java new file mode 100644 index 0000000000000..ed4cc1e8f19d4 --- /dev/null +++ b/client/rest-http-client/src/main/java/org/opensearch/internal/httpclient/Request.java @@ -0,0 +1,160 @@ +/* + * SPDX-License-Identifier: Apache-2.0 + * + * The OpenSearch Contributors require contributions made to + * this file be licensed under the Apache-2.0 license or a + * compatible open source license. + */ + +package org.opensearch.internal.httpclient; + +import java.net.http.HttpRequest.BodyPublisher; +import java.net.http.HttpRequest.BodyPublishers; +import java.util.HashMap; +import java.util.Map; +import java.util.Objects; + +import static java.util.Collections.unmodifiableMap; + +/** + * HTTP Request to OpenSearch. + * Note: This is an experimental API. + */ +public record Request(String method, String endpoint, Map parameters, BodyPublisher entity, RequestOptions options) { + + public Request { + method = Objects.requireNonNull(method, "method cannot be null"); + endpoint = Objects.requireNonNull(endpoint, "endpoint cannot be null"); + parameters = parameters == null ? new HashMap<>() : parameters; + options = options == null ? RequestOptions.DEFAULT : options; + } + + /** + * Query string parameters. The returned map is an unmodifiable view of the + * map in the request so calls to {@link Map#put(Object, Object)} + * will change it. + */ + @Override + public Map parameters() { + if (options.getParameters().isEmpty()) { + return unmodifiableMap(parameters); + } else { + Map combinedParameters = new HashMap<>(parameters); + combinedParameters.putAll(options.getParameters()); + return unmodifiableMap(combinedParameters); + } + } + + public static Request.Builder newRequest(String method, String endpoint) { + return new Request.Builder(method, endpoint); + } + + public static final class Builder { + private final String method; + private final String endpoint; + private final Map parameters = new HashMap<>(); + private BodyPublisher entity; + private RequestOptions options = RequestOptions.DEFAULT; + + private Builder(String method, String endpoint) { + this.method = Objects.requireNonNull(method, "method cannot be null"); + this.endpoint = Objects.requireNonNull(endpoint, "endpoint cannot be null"); + } + + /** + * Set the portion of an HTTP request to OpenSearch that can be + * manipulated without changing OpenSearch's behavior. + * + * @param options the options to be set. + * @throws NullPointerException if {@code options} is null. + */ + public void setOptions(RequestOptions options) { + Objects.requireNonNull(options, "options cannot be null"); + this.options = options; + } + + /** + * Add a query string parameter. + * @param name the name of the url parameter. Must not be null. + * @param value the value of the url url parameter. If {@code null} then + * the parameter is sent as {@code name} rather than {@code name=value} + * @throws IllegalArgumentException if a parameter with that name has + * already been set + */ + public Builder withParameter(String name, String value) { + Objects.requireNonNull(name, "url parameter name cannot be null"); + if (parameters.containsKey(name)) { + throw new IllegalArgumentException("url parameter [" + name + "] has already been set to [" + parameters.get(name) + "]"); + } else { + parameters.put(name, value); + } + return this; + } + + /** + * Add query parameters using the provided map of key value pairs. + * + * @param paramSource a map of key value pairs where the key is the url parameter. + * @throws IllegalArgumentException if a parameter with that name has already been set. + */ + public Builder withParameters(Map paramSource) { + paramSource.forEach(this::withParameter); + return this; + } + + /** + * Set the portion of an HTTP request to OpenSearch that can be + * manipulated without changing OpenSearch's behavior. + * + * @param options the options to be set. + * @throws NullPointerException if {@code options} is null. + */ + public Builder withOptions(RequestOptions.Builder options) { + Objects.requireNonNull(options, "options cannot be null"); + this.options = options.build(); + return this; + } + + /** + * Set the portion of an HTTP request to OpenSearch that can be + * manipulated without changing OpenSearch's behavior. + * + * @param options the options to be set. + * @throws NullPointerException if {@code options} is null. + */ + public Builder withOptions(RequestOptions options) { + Objects.requireNonNull(options, "options cannot be null"); + this.options = options; + return this; + } + + /** + * Set the body of the request. If not set or set to {@code null} then no + * body is sent with the request. + * + * @param entity the {@link BodyPublisher} to be set as the body of the request. + */ + public Builder withEntity(BodyPublisher entity) { + this.entity = entity; + return this; + } + + /** + * Set the body of the request to a string. If not set or set to + * {@code null} then no body is sent with the request. The + * {@code Content-Type} will be sent as {@code application/json}. + * If you need a different content type then use + * {@link #withEntity(BodyPublisher)}. + * + * @param entity JSON string to be set as the entity body of the request. + */ + public Builder withEntity(String entity) { + withEntity(entity == null ? BodyPublishers.noBody() : BodyPublishers.ofString(entity)); + return this; + } + + public Request build() { + return new Request(method, endpoint, parameters, entity, options); + } + } +} diff --git a/client/rest-http-client/src/main/java/org/opensearch/internal/httpclient/RequestLine.java b/client/rest-http-client/src/main/java/org/opensearch/internal/httpclient/RequestLine.java new file mode 100644 index 0000000000000..06899bbf86d5a --- /dev/null +++ b/client/rest-http-client/src/main/java/org/opensearch/internal/httpclient/RequestLine.java @@ -0,0 +1,68 @@ +/* + * SPDX-License-Identifier: Apache-2.0 + * + * The OpenSearch Contributors require contributions made to + * this file be licensed under the Apache-2.0 license or a + * compatible open source license. + */ + +package org.opensearch.internal.httpclient; + +import java.io.Serializable; +import java.net.URI; +import java.net.http.HttpClient.Version; +import java.net.http.HttpRequest; +import java.util.Objects; + +/** + * Request line (protocol, method, uri) + * Note: This is an experimental API. + */ +public record RequestLine(String method, String uri, Version protocolVersion) implements Serializable { + private static final long serialVersionUID = 2810581718468737193L; + + public RequestLine { + method = Objects.requireNonNull(method, "Method"); + } + + /** + * Create a new instance from the request + * @param request HTTP request + */ + public RequestLine(final HttpRequest request) { + this(Objects.requireNonNull(request, "Request").method(), buildUri(request.uri()), request.version().orElse(Version.HTTP_1_1)); + } + + private static String buildUri(URI uri) { + final String query = uri.getQuery(); + if (query != null && query.isBlank() == false) { + return uri.getPath() + "?" + query; + } else { + return uri.getPath(); + } + } + + /** + * Creates new request line instance + * @param method request HTTP method + * @param uri request uri + * @param version HTTP protocol + */ + public RequestLine(final String method, final URI uri, final Version version) { + this( + Objects.requireNonNull(method, "Method"), + Objects.requireNonNull(uri, "URI").getPath(), + version != null ? version : Version.HTTP_1_1 + ); + } + + /** + * Converts the request line to string + */ + @Override + public String toString() { + final StringBuilder buf = new StringBuilder(); + buf.append(this.method).append(" ").append(this.uri).append(" ").append(this.protocolVersion); + return buf.toString(); + } +} diff --git a/client/rest-http-client/src/main/java/org/opensearch/internal/httpclient/RequestLogger.java b/client/rest-http-client/src/main/java/org/opensearch/internal/httpclient/RequestLogger.java new file mode 100644 index 0000000000000..1120669b889ba --- /dev/null +++ b/client/rest-http-client/src/main/java/org/opensearch/internal/httpclient/RequestLogger.java @@ -0,0 +1,190 @@ +/* + * SPDX-License-Identifier: Apache-2.0 + * + * The OpenSearch Contributors require contributions made to + * this file be licensed under the Apache-2.0 license or a + * compatible open source license. + */ + +/* + * Licensed to Elasticsearch under one or more contributor + * license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright + * ownership. Elasticsearch 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. + */ + +/* + * Modifications Copyright OpenSearch Contributors. See + * GitHub history for details. + */ + +package org.opensearch.internal.httpclient; + +import org.apache.commons.logging.Log; +import org.apache.commons.logging.LogFactory; + +import java.io.IOException; +import java.net.http.HttpRequest; +import java.net.http.HttpResponse; +import java.nio.ByteBuffer; +import java.nio.charset.StandardCharsets; +import java.util.List; +import java.util.Map; +import java.util.concurrent.Flow; +import java.util.function.Function; +import java.util.stream.Collectors; + +/** + * Helper class that exposes static methods to unify the way requests are logged. + * Includes trace logging to log complete requests and responses in curl format. + * Useful for debugging, manually sending logged requests via curl and checking their responses. + * Trace logging is a feature that all the language clients provide. + */ +final class RequestLogger { + + private static final Log tracer = LogFactory.getLog("tracer"); + + private RequestLogger() {} + + /** + * Logs a streaming request that yielded a streaming response + */ + static void logStreamingResponse( + Log logger, + HttpRequest request, + HttpHost host, + HttpResponse>> httpResponse + ) { + logResponse(logger, request, host, httpResponse, List.of()); + } + + /** + * Logs a request that yielded a response + */ + static void logResponse(Log logger, HttpRequest request, HttpHost host, HttpResponse> httpResponse) { + logResponse(logger, request, host, httpResponse, httpResponse.body()); + } + + /** + * Logs a request that failed + */ + static void logFailedRequest(Log logger, Function request, Node node, Exception e) { + if (logger.isDebugEnabled()) { + final HttpRequest r = request.apply(node); + logger.debug("request [" + r.method() + " " + node.getHost() + r.uri() + "] failed", e); + } + if (tracer.isTraceEnabled()) { + String traceRequest; + try { + traceRequest = buildTraceRequest(request, node); + } catch (IOException e1) { + tracer.trace("error while reading request for trace purposes", e); + traceRequest = ""; + } + tracer.trace(traceRequest); + } + } + + static String buildWarningMessage(HttpRequest request, HttpHost host, List warnings) { + StringBuilder message = new StringBuilder("request [").append(request.method()) + .append(" ") + .append(host) + .append(request.uri()) + .append("] returned ") + .append(warnings.size()) + .append(" warnings: "); + for (int i = 0; i < warnings.size(); i++) { + if (i > 0) { + message.append(","); + } + message.append("[").append(warnings.get(i)).append("]"); + } + return message.toString(); + } + + /** + * Creates curl output for given request + */ + static String buildTraceRequest(Function request, Node node) throws IOException { + final HttpRequest r = request.apply(node); + return buildTraceRequest(r, node.getHost()); + } + + static String buildTraceRequest(HttpRequest request, HttpHost host) throws IOException { + String requestLine = "curl -iX " + request.method() + " '" + request.uri() + "'"; + final String body = BodyUtils.getBodyAsString(request).block(); + if (body != null) { + requestLine += " -d '"; + requestLine += body + "'"; + } + return requestLine; + } + + /** + * Creates curl output for given response + */ + static String buildTraceResponse(HttpResponse httpResponse, List body) throws IOException { + StringBuilder responseLine = new StringBuilder(); + responseLine.append("# ").append(new StatusLine(httpResponse)); + for (Map.Entry> header : httpResponse.headers().map().entrySet()) { + responseLine.append("\n# ") + .append(header.getKey()) + .append(": ") + .append(header.getValue().stream().collect(Collectors.joining(","))); + } + responseLine.append("\n#"); + + if (body != null && body.isEmpty() == false) { + for (ByteBuffer chunk : body) { + responseLine.append("\n# ").append(StandardCharsets.UTF_8.decode(chunk).toString()); + } + } + + return responseLine.toString(); + } + + /** + * Logs a request that yielded a response + */ + private static void logResponse(Log logger, HttpRequest request, HttpHost host, HttpResponse httpResponse, List body) { + if (logger.isDebugEnabled()) { + logger.debug("request [" + request.method() + " " + request.uri() + "] returned [" + new StatusLine(httpResponse) + "]"); + } + if (logger.isWarnEnabled()) { + List warnings = httpResponse.headers().allValues("Warning"); + if (warnings != null && warnings.size() > 0) { + logger.warn(buildWarningMessage(request, host, warnings)); + } + } + if (tracer.isTraceEnabled()) { + String requestLine; + try { + requestLine = buildTraceRequest(request, host); + } catch (IOException e) { + requestLine = ""; + tracer.trace("error while reading request for trace purposes", e); + } + String responseLine; + try { + responseLine = buildTraceResponse(httpResponse, body); + } catch (IOException e) { + responseLine = ""; + tracer.trace("error while reading response for trace purposes", e); + } + tracer.trace(requestLine + '\n' + responseLine); + } + } + +} diff --git a/client/rest-http-client/src/main/java/org/opensearch/internal/httpclient/RequestOptions.java b/client/rest-http-client/src/main/java/org/opensearch/internal/httpclient/RequestOptions.java new file mode 100644 index 0000000000000..c814b257c6100 --- /dev/null +++ b/client/rest-http-client/src/main/java/org/opensearch/internal/httpclient/RequestOptions.java @@ -0,0 +1,289 @@ +/* + * SPDX-License-Identifier: Apache-2.0 + * + * The OpenSearch Contributors require contributions made to + * this file be licensed under the Apache-2.0 license or a + * compatible open source license. + */ + +/* + * Licensed to Elasticsearch under one or more contributor + * license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright + * ownership. Elasticsearch 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. + */ + +/* + * Modifications Copyright OpenSearch Contributors. See + * GitHub history for details. + */ + +package org.opensearch.internal.httpclient; + +import java.time.Duration; +import java.util.ArrayList; +import java.util.Collections; +import java.util.HashMap; +import java.util.List; +import java.util.Map; +import java.util.Objects; +import java.util.stream.Collectors; + +/** + * The portion of an HTTP request to OpenSearch that can be + * manipulated without changing OpenSearch's behavior. + * Note: This is an experimental API. + */ +public final class RequestOptions { + /** + * Default request options. + */ + public static final RequestOptions DEFAULT = new Builder( + Collections.emptyMap(), + Collections.emptyMap(), + null, + Duration.ofMillis(RestHttpClientBuilder.DEFAULT_RESPONSE_TIMEOUT_MILLIS) + ).build(); + + private final Map> headers; + private final Map parameters; + private final WarningsHandler warningsHandler; + private final Duration timeout; + + private RequestOptions(Builder builder) { + this.headers = Collections.unmodifiableMap(new HashMap<>(builder.headers)); + this.parameters = Collections.unmodifiableMap(new HashMap<>(builder.parameters)); + this.warningsHandler = builder.warningsHandler; + this.timeout = builder.timeout; + } + + /** + * Create a builder that contains these options but can be modified. + */ + public Builder toBuilder() { + return new Builder(headers, parameters, warningsHandler, timeout); + } + + /** + * Headers to attach to the request. + */ + public Map> getHeaders() { + return headers; + } + + /** + * Query parameters to attach to the request. Any parameters present here + * will override matching parameters in the {@link Request}, if they exist. + */ + public Map getParameters() { + return parameters; + } + + /** + * How this request should handle warnings. If null (the default) then + * this request will default to the behavior dictacted by + * {@link RestHttpClientBuilder#setStrictDeprecationMode}. + *

    + * This can be set to {@link WarningsHandler#PERMISSIVE} if the client + * should ignore all warnings which is the same behavior as setting + * strictDeprecationMode to true. It can be set to + * {@link WarningsHandler#STRICT} if the client should fail if there are + * any warnings which is the same behavior as settings + * strictDeprecationMode to false. + *

    + * It can also be set to a custom implementation of + * {@linkplain WarningsHandler} to permit only certain warnings or to + * fail the request if the warnings returned don't + * exactly match some set. + */ + public WarningsHandler getWarningsHandler() { + return warningsHandler; + } + + /** + * Gets the request timeout + * @return request timeout + */ + public Duration getTimeout() { + return timeout; + } + + /** + * Convert request options to string representation + */ + @Override + public String toString() { + StringBuilder b = new StringBuilder(); + b.append("RequestOptions{"); + boolean comma = false; + if (headers.size() > 0) { + comma = true; + b.append("headers="); + b.append(headers.entrySet().stream().map(e -> e.getKey() + "=" + e.getValue()).collect(Collectors.joining(","))); + } + if (timeout != null) { + if (comma) b.append(", "); + comma = true; + b.append("timeout=").append(timeout.toMillis()).append("ms"); + } + if (parameters.size() > 0) { + if (comma) b.append(", "); + comma = true; + b.append("parameters="); + b.append(parameters.entrySet().stream().map(e -> e.getKey() + "=" + e.getValue()).collect(Collectors.joining(","))); + } + if (warningsHandler != null) { + if (comma) b.append(", "); + comma = true; + b.append("warningsHandler=").append(warningsHandler); + } + return b.append('}').toString(); + } + + /** + * Compare two request options for equality + * @param obj request options instance to compare with + */ + @Override + public boolean equals(Object obj) { + if (obj == null || (obj.getClass() != getClass())) { + return false; + } + if (obj == this) { + return true; + } + + RequestOptions other = (RequestOptions) obj; + return headers.equals(other.headers) + && parameters.equals(other.parameters) + && Objects.equals(timeout, other.timeout) + && Objects.equals(warningsHandler, other.warningsHandler); + } + + /** + * Calculate the hash code of the request options + */ + @Override + public int hashCode() { + return Objects.hash(headers, parameters, warningsHandler); + } + + /** + * Builds {@link RequestOptions}. Get one by calling + * {@link RequestOptions#toBuilder} on {@link RequestOptions#DEFAULT} or + * any other {@linkplain RequestOptions}. + */ + public static class Builder { + private final Map> headers; + private final Map parameters; + private WarningsHandler warningsHandler; + private Duration timeout = Duration.ofMillis(RestHttpClientBuilder.DEFAULT_RESPONSE_TIMEOUT_MILLIS); + + private Builder( + Map> headers, + Map parameters, + WarningsHandler warningsHandler, + Duration timeout + ) { + this.headers = new HashMap<>(headers); + this.parameters = new HashMap<>(parameters); + this.warningsHandler = warningsHandler; + this.timeout = timeout; + } + + /** + * Build the {@linkplain RequestOptions}. + */ + public RequestOptions build() { + return new RequestOptions(this); + } + + /** + * Sets request timeout + * + * @param timeout timeout + */ + public void setTimeout(Duration timeout) { + Objects.requireNonNull(timeout, "timeout cannot be null"); + this.timeout = timeout; + } + + /** + * Add the provided headers to the request. + * + * @param headers headers to add + * @throws NullPointerException if {@code name} or {@code value} is null. + */ + public Builder addHeaders(Map> headers) { + Objects.requireNonNull(headers, "headers cannot be null"); + for (Map.Entry> header : headers.entrySet()) { + header.getValue().forEach(v -> addHeader(header.getKey(), v)); + } + return this; + } + + /** + * Add the provided header to the request. + * + * @param name the header name + * @param value the header value + * @throws NullPointerException if {@code name} or {@code value} is null. + */ + public Builder addHeader(String name, String value) { + Objects.requireNonNull(name, "header name cannot be null"); + Objects.requireNonNull(value, "header value cannot be null"); + this.headers.computeIfAbsent(name, v -> new ArrayList<>()).add(value); + return this; + } + + /** + * Add the provided query parameter to the request. Any parameters added here + * will override matching parameters in the {@link Request}, if they exist. + * + * @param name the query parameter name + * @param value the query parameter value + * @throws NullPointerException if {@code name} or {@code value} is null. + */ + public Builder addParameter(String name, String value) { + Objects.requireNonNull(name, "query parameter name cannot be null"); + Objects.requireNonNull(value, "query parameter value cannot be null"); + this.parameters.put(name, value); + return this; + } + + /** + * How this request should handle warnings. If null (the default) then + * this request will default to the behavior dictacted by + * {@link RestHttpClientBuilder#setStrictDeprecationMode}. + *

    + * This can be set to {@link WarningsHandler#PERMISSIVE} if the client + * should ignore all warnings which is the same behavior as setting + * strictDeprecationMode to true. It can be set to + * {@link WarningsHandler#STRICT} if the client should fail if there are + * any warnings which is the same behavior as settings + * strictDeprecationMode to false. + *

    + * It can also be set to a custom implementation of + * {@linkplain WarningsHandler} to permit only certain warnings or to + * fail the request if the warnings returned don't + * exactly match some set. + * + * @param warningsHandler the {@link WarningsHandler} to be used + */ + public void setWarningsHandler(WarningsHandler warningsHandler) { + this.warningsHandler = warningsHandler; + } + } +} diff --git a/client/rest-http-client/src/main/java/org/opensearch/internal/httpclient/Response.java b/client/rest-http-client/src/main/java/org/opensearch/internal/httpclient/Response.java new file mode 100644 index 0000000000000..dccfb8d24a79b --- /dev/null +++ b/client/rest-http-client/src/main/java/org/opensearch/internal/httpclient/Response.java @@ -0,0 +1,117 @@ +/* + * SPDX-License-Identifier: Apache-2.0 + * + * The OpenSearch Contributors require contributions made to + * this file be licensed under the Apache-2.0 license or a + * compatible open source license. + */ + +package org.opensearch.internal.httpclient; + +import java.io.InputStream; +import java.net.http.HttpHeaders; +import java.net.http.HttpResponse; +import java.nio.ByteBuffer; +import java.util.List; +import java.util.concurrent.Flow; + +/** + * Holds an OpenSearch response. It wraps the {@link HttpResponse} returned and associates it with + * its corresponding {@link RequestLine} and {@link HttpHost}. + * Note: This is an experimental API. + */ +public sealed interface Response permits CompressedResponse, NonCompressedResponse { + /** + * Create response from streaming conversation + * @param requestLine request line + * @param host host + * @param response underlying HTTP response + * @return new response instance + */ + static Response fromStreaming(RequestLine requestLine, HttpHost host, HttpResponse>> response) { + return new NonCompressedResponse(requestLine, host, response, List.of() /* streaming body could be very large */); + } + + /** + * Create response from non-streaming conversation + * @param requestLine request line + * @param host host + * @param response underlying HTTP response + * @return new response instance + */ + static Response from(RequestLine requestLine, HttpHost host, HttpResponse> response) { + final boolean compressed = response.headers() + .firstValue("Content-Encoding") + .filter("gzip"::equalsIgnoreCase) + .map(h -> true) + .orElse(false); + if (compressed == false) { + return new NonCompressedResponse(requestLine, host, response, response.body()); + } else { + return new CompressedResponse(requestLine, host, response, response.body()); + } + } + + /** + * Returns the request line that generated this response + */ + RequestLine requestLine(); + + /** + * Returns the node that returned this response + */ + HttpHost host(); + + /** + * Returns the status line of the current response + */ + default StatusLine statusLine() { + return new StatusLine(httpResponse()); + } + + /** + * Returns all the response headers + */ + default HttpHeaders headers() { + return httpResponse().headers(); + } + + /** + * Returns the response body available, null otherwise + * @see InputStream + */ + List entity(); + + /** + * Returns the value of the first header with a specified name of this message. + * If there is more than one matching header in the message the first element is returned. + * If there is no matching header in the message null is returned. + * + * @param name header name + */ + default String header(String name) { + return headers().firstValue(name).orElse(null); + } + + /** + * Returns a list of all warning headers returned in the response. + */ + default List warnings() { + return ResponseWarningsExtractor.getWarnings(httpResponse()); + } + + /** + * Returns true if there is at least one warning header returned in the + * response. + */ + default boolean hasWarnings() { + List warnings = headers().allValues("Warning"); + return warnings != null && warnings.size() > 0; + } + + /** + * Returns underlying HTTP response instance + * @return + */ + HttpResponse httpResponse(); +} diff --git a/client/rest-http-client/src/main/java/org/opensearch/internal/httpclient/ResponseException.java b/client/rest-http-client/src/main/java/org/opensearch/internal/httpclient/ResponseException.java new file mode 100644 index 0000000000000..05d95bd9014ab --- /dev/null +++ b/client/rest-http-client/src/main/java/org/opensearch/internal/httpclient/ResponseException.java @@ -0,0 +1,61 @@ +/* + * SPDX-License-Identifier: Apache-2.0 + * + * The OpenSearch Contributors require contributions made to + * this file be licensed under the Apache-2.0 license or a + * compatible open source license. + */ + +package org.opensearch.internal.httpclient; + +import java.io.IOException; +import java.nio.ByteBuffer; +import java.util.List; +import java.util.Locale; + +/** + * Exception thrown when an opensearch node responds to a request with a status code that indicates an error. + * Holds the response that was returned. + */ +public final class ResponseException extends IOException { + private static final long serialVersionUID = 1L; + private final Response response; + + /** + * Creates a ResponseException containing the given {@code Response}. + * + * @param response The error response. + */ + public ResponseException(Response response) throws IOException { + super(buildMessage(response)); + this.response = response; + } + + static String buildMessage(Response response) throws IOException { + String message = String.format( + Locale.ROOT, + "method [%s], host [%s], URI [%s], status line [%s]", + response.requestLine().method(), + response.host(), + response.requestLine().uri(), + response.statusLine().toString() + ); + + if (response.hasWarnings()) { + message += "\n" + "Warnings: " + response.warnings(); + } + + List entity = response.entity(); + if (entity != null) { + message += "\n" + BodyUtils.getBodyAsString(response); + } + return message; + } + + /** + * Returns the {@link Response} that caused this exception to be thrown. + */ + public Response getResponse() { + return response; + } +} diff --git a/client/rest-http-client/src/main/java/org/opensearch/internal/httpclient/ResponseListener.java b/client/rest-http-client/src/main/java/org/opensearch/internal/httpclient/ResponseListener.java new file mode 100644 index 0000000000000..6775322ff5e71 --- /dev/null +++ b/client/rest-http-client/src/main/java/org/opensearch/internal/httpclient/ResponseListener.java @@ -0,0 +1,61 @@ +/* + * SPDX-License-Identifier: Apache-2.0 + * + * The OpenSearch Contributors require contributions made to + * this file be licensed under the Apache-2.0 license or a + * compatible open source license. + */ + +/* + * Licensed to Elasticsearch under one or more contributor + * license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright + * ownership. Elasticsearch 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. + */ + +/* + * Modifications Copyright OpenSearch Contributors. See + * GitHub history for details. + */ + +package org.opensearch.internal.httpclient; + +/** + * Listener to be provided when calling async performRequest methods provided by {@link RestHttpClient}. + * Those methods that do accept a listener will return immediately, execute asynchronously, and notify + * the listener whenever the request yielded a response, or failed with an exception. + * + *

    + * Note that it is not safe to call {@link RestHttpClient#close()} from either of these + * callbacks. + */ +public interface ResponseListener { + + /** + * Method invoked if the request yielded a successful response. + * + * @param response success response + */ + void onSuccess(Response response); + + /** + * Method invoked if the request failed. There are two main categories of failures: connection failures (usually + * {@link java.io.IOException}s, or responses that were treated as errors based on their error response code + * ({@link ResponseException}s). + * + * @param exception the failure exception + */ + void onFailure(Exception exception); +} diff --git a/client/rest-http-client/src/main/java/org/opensearch/internal/httpclient/ResponseWarningsExtractor.java b/client/rest-http-client/src/main/java/org/opensearch/internal/httpclient/ResponseWarningsExtractor.java new file mode 100644 index 0000000000000..296538333fea3 --- /dev/null +++ b/client/rest-http-client/src/main/java/org/opensearch/internal/httpclient/ResponseWarningsExtractor.java @@ -0,0 +1,96 @@ +/* + * SPDX-License-Identifier: Apache-2.0 + * + * The OpenSearch Contributors require contributions made to + * this file be licensed under the Apache-2.0 license or a + * compatible open source license. + */ + +package org.opensearch.internal.httpclient; + +import java.net.http.HttpResponse; +import java.util.ArrayList; +import java.util.List; +import java.util.regex.Matcher; +import java.util.regex.Pattern; + +final class ResponseWarningsExtractor { + + /** + * Optimized regular expression to test if a string matches the RFC 1123 date + * format (with quotes and leading space). Start/end of line characters and + * atomic groups are used to prevent backtracking. + */ + private static final Pattern WARNING_HEADER_DATE_PATTERN = Pattern.compile("^ " + // start of line, leading space + // quoted RFC 1123 date format + "\"" + // opening quote + "(?>Mon|Tue|Wed|Thu|Fri|Sat|Sun), " + // day of week, atomic group to prevent backtracking + "\\d{2} " + // 2-digit day + "(?>Jan|Feb|Mar|Apr|May|Jun|Jul|Aug|Sep|Oct|Nov|Dec) " + // month, atomic group to prevent backtracking + "\\d{4} " + // 4-digit year + "\\d{2}:\\d{2}:\\d{2} " + // (two-digit hour):(two-digit minute):(two-digit second) + "GMT" + // GMT + "\"$"); // closing quote (optional, since an older version can still send a warn-date), end of line + + /** + * Length of RFC 1123 format (with quotes and leading space), used in + * matchWarningHeaderPatternByPrefix(String). + */ + private static final int WARNING_HEADER_DATE_LENGTH = 0 + 1 + 1 + 3 + 1 + 1 + 2 + 1 + 3 + 1 + 4 + 1 + 2 + 1 + 2 + 1 + 2 + 1 + 3 + 1; + + private ResponseWarningsExtractor() {} + + /** + * Returns a list of all warning headers returned in the response. + * @param response HTTP response + */ + static List getWarnings(final HttpResponse response) { + List warnings = new ArrayList<>(); + for (String warning : response.headers().allValues("Warning")) { + if (matchWarningHeaderPatternByPrefix(warning)) { + warnings.add(extractWarningValueFromWarningHeader(warning)); + } else { + warnings.add(warning); + } + } + return warnings; + } + + /** + * Tests if a string matches the RFC 7234 specification for warning headers. + * This assumes that the warn code is always 299 and the warn agent is always + * OpenSearch. + * + * @param s the value of a warning header formatted according to RFC 7234 + * @return {@code true} if the input string matches the specification + */ + private static boolean matchWarningHeaderPatternByPrefix(final String s) { + return s.startsWith("299 OpenSearch-"); + } + + /** + * Refer to org.opensearch.common.logging.DeprecationLogger + */ + private static String extractWarningValueFromWarningHeader(final String s) { + String warningHeader = s; + + /* + * The following block tests for the existence of a RFC 1123 date in the warning header. If the date exists, it is removed for + * extractWarningValueFromWarningHeader(String) to work properly (as it does not handle dates). + */ + if (s.length() > WARNING_HEADER_DATE_LENGTH) { + final String possibleDateString = s.substring(s.length() - WARNING_HEADER_DATE_LENGTH); + final Matcher matcher = WARNING_HEADER_DATE_PATTERN.matcher(possibleDateString); + + if (matcher.matches()) { + warningHeader = warningHeader.substring(0, s.length() - WARNING_HEADER_DATE_LENGTH); + } + } + + final int firstQuote = warningHeader.indexOf('\"'); + final int lastQuote = warningHeader.length() - 1; + final String warningValue = warningHeader.substring(firstQuote + 1, lastQuote); + return warningValue; + } + +} diff --git a/client/rest-http-client/src/main/java/org/opensearch/internal/httpclient/RestHttpClient.java b/client/rest-http-client/src/main/java/org/opensearch/internal/httpclient/RestHttpClient.java new file mode 100644 index 0000000000000..446b537d68826 --- /dev/null +++ b/client/rest-http-client/src/main/java/org/opensearch/internal/httpclient/RestHttpClient.java @@ -0,0 +1,1096 @@ +/* + * SPDX-License-Identifier: Apache-2.0 + * + * The OpenSearch Contributors require contributions made to + * this file be licensed under the Apache-2.0 license or a + * compatible open source license. + */ + +package org.opensearch.internal.httpclient; + +import org.apache.commons.logging.Log; +import org.apache.commons.logging.LogFactory; + +import javax.net.ssl.SSLHandshakeException; + +import java.io.Closeable; +import java.io.IOException; +import java.net.ConnectException; +import java.net.SocketTimeoutException; +import java.net.URI; +import java.net.URISyntaxException; +import java.net.URLEncoder; +import java.net.http.HttpClient; +import java.net.http.HttpRequest; +import java.net.http.HttpRequest.BodyPublisher; +import java.net.http.HttpRequest.BodyPublishers; +import java.net.http.HttpResponse; +import java.net.http.HttpResponse.BodyHandler; +import java.net.http.HttpResponse.BodyHandlers; +import java.net.http.HttpTimeoutException; +import java.nio.ByteBuffer; +import java.nio.channels.ClosedChannelException; +import java.nio.charset.StandardCharsets; +import java.time.Duration; +import java.util.ArrayList; +import java.util.Arrays; +import java.util.Base64; +import java.util.Collection; +import java.util.Collections; +import java.util.HashMap; +import java.util.HashSet; +import java.util.Iterator; +import java.util.LinkedHashMap; +import java.util.List; +import java.util.Map; +import java.util.Objects; +import java.util.Set; +import java.util.concurrent.CancellationException; +import java.util.concurrent.CompletableFuture; +import java.util.concurrent.CompletionException; +import java.util.concurrent.ConcurrentHashMap; +import java.util.concurrent.ConcurrentMap; +import java.util.concurrent.ExecutionException; +import java.util.concurrent.Flow; +import java.util.concurrent.Future; +import java.util.concurrent.atomic.AtomicInteger; +import java.util.function.Function; +import java.util.stream.Collectors; + +import org.reactivestreams.Publisher; +import reactor.adapter.JdkFlowAdapter; +import reactor.core.publisher.Flux; +import reactor.core.publisher.Mono; + +import static java.nio.charset.StandardCharsets.UTF_8; +import static java.util.Collections.singletonList; + +/** + * Client that connects to an OpenSearch cluster through HTTP. + *

    + * Must be created using {@link RestHttpClientBuilder}, which allows to set all the different options or just rely on defaults. + * The hosts that are part of the cluster need to be provided at creation time, but can also be replaced later + * by calling {@link #setNodes(Collection)}. + *

    + * The method {@link #performRequest(Request)} allows to send a request to the cluster. When + * sending a request, a host gets selected out of the provided ones in a round-robin fashion. Failing hosts are marked dead and + * retried after a certain amount of time (minimum 1 minute, maximum 30 minutes), depending on how many times they previously + * failed (the more failures, the later they will be retried). In case of failures all of the alive nodes (or dead nodes that + * deserve a retry) are retried until one responds or none of them does, in which case an {@link IOException} will be thrown. + *

    + * Requests can be either synchronous or asynchronous. The asynchronous variants all end with {@code Async}. + *

    + * Requests can be traced by enabling trace logging for "tracer". The trace logger outputs requests and responses in curl format. + * + * Note: This is an experimental API. + */ +public class RestHttpClient implements Closeable { + + private static final Log logger = LogFactory.getLog(RestHttpClient.class); + + private final HttpClient client; + final Map> defaultHeaders; + private final String pathPrefix; + private final AtomicInteger lastNodeIndex = new AtomicInteger(0); + private final ConcurrentMap denylist = new ConcurrentHashMap<>(); + private final FailureListener failureListener; + private final NodeSelector nodeSelector; + private volatile List nodes; + private final WarningsHandler warningsHandler; + private final boolean compressionEnabled; + + RestHttpClient( + HttpClient client, + Map> defaultHeaders, + List nodes, + String pathPrefix, + FailureListener failureListener, + NodeSelector nodeSelector, + boolean strictDeprecationMode, + boolean compressionEnabled + ) { + this.client = client; + this.defaultHeaders = Collections.unmodifiableMap(defaultHeaders); + this.failureListener = failureListener; + this.pathPrefix = pathPrefix; + this.nodeSelector = nodeSelector; + this.warningsHandler = strictDeprecationMode ? WarningsHandler.STRICT : WarningsHandler.PERMISSIVE; + this.compressionEnabled = compressionEnabled; + setNodes(nodes); + } + + /** + * Returns a new {@link RestHttpClientBuilder} to help with {@link RestHttpClient} creation. + * Creates a new builder instance and sets the nodes that the client will send requests to. + * + * @param cloudId a valid elastic cloud cloudId that will route to a cluster. The cloudId is located in + * the user console https://cloud.elastic.co and will resemble a string like the following + * optionalHumanReadableName:dXMtZWFzdC0xLmF3cy5mb3VuZC5pbyRlbGFzdGljc2VhcmNoJGtpYmFuYQ== + */ + public static RestHttpClientBuilder builder(String cloudId) { + // there is an optional first portion of the cloudId that is a human readable string, but it is not used. + if (cloudId.contains(":")) { + if (cloudId.indexOf(":") == cloudId.length() - 1) { + throw new IllegalStateException("cloudId " + cloudId + " must begin with a human readable identifier followed by a colon"); + } + cloudId = cloudId.substring(cloudId.indexOf(":") + 1); + } + + String decoded = new String(Base64.getDecoder().decode(cloudId), UTF_8); + // once decoded the parts are separated by a $ character. + // they are respectively domain name and optional port, opensearch id, opensearch-dashboards id + String[] decodedParts = decoded.split("\\$"); + if (decodedParts.length != 3) { + throw new IllegalStateException("cloudId " + cloudId + " did not decode to a cluster identifier correctly"); + } + + // domain name and optional port + String[] domainAndMaybePort = decodedParts[0].split(":", 2); + String domain = domainAndMaybePort[0]; + int port; + + if (domainAndMaybePort.length == 2) { + try { + port = Integer.parseInt(domainAndMaybePort[1]); + } catch (NumberFormatException nfe) { + throw new IllegalStateException("cloudId " + cloudId + " does not contain a valid port number"); + } + } else { + port = 443; + } + + String url = decodedParts[1] + "." + domain; + return builder(new HttpHost("https", url, port)); + } + + /** + * Returns a new {@link RestHttpClientBuilder} to help with {@link RestHttpClient} creation. + * Creates a new builder instance and sets the hosts that the client will send requests to. + *

    + * Prefer this to {@link #builder(HttpHost...)} if you have metadata up front about the nodes. + * If you don't either one is fine. + * + * @param nodes The nodes that the client will send requests to. + */ + public static RestHttpClientBuilder builder(Node... nodes) { + return new RestHttpClientBuilder(nodes == null ? null : Arrays.asList(nodes)); + } + + /** + * Returns a new {@link RestHttpClientBuilder} to help with {@link RestHttpClient} creation. + * Creates a new builder instance and sets the nodes that the client will send requests to. + *

    + * You can use this if you do not have metadata up front about the nodes. If you do, prefer + * {@link #builder(Node...)}. + * @see Node#Node(HttpHost) + * + * @param hosts The hosts that the client will send requests to. + */ + public static RestHttpClientBuilder builder(HttpHost... hosts) { + if (hosts == null || hosts.length == 0) { + throw new IllegalArgumentException("hosts must not be null nor empty"); + } + List nodes = Arrays.stream(hosts).map(Node::new).collect(Collectors.toList()); + return new RestHttpClientBuilder(nodes); + } + + /** + * Replaces the nodes with which the client communicates. + * + * @param nodes the new nodes to communicate with. + */ + public synchronized void setNodes(Collection nodes) { + if (nodes == null || nodes.isEmpty()) { + throw new IllegalArgumentException("nodes must not be null or empty"); + } + + Map nodesByHost = new LinkedHashMap<>(); + for (Node node : nodes) { + Objects.requireNonNull(node, "node cannot be null"); + // TODO should we throw an IAE if we have two nodes with the same host? + nodesByHost.put(node.getHost(), node); + } + this.nodes = Collections.unmodifiableList(new ArrayList<>(nodesByHost.values())); + this.denylist.clear(); + } + + /** + * Get the list of nodes that the client knows about. The list is + * unmodifiable. + */ + public List getNodes() { + return nodes; + } + + /** + * check client running status + * @return client running status + */ + public boolean isRunning() { + return client.isTerminated() == false; + } + + /** + * Sends a streaming request to the OpenSearch cluster that the client points to and returns streaming response. This is an experimental API. + * @param request streaming request + * @return streaming response + * @throws IOException IOException + */ + public StreamingResponse streamRequest(StreamingRequest request) throws IOException { + final InternalStreamingRequest internalRequest = new InternalStreamingRequest(request); + return streamRequest(nextNodes(), internalRequest); + } + + /** + * Sends a request to the OpenSearch cluster that the client points to. + * Blocks until the request is completed and returns its response or fails + * by throwing an exception. Selects a host out of the provided ones in a + * round-robin fashion. Failing hosts are marked dead and retried after a + * certain amount of time (minimum 1 minute, maximum 30 minutes), depending + * on how many times they previously failed (the more failures, the later + * they will be retried). In case of failures all of the alive nodes (or + * dead nodes that deserve a retry) are retried until one responds or none + * of them does, in which case an {@link IOException} will be thrown. + *

    + * This method works by performing an asynchronous call and waiting + * for the result. If the asynchronous call throws an exception we wrap + * it and rethrow it so that the stack trace attached to the exception + * contains the call site. While we attempt to preserve the original + * exception this isn't always possible and likely haven't covered all of + * the cases. You can get the original exception from + * {@link Exception#getCause()}. + * + * @param request the request to perform + * @return the response returned by OpenSearch + * @throws IOException in case of a problem or the connection was aborted + * @throws ResponseException in case OpenSearch responded with a status code that indicated an error + */ + public Response performRequest(Request request) throws IOException { + InternalRequest internalRequest = new InternalRequest(request); + return performRequest(nextNodes(), internalRequest, null); + } + + private Response performRequest(final Iterator nodes, final InternalRequest request, Exception previousException) + throws IOException { + Node node = nodes.next(); + RequestContext> context = request.createContextForNextAttempt(node); + HttpResponse> httpResponse; + try { + httpResponse = client.send(context.requestProducer(), context.asyncResponseConsumer()); + } catch (Exception e) { + RequestLogger.logFailedRequest(logger, request.httpRequest, context.node(), e); + onFailure(context.node()); + Exception cause = extractAndWrapCause(e); + addSuppressedException(previousException, cause); + if (nodes.hasNext()) { + return performRequest(nodes, request, cause); + } + if (cause instanceof IOException) { + throw (IOException) cause; + } + if (cause instanceof RuntimeException) { + throw (RuntimeException) cause; + } + throw new IllegalStateException("unexpected exception type: must be either RuntimeException or IOException", cause); + } + ResponseOrResponseException responseOrResponseException = convertResponse(request, context.node(), httpResponse); + if (responseOrResponseException.responseException == null) { + return responseOrResponseException.response; + } + addSuppressedException(previousException, responseOrResponseException.responseException); + if (nodes.hasNext()) { + return performRequest(nodes, request, responseOrResponseException.responseException); + } + throw responseOrResponseException.responseException; + } + + private Publisher>>> streamRequest( + final Node node, + final Iterator nodes, + final InternalStreamingRequest request + ) throws IOException { + return request.cancellable.callIfNotCancelled(() -> { + final RequestContext>> context = request.createContextForNextAttempt(node); + final CompletableFuture>>> future = client.sendAsync( + context.requestProducer(), + context.asyncResponseConsumer() + ); + request.setCancellable(future); + + final Mono>>> publisher = Mono.fromCompletionStage(future).flatMap(response -> { + try { + final ResponseOrResponseException responseOrResponseException = convertResponse(request, node, response); + if (responseOrResponseException.responseException == null) { + return Mono.just(response); + } else { + if (nodes.hasNext()) { + return Mono.from(streamRequest(nodes.next(), nodes, request)); + } else { + return Mono.error(responseOrResponseException.responseException); + } + } + } catch (final Exception ex) { + return Mono.error(ex); + } + }); + + return publisher; + }); + } + + private StreamingResponse streamRequest(final Iterator nodes, final InternalStreamingRequest request) throws IOException { + return request.cancellable.callIfNotCancelled(() -> { + final Node node = nodes.next(); + return new StreamingResponse(new RequestLine(request.httpRequest.apply(node)), streamRequest(node, nodes, request)); + }); + } + + private ResponseOrResponseException convertResponse(InternalRequest request, Node node, HttpResponse> httpResponse) + throws IOException { + + final HttpRequest httpRequest = request.httpRequest.apply(node); + RequestLogger.logResponse(logger, httpRequest, node.getHost(), httpResponse); + int statusCode = httpResponse.statusCode(); + + Response response = Response.from(new RequestLine(httpRequest), node.getHost(), httpResponse); + if (isSuccessfulResponse(statusCode) || request.ignoreErrorCodes.contains(response.statusLine().statusCode())) { + onResponse(node); + if (request.warningsHandler.warningsShouldFailRequest(response.warnings())) { + throw new WarningFailureException(response); + } + return new ResponseOrResponseException(response); + } + ResponseException responseException = new ResponseException(response); + if (isRetryStatus(statusCode)) { + // mark host dead and retry against next one + onFailure(node); + return new ResponseOrResponseException(responseException); + } + // mark host alive and don't retry, as the error should be a request problem + onResponse(node); + throw responseException; + } + + private ResponseOrResponseException convertResponse( + InternalStreamingRequest request, + Node node, + HttpResponse>> httpResponse + ) throws IOException { + + // Streaming Response could accumulate a lot of data so we may not be able to fully consume it. + final HttpRequest httpRequest = request.httpRequest.apply(node); + final Response response = Response.fromStreaming(new RequestLine(httpRequest), node.getHost(), httpResponse); + + RequestLogger.logStreamingResponse(logger, request.httpRequest.apply(node), node.getHost(), httpResponse); + int statusCode = httpResponse.statusCode(); + + if (isSuccessfulResponse(statusCode) || request.ignoreErrorCodes.contains(response.statusLine().statusCode())) { + onResponse(node); + if (request.warningsHandler.warningsShouldFailRequest(response.warnings())) { + throw new WarningFailureException(response); + } + return new ResponseOrResponseException(response); + } + ResponseException responseException = new ResponseException(response); + if (isRetryStatus(statusCode)) { + // mark host dead and retry against next one + onFailure(node); + return new ResponseOrResponseException(responseException); + } + // mark host alive and don't retry, as the error should be a request problem + onResponse(node); + throw responseException; + } + + /** + * Sends a request to the OpenSearch cluster that the client points to. + * The request is executed asynchronously and the provided + * {@link ResponseListener} gets notified upon request completion or + * failure. Selects a host out of the provided ones in a round-robin + * fashion. Failing hosts are marked dead and retried after a certain + * amount of time (minimum 1 minute, maximum 30 minutes), depending on how + * many times they previously failed (the more failures, the later they + * will be retried). In case of failures all of the alive nodes (or dead + * nodes that deserve a retry) are retried until one responds or none of + * them does, in which case an {@link IOException} will be thrown. + * + * @param request the request to perform + * @param responseListener the {@link ResponseListener} to notify when the + * request is completed or fails + */ + public Cancellable performRequestAsync(Request request, ResponseListener responseListener) { + try { + FailureTrackingResponseListener failureTrackingResponseListener = new FailureTrackingResponseListener(responseListener); + InternalRequest internalRequest = new InternalRequest(request); + performRequestAsync(nextNodes(), internalRequest, failureTrackingResponseListener); + return internalRequest.cancellable; + } catch (Exception e) { + responseListener.onFailure(e); + return Cancellable.NO_OP; + } + } + + private void performRequestAsync( + final Iterator nodes, + final InternalRequest request, + final FailureTrackingResponseListener listener + ) { + request.cancellable.runIfNotCancelled(() -> { + final RequestContext> context = request.createContextForNextAttempt(nodes.next()); + CompletableFuture>> future = client.sendAsync( + context.requestProducer(), + context.asyncResponseConsumer() + ); + + request.setCancellable(future); + future.whenComplete((httpResponse, throwable) -> { + if (httpResponse != null) { + try { + ResponseOrResponseException responseOrResponseException = convertResponse(request, context.node(), httpResponse); + if (responseOrResponseException.responseException == null) { + listener.onSuccess(responseOrResponseException.response); + } else { + if (nodes.hasNext()) { + listener.trackFailure(responseOrResponseException.responseException); + performRequestAsync(nodes, request, listener); + } else { + listener.onDefinitiveFailure(responseOrResponseException.responseException); + } + } + } catch (Exception e) { + listener.onDefinitiveFailure(e); + } + } else if (throwable instanceof Exception failure) { + if (failure instanceof CancellationException) { + listener.onDefinitiveFailure(Cancellable.newCancellationException()); + } else { + Exception cause = failure; + if (failure instanceof CompletionException && failure.getCause() instanceof Exception ce) { + cause = ce; + } + try { + RequestLogger.logFailedRequest(logger, request.httpRequest, context.node(), failure); + onFailure(context.node()); + if (nodes.hasNext()) { + listener.trackFailure(failure); + performRequestAsync(nodes, request, listener); + } else { + listener.onDefinitiveFailure(cause); + } + } catch (Exception e) { + listener.onDefinitiveFailure(e); + } + } + } + }); + }); + + } + + /** + * Returns a non-empty {@link Iterator} of nodes to be used for a request + * that match the {@link NodeSelector}. + *

    + * If there are no living nodes that match the {@link NodeSelector} + * this will return the dead node that matches the {@link NodeSelector} + * that is closest to being revived. + * @throws IOException if no nodes are available + */ + private Iterator nextNodes() throws IOException { + List nodes = this.nodes; + Iterable hosts = selectNodes(nodes, denylist, lastNodeIndex, nodeSelector); + return hosts.iterator(); + } + + /** + * Select nodes to try and sorts them so that the first one will be tried initially, then the following ones + * if the previous attempt failed and so on. Package private for testing. + */ + static Iterable selectNodes( + List nodes, + Map denylist, + AtomicInteger lastNodeIndex, + NodeSelector nodeSelector + ) throws IOException { + /* + * Sort the nodes into living and dead lists. + */ + List livingNodes = new ArrayList<>(Math.max(0, nodes.size() - denylist.size())); + List deadNodes = new ArrayList<>(denylist.size()); + for (Node node : nodes) { + DeadHostState deadness = denylist.get(node.getHost()); + if (deadness == null || deadness.shallBeRetried()) { + livingNodes.add(node); + } else { + deadNodes.add(new DeadNode(node, deadness)); + } + } + + if (false == livingNodes.isEmpty()) { + /* + * Normal state: there is at least one living node. If the + * selector is ok with any over the living nodes then use them + * for the request. + */ + List selectedLivingNodes = new ArrayList<>(livingNodes); + nodeSelector.select(selectedLivingNodes); + if (false == selectedLivingNodes.isEmpty()) { + /* + * Rotate the list using a global counter as the distance so subsequent + * requests will try the nodes in a different order. + */ + Collections.rotate(selectedLivingNodes, lastNodeIndex.getAndIncrement()); + return selectedLivingNodes; + } + } + + /* + * Last resort: there are no good nodes to use, either because + * the selector rejected all the living nodes or because there aren't + * any living ones. Either way, we want to revive a single dead node + * that the NodeSelectors are OK with. We do this by passing the dead + * nodes through the NodeSelector so it can have its say in which nodes + * are ok. If the selector is ok with any of the nodes then we will take + * the one in the list that has the lowest revival time and try it. + */ + if (false == deadNodes.isEmpty()) { + final List selectedDeadNodes = new ArrayList<>(deadNodes); + /* + * We'd like NodeSelectors to remove items directly from deadNodes + * so we can find the minimum after it is filtered without having + * to compare many things. This saves us a sort on the unfiltered + * list. + */ + nodeSelector.select(() -> new DeadNodeIteratorAdapter(selectedDeadNodes.iterator())); + if (false == selectedDeadNodes.isEmpty()) { + return singletonList(Collections.min(selectedDeadNodes).node); + } + } + throw new IOException( + "NodeSelector [" + nodeSelector + "] rejected all nodes, " + "living " + livingNodes + " and dead " + deadNodes + ); + } + + /** + * Called after each successful request call. + * Receives as an argument the host that was used for the successful request. + */ + private void onResponse(Node node) { + DeadHostState removedHost = this.denylist.remove(node.getHost()); + if (logger.isDebugEnabled() && removedHost != null) { + logger.debug("removed [" + node + "] from denylist"); + } + } + + /** + * Called after each failed attempt. + * Receives as an argument the host that was used for the failed attempt. + */ + private void onFailure(Node node) { + while (true) { + DeadHostState previousDeadHostState = denylist.putIfAbsent( + node.getHost(), + new DeadHostState(DeadHostState.DEFAULT_TIME_SUPPLIER) + ); + if (previousDeadHostState == null) { + if (logger.isDebugEnabled()) { + logger.debug("added [" + node + "] to denylist"); + } + break; + } + if (denylist.replace(node.getHost(), previousDeadHostState, new DeadHostState(previousDeadHostState))) { + if (logger.isDebugEnabled()) { + logger.debug("updated [" + node + "] already in denylist"); + } + break; + } + } + failureListener.onFailure(node); + } + + /** + * Close the underlying {@link HttpClient} instance + */ + @Override + public void close() throws IOException { + client.shutdownNow(); + client.close(); + } + + private static boolean isSuccessfulResponse(int statusCode) { + return statusCode < 300; + } + + private static boolean isRetryStatus(int statusCode) { + switch (statusCode) { + case 502: + case 503: + case 504: + return true; + } + return false; + } + + private static void addSuppressedException(Exception suppressedException, Exception currentException) { + if (suppressedException != null) { + currentException.addSuppressed(suppressedException); + } + } + + private HttpRequest.Builder createHttpRequest( + Node node, + String method, + URI uri, + BodyPublisher body, + Duration timeout, + boolean compressed + ) { + return HttpRequest.newBuilder() + .uri(URI.create(node.getHost().toString()).resolve(uri)) + .timeout(timeout) + .method( + method, + (body == null) ? BodyPublishers.noBody() + : compressed == false ? body + : BodyPublishers.fromPublisher( + JdkFlowAdapter.publisherToFlowPublisher( + JdkFlowAdapter.flowPublisherToFlux(body).buffer().map(BodyUtils::compress).flatMap(Flux::fromIterable) + ) + ) + ); + } + + private HttpRequest.Builder createStreamingHttpRequest( + Node node, + String method, + URI uri, + Publisher body, + Duration timeout, + boolean compressed + ) { + return HttpRequest.newBuilder() + .uri(URI.create(node.getHost().toString()).resolve(uri)) + .timeout(timeout) + .method( + method, + (body == null) ? BodyPublishers.noBody() + : compressed == false ? BodyPublishers.fromPublisher(JdkFlowAdapter.publisherToFlowPublisher(body)) + : BodyPublishers.fromPublisher(JdkFlowAdapter.publisherToFlowPublisher(Flux.from(body).map(BodyUtils::compress))) + ); + } + + static URI buildUri(String pathPrefix, String path, Map params) { + Objects.requireNonNull(path, "path must not be null"); + try { + String fullPath; + if (pathPrefix != null && pathPrefix.isEmpty() == false) { + if (pathPrefix.endsWith("/") && path.startsWith("/")) { + fullPath = pathPrefix.substring(0, pathPrefix.length() - 1) + path; + } else if (pathPrefix.endsWith("/") || path.startsWith("/")) { + fullPath = pathPrefix + path; + } else { + fullPath = pathPrefix + "/" + path; + } + } else { + fullPath = path; + } + + final String additionalQuery = params.entrySet().stream().map(e -> { + if (e.getValue() == null) { + return e.getKey(); + } else { + return e.getKey() + "=" + URLEncoder.encode(e.getValue(), StandardCharsets.UTF_8); + } + }).collect(Collectors.joining("&")); + final URI uri = URI.create(fullPath); + + String newQuery = uri.getQuery(); + if (newQuery == null) { + newQuery = additionalQuery; + } else { + newQuery += "&" + additionalQuery; + } + + return new URI(uri.getScheme(), uri.getAuthority(), uri.getPath(), newQuery, uri.getFragment()); + } catch (URISyntaxException e) { + throw new IllegalArgumentException(e.getMessage(), e); + } + } + + /** + * Listener used in any async call to wrap the provided user listener (or SyncResponseListener in sync calls). + * Allows to track potential failures coming from the different retry attempts and returning to the original listener + * only when we got a response (successful or not to be retried) or there are no hosts to retry against. + */ + static class FailureTrackingResponseListener { + private final ResponseListener responseListener; + private volatile Exception exception; + + FailureTrackingResponseListener(ResponseListener responseListener) { + this.responseListener = responseListener; + } + + /** + * Notifies the caller of a response through the wrapped listener + */ + void onSuccess(Response response) { + responseListener.onSuccess(response); + } + + /** + * Tracks one last definitive failure and returns to the caller by notifying the wrapped listener + */ + void onDefinitiveFailure(Exception exception) { + trackFailure(exception); + responseListener.onFailure(this.exception); + } + + /** + * Tracks an exception, which caused a retry hence we should not return yet to the caller + */ + void trackFailure(Exception exception) { + addSuppressedException(this.exception, exception); + this.exception = exception; + } + } + + /** + * Listener that allows to be notified whenever a failure happens. Useful when sniffing is enabled, so that we can sniff on failure. + * The default implementation is a no-op. + */ + public static class FailureListener { + /** + * Create a {@link FailureListener} instance. + */ + public FailureListener() {} + + /** + * Notifies that the node provided as argument has just failed. + * + * @param node The node which has failed. + */ + public void onFailure(Node node) {} + } + + /** + * Contains a reference to a denylisted node and the time until it is + * revived. We use this so we can do a single pass over the denylist. + */ + private static class DeadNode implements Comparable { + final Node node; + final DeadHostState deadness; + + DeadNode(Node node, DeadHostState deadness) { + this.node = node; + this.deadness = deadness; + } + + @Override + public String toString() { + return node.toString(); + } + + @Override + public int compareTo(DeadNode rhs) { + return deadness.compareTo(rhs.deadness); + } + } + + /** + * Adapts an Iterator<DeadNodeAndRevival> into an + * Iterator<Node>. + */ + private static class DeadNodeIteratorAdapter implements Iterator { + private final Iterator itr; + + private DeadNodeIteratorAdapter(Iterator itr) { + this.itr = itr; + } + + @Override + public boolean hasNext() { + return itr.hasNext(); + } + + @Override + public Node next() { + return itr.next().node; + } + + @Override + public void remove() { + itr.remove(); + } + } + + private class InternalStreamingRequest { + private final Set ignoreErrorCodes; + private final Function httpRequest; + private final WarningsHandler warningsHandler; + private volatile Cancellable cancellable = Cancellable.fromFuture(new CompletableFuture<>()); + + InternalStreamingRequest(StreamingRequest request) { + Map params = new HashMap<>(request.parameters()); + // ignore is a special parameter supported by the clients, shouldn't be sent to es + String ignoreString = params.remove("ignore"); + this.ignoreErrorCodes = getIgnoreErrorCodes(ignoreString, request.method()); + + this.httpRequest = node -> { + URI uri = buildUri(pathPrefix, request.endpoint(), params); + HttpRequest.Builder builder = createStreamingHttpRequest( + node, + request.method(), + uri, + request.body(), + request.options().getTimeout(), + compressionEnabled + ); + setHeaders(builder, request.options().getHeaders()); + return builder.build(); + }; + + this.warningsHandler = request.options().getWarningsHandler() == null + ? RestHttpClient.this.warningsHandler + : request.options().getWarningsHandler(); + } + + private void setHeaders(HttpRequest.Builder httpRequest, Map> requestHeaders) { + // request headers override default headers, so we don't add default headers if they exist as request headers + final Set requestNames = new HashSet<>(requestHeaders.size()); + for (Map.Entry> requestHeader : requestHeaders.entrySet()) { + requestHeader.getValue().forEach(v -> httpRequest.header(requestHeader.getKey(), v)); + requestNames.add(requestHeader.getKey()); + } + for (Map.Entry> defaultHeader : defaultHeaders.entrySet()) { + if (requestNames.contains(defaultHeader.getKey()) == false) { + defaultHeader.getValue().forEach(v -> httpRequest.header(defaultHeader.getKey(), v)); + } + } + if (compressionEnabled) { + httpRequest.header("Content-Encoding", "gzip"); + httpRequest.header("Accept-Encoding", "gzip"); + } + } + + private void setCancellable(Future f) { + cancellable = Cancellable.fromFuture(f); + } + + RequestContext>> createContextForNextAttempt(Node node) { + return new ReactiveRequestContext(this, node); + } + } + + private class InternalRequest { + private final Set ignoreErrorCodes; + private final Function httpRequest; + private final WarningsHandler warningsHandler; + private volatile Cancellable cancellable = Cancellable.fromFuture(new CompletableFuture<>()); + + InternalRequest(Request request) { + Map params = new HashMap<>(request.parameters()); + // ignore is a special parameter supported by the clients, shouldn't be sent to es + String ignoreString = params.remove("ignore"); + this.ignoreErrorCodes = getIgnoreErrorCodes(ignoreString, request.method()); + this.httpRequest = node -> { + URI uri = buildUri(pathPrefix, request.endpoint(), params); + final HttpRequest.Builder builder = createHttpRequest( + node, + request.method(), + uri, + request.entity(), + request.options().getTimeout(), + compressionEnabled + ); + setHeaders(builder, request.options().getHeaders()); + return builder.build(); + }; + this.warningsHandler = request.options().getWarningsHandler() == null + ? RestHttpClient.this.warningsHandler + : request.options().getWarningsHandler(); + } + + private void setCancellable(Future f) { + cancellable = Cancellable.fromFuture(f); + } + + private void setHeaders(HttpRequest.Builder httpRequest, Map> requestHeaders) { + // request headers override default headers, so we don't add default headers if they exist as request headers + final Set requestNames = new HashSet<>(requestHeaders.size()); + + for (Map.Entry> requestHeader : requestHeaders.entrySet()) { + requestHeader.getValue().forEach(v -> httpRequest.header(requestHeader.getKey(), v)); + requestNames.add(requestHeader.getKey()); + } + for (Map.Entry> defaultHeader : defaultHeaders.entrySet()) { + if (requestNames.contains(defaultHeader.getKey()) == false) { + defaultHeader.getValue().forEach(v -> httpRequest.header(defaultHeader.getKey(), v)); + } + } + if (compressionEnabled) { + httpRequest.header("Content-Encoding", "gzip"); + httpRequest.header("Accept-Encoding", "gzip"); + } + } + + RequestContext> createContextForNextAttempt(Node node) { + return new AsyncRequestContext(this, node); + } + } + + private interface RequestContext { + Node node(); + + HttpRequest requestProducer(); + + BodyHandler asyncResponseConsumer(); + } + + private static class ReactiveRequestContext implements RequestContext>> { + private final Node node; + private final HttpRequest requestProducer; + private final BodyHandler>> asyncResponseConsumer; + + ReactiveRequestContext(InternalStreamingRequest request, Node node) { + this.node = node; + // we stream the request body if the entity allows for it + this.requestProducer = request.httpRequest.apply(node); + this.asyncResponseConsumer = BodyHandlers.ofPublisher(); + } + + @Override + public BodyHandler>> asyncResponseConsumer() { + return asyncResponseConsumer; + } + + @Override + public Node node() { + return node; + } + + @Override + public HttpRequest requestProducer() { + return requestProducer; + } + } + + private static class AsyncRequestContext implements RequestContext> { + private final Node node; + private final HttpRequest requestProducer; + private final BodyHandler> asyncResponseConsumer; + + AsyncRequestContext(InternalRequest request, Node node) { + this.node = node; + this.requestProducer = request.httpRequest.apply(node); + this.asyncResponseConsumer = BodyHandlers.fromSubscriber(new AsyncResponseProducer(), AsyncResponseProducer::getResult); + } + + @Override + public BodyHandler> asyncResponseConsumer() { + return asyncResponseConsumer; + } + + @Override + public Node node() { + return node; + } + + @Override + public HttpRequest requestProducer() { + return requestProducer; + } + } + + private static Set getIgnoreErrorCodes(String ignoreString, String requestMethod) { + Set ignoreErrorCodes; + if (ignoreString == null) { + if ("HEAD".equalsIgnoreCase(requestMethod)) { + // 404 never causes error if returned for a HEAD request + ignoreErrorCodes = Collections.singleton(404); + } else { + ignoreErrorCodes = Collections.emptySet(); + } + } else { + String[] ignoresArray = ignoreString.split(","); + ignoreErrorCodes = new HashSet<>(); + if ("HEAD".equalsIgnoreCase(requestMethod)) { + // 404 never causes error if returned for a HEAD request + ignoreErrorCodes.add(404); + } + for (String ignoreCode : ignoresArray) { + try { + ignoreErrorCodes.add(Integer.valueOf(ignoreCode)); + } catch (NumberFormatException e) { + throw new IllegalArgumentException("ignore value should be a number, found [" + ignoreString + "] instead", e); + } + } + } + return ignoreErrorCodes; + } + + private static class ResponseOrResponseException { + private final Response response; + private final ResponseException responseException; + + ResponseOrResponseException(Response response) { + this.response = Objects.requireNonNull(response); + this.responseException = null; + } + + ResponseOrResponseException(ResponseException responseException) { + this.responseException = Objects.requireNonNull(responseException); + this.response = null; + } + } + + /** + * Wrap the exception so the caller's signature shows up in the stack trace, taking care to copy the original type and message + * where possible so async and sync code don't have to check different exceptions. + */ + private static Exception extractAndWrapCause(Exception exception) { + if (exception instanceof InterruptedException) { + throw new RuntimeException("thread waiting for the response was interrupted", exception); + } + if (exception instanceof ExecutionException || exception instanceof CompletionException) { + Throwable t = exception.getCause() == null ? exception : exception.getCause(); + if (t instanceof Error) { + throw (Error) t; + } + exception = (Exception) t; + } + if (exception instanceof HttpTimeoutException) { + HttpTimeoutException e = new HttpTimeoutException(exception.getMessage()); + e.initCause(exception); + return e; + } + if (exception instanceof ClosedChannelException) { + ClosedChannelException e = new ClosedChannelException(); + e.initCause(exception); + return e; + } + if (exception instanceof SocketTimeoutException) { + SocketTimeoutException e = new SocketTimeoutException(exception.getMessage()); + e.initCause(exception); + return e; + } + if (exception instanceof SSLHandshakeException) { + SSLHandshakeException e = new SSLHandshakeException( + exception.getMessage() + "\nSee https://opensearch.org/docs/latest/clients/java-rest-high-level/ for troubleshooting." + ); + e.initCause(exception); + return e; + } + if (exception instanceof ConnectException) { + ConnectException e = new ConnectException(exception.getMessage()); + e.initCause(exception); + return e; + } + if (exception instanceof IOException) { + return new IOException(exception.getMessage(), exception); + } + if (exception instanceof RuntimeException) { + return new RuntimeException(exception.getMessage(), exception); + } + return new RuntimeException("error while performing request", exception); + } +} diff --git a/client/rest-http-client/src/main/java/org/opensearch/internal/httpclient/RestHttpClientBuilder.java b/client/rest-http-client/src/main/java/org/opensearch/internal/httpclient/RestHttpClientBuilder.java new file mode 100644 index 0000000000000..367e2efc05da4 --- /dev/null +++ b/client/rest-http-client/src/main/java/org/opensearch/internal/httpclient/RestHttpClientBuilder.java @@ -0,0 +1,248 @@ +/* + * SPDX-License-Identifier: Apache-2.0 + * + * The OpenSearch Contributors require contributions made to + * this file be licensed under the Apache-2.0 license or a + * compatible open source license. + */ + +/* + * Modifications Copyright OpenSearch Contributors. See + * GitHub history for details. + */ + +package org.opensearch.internal.httpclient; + +import javax.net.ssl.SSLContext; + +import java.net.Authenticator; +import java.net.http.HttpClient; +import java.security.AccessController; +import java.security.NoSuchAlgorithmException; +import java.security.PrivilegedAction; +import java.time.Duration; +import java.util.Collections; +import java.util.List; +import java.util.Map; +import java.util.Objects; + +/** + * Helps creating a new {@link RestHttpClient}. Allows to set the most common http client configuration options when internally + * creating the underlying {@link HttpClient}. Also allows to provide an externally created + * {@link HttpClient} in case additional customization is needed. + * + * Note: This is an experimental API. + */ +public final class RestHttpClientBuilder { + /** + * The default connection timeout in milliseconds. + */ + public static final int DEFAULT_CONNECT_TIMEOUT_MILLIS = 1000; + + /** + * The default response timeout in milliseconds. + */ + public static final int DEFAULT_RESPONSE_TIMEOUT_MILLIS = 30000; + + private final List nodes; + private Map> defaultHeaders = Map.of(); + private RestHttpClient.FailureListener failureListener; + private HttpClientConfigCallback httpClientConfigCallback; + private String pathPrefix; + private NodeSelector nodeSelector = NodeSelector.ANY; + private boolean strictDeprecationMode = false; + private boolean compressionEnabled = false; + + /** + * Creates a new builder instance and sets the hosts that the client will send requests to. + * + * @throws IllegalArgumentException if {@code nodes} is {@code null} or empty. + */ + RestHttpClientBuilder(List nodes) { + if (nodes == null || nodes.isEmpty()) { + throw new IllegalArgumentException("nodes must not be null or empty"); + } + for (Node node : nodes) { + if (node == null) { + throw new IllegalArgumentException("node cannot be null"); + } + } + this.nodes = nodes; + } + + /** + * Sets the default request headers, which will be sent along with each request. + *

    + * Request-time headers will always overwrite any default headers. + * + * @param defaultHeaders array of default header + * @throws NullPointerException if {@code defaultHeaders} or any header is {@code null}. + */ + public RestHttpClientBuilder setDefaultHeaders(Map> defaultHeaders) { + Objects.requireNonNull(defaultHeaders, "defaultHeaders must not be null"); + this.defaultHeaders = Collections.unmodifiableMap(defaultHeaders); + return this; + } + + /** + * Sets the {@link RestHttpClient.FailureListener} to be notified for each request failure + * + * @param failureListener the {@link RestHttpClient.FailureListener} for each failure + * @throws NullPointerException if {@code failureListener} is {@code null}. + */ + public RestHttpClientBuilder setFailureListener(RestHttpClient.FailureListener failureListener) { + Objects.requireNonNull(failureListener, "failureListener must not be null"); + this.failureListener = failureListener; + return this; + } + + /** + * Sets the {@link HttpClientConfigCallback} to be used to customize http client configuration + * + * @param httpClientConfigCallback the {@link HttpClientConfigCallback} to be used + * @throws NullPointerException if {@code httpClientConfigCallback} is {@code null}. + */ + public RestHttpClientBuilder setHttpClientConfigCallback(HttpClientConfigCallback httpClientConfigCallback) { + Objects.requireNonNull(httpClientConfigCallback, "httpClientConfigCallback must not be null"); + this.httpClientConfigCallback = httpClientConfigCallback; + return this; + } + + /** + * Sets the path's prefix for every request used by the http client. + *

    + * For example, if this is set to "/my/path", then any client request will become "/my/path/" + endpoint. + *

    + * In essence, every request's {@code endpoint} is prefixed by this {@code pathPrefix}. The path prefix is useful for when + * OpenSearch is behind a proxy that provides a base path or a proxy that requires all paths to start with '/'; + * it is not intended for other purposes and it should not be supplied in other scenarios. + * + * @param pathPrefix the path prefix for every request. + * @throws NullPointerException if {@code pathPrefix} is {@code null}. + * @throws IllegalArgumentException if {@code pathPrefix} is empty, or ends with more than one '/'. + */ + public RestHttpClientBuilder setPathPrefix(String pathPrefix) { + this.pathPrefix = cleanPathPrefix(pathPrefix); + return this; + } + + /** + * Cleans up the given path prefix to ensure that looks like "/base/path". + * + * @param pathPrefix the path prefix to be cleaned up. + * @return the cleaned up path prefix. + * @throws NullPointerException if {@code pathPrefix} is {@code null}. + * @throws IllegalArgumentException if {@code pathPrefix} is empty, or ends with more than one '/'. + */ + public static String cleanPathPrefix(String pathPrefix) { + Objects.requireNonNull(pathPrefix, "pathPrefix must not be null"); + + if (pathPrefix.isEmpty()) { + throw new IllegalArgumentException("pathPrefix must not be empty"); + } + + String cleanPathPrefix = pathPrefix; + if (cleanPathPrefix.startsWith("/") == false) { + cleanPathPrefix = "/" + cleanPathPrefix; + } + + // best effort to ensure that it looks like "/base/path" rather than "/base/path/" + if (cleanPathPrefix.endsWith("/") && cleanPathPrefix.length() > 1) { + cleanPathPrefix = cleanPathPrefix.substring(0, cleanPathPrefix.length() - 1); + + if (cleanPathPrefix.endsWith("/")) { + throw new IllegalArgumentException("pathPrefix is malformed. too many trailing slashes: [" + pathPrefix + "]"); + } + } + return cleanPathPrefix; + } + + /** + * Sets the {@link NodeSelector} to be used for all requests. + * + * @param nodeSelector the {@link NodeSelector} to be used + * @throws NullPointerException if the provided nodeSelector is null + */ + public RestHttpClientBuilder setNodeSelector(NodeSelector nodeSelector) { + Objects.requireNonNull(nodeSelector, "nodeSelector must not be null"); + this.nodeSelector = nodeSelector; + return this; + } + + /** + * Whether the REST client should return any response containing at least + * one warning header as a failure. + * + * @param strictDeprecationMode flag for enabling strict deprecation mode + */ + public RestHttpClientBuilder setStrictDeprecationMode(boolean strictDeprecationMode) { + this.strictDeprecationMode = strictDeprecationMode; + return this; + } + + /** + * Whether the REST client should compress requests using gzip content encoding and add the "Accept-Encoding: gzip" + * header to receive compressed responses. + * + * @param compressionEnabled flag for enabling compression + */ + public RestHttpClientBuilder setCompressionEnabled(boolean compressionEnabled) { + this.compressionEnabled = compressionEnabled; + return this; + } + + /** + * Creates a new {@link RestHttpClient} based on the provided configuration. + */ + public RestHttpClient build() { + if (failureListener == null) { + failureListener = new RestHttpClient.FailureListener(); + } + @SuppressWarnings("removal") + HttpClient httpClient = AccessController.doPrivileged((PrivilegedAction) this::createHttpClient); + + return new RestHttpClient( + httpClient, + defaultHeaders, + nodes, + pathPrefix, + failureListener, + nodeSelector, + strictDeprecationMode, + compressionEnabled + ); + } + + @SuppressWarnings("removal") + private HttpClient createHttpClient() { + try { + HttpClient.Builder httpClientBuilder = HttpClient.newBuilder() + .connectTimeout(Duration.ofMillis(DEFAULT_CONNECT_TIMEOUT_MILLIS)) + .sslContext(SSLContext.getDefault()); + + if (httpClientConfigCallback != null) { + httpClientBuilder = httpClientConfigCallback.customizeHttpClient(httpClientBuilder); + } + + final HttpClient.Builder finalBuilder = httpClientBuilder; + return AccessController.doPrivileged((PrivilegedAction) finalBuilder::build); + } catch (NoSuchAlgorithmException e) { + throw new IllegalStateException("could not create the default ssl context", e); + } + } + + /** + * Callback used to customize the {@link HttpClient} instance used by a {@link RestHttpClient} instance. + */ + public interface HttpClientConfigCallback { + /** + * Allows to customize the {@link HttpClient} being created and used by the {@link RestHttpClient}. + * Commonly used to customize the default {@link Authenticator} for authentication for communication + * through TLS/SSL without losing any other useful default value that the {@link RestHttpClientBuilder} internally + * sets, like connection pooling. + * + * @param httpClientBuilder the {@link HttpClient.Builder} for customizing the client instance. + */ + HttpClient.Builder customizeHttpClient(HttpClient.Builder httpClientBuilder); + } +} diff --git a/client/rest-http-client/src/main/java/org/opensearch/internal/httpclient/StatusLine.java b/client/rest-http-client/src/main/java/org/opensearch/internal/httpclient/StatusLine.java new file mode 100644 index 0000000000000..120b042c01357 --- /dev/null +++ b/client/rest-http-client/src/main/java/org/opensearch/internal/httpclient/StatusLine.java @@ -0,0 +1,38 @@ +/* + * SPDX-License-Identifier: Apache-2.0 + * + * The OpenSearch Contributors require contributions made to + * this file be licensed under the Apache-2.0 license or a + * compatible open source license. + */ + +package org.opensearch.internal.httpclient; + +import java.net.http.HttpClient.Version; +import java.net.http.HttpResponse; +import java.util.Objects; + +/** + * Response status line (protocol, status code) + * Note: This is an experimental API. + */ +public record StatusLine(Version protoVersion, int statusCode) { + /** + * Creates a new status line with the given version and status. + * + * @param protoVersion the protocol version of the response + * @param statusCode the status code of the response + */ + public StatusLine { + protoVersion = protoVersion != null ? protoVersion : Version.HTTP_1_1; + } + + /** + * Creates a new status line from the response + * + * @param response HTTP response + */ + public StatusLine(final HttpResponse response) { + this(Objects.requireNonNull(response, "Response").version(), Objects.requireNonNull(response, "Response").statusCode()); + } +} diff --git a/client/rest-http-client/src/main/java/org/opensearch/internal/httpclient/StreamingRequest.java b/client/rest-http-client/src/main/java/org/opensearch/internal/httpclient/StreamingRequest.java new file mode 100644 index 0000000000000..69c943a6123f3 --- /dev/null +++ b/client/rest-http-client/src/main/java/org/opensearch/internal/httpclient/StreamingRequest.java @@ -0,0 +1,212 @@ +/* + * SPDX-License-Identifier: Apache-2.0 + * + * The OpenSearch Contributors require contributions made to + * this file be licensed under the Apache-2.0 license or a + * compatible open source license. + */ + +package org.opensearch.internal.httpclient; + +import java.nio.ByteBuffer; +import java.util.HashMap; +import java.util.Map; +import java.util.Objects; + +import org.reactivestreams.Publisher; + +import static java.util.Collections.unmodifiableMap; + +/** + * HTTP Streaming Request to OpenSearch. + * Note: This is an experimental API. + */ +public final class StreamingRequest { + private final String method; + private final String endpoint; + private final Map parameters = new HashMap<>(); + private final RequestOptions options; + private final Publisher publisher; + + /** + * Constructor + * @param method method + * @param endpoint endpoint + * @param publisher publisher + */ + private StreamingRequest(String method, String endpoint, Publisher publisher) { + this(method, endpoint, Map.of(), publisher, RequestOptions.DEFAULT); + } + + /** + * Constructor + * @param method method + * @param endpoint endpoint + * @param builder request options builder + * @param publisher publisher + */ + private StreamingRequest( + String method, + String endpoint, + Map parameters, + Publisher publisher, + RequestOptions.Builder builder + ) { + this(method, endpoint, parameters, publisher, builder.build()); + } + + /** + * Constructor + * @param method method + * @param endpoint endpoint + * @param options request options + * @param publisher publisher + */ + private StreamingRequest( + String method, + String endpoint, + Map parameters, + Publisher publisher, + RequestOptions options + ) { + this.method = method; + this.endpoint = endpoint; + this.publisher = publisher; + this.options = options; + } + + public static StreamingRequest.Builder newRequest(String method, String endpoint, Publisher publisher) { + return new StreamingRequest.Builder(method, endpoint, publisher); + } + + /** + * Get endpoint + * @return endpoint + */ + public String endpoint() { + return endpoint; + } + + /** + * Get method + * @return method + */ + public String method() { + return method; + } + + /** + * Get options + * @return options + */ + public RequestOptions options() { + return options; + } + + /** + * Get parameters + * @return parameters + */ + public Map parameters() { + if (options.getParameters().isEmpty()) { + return unmodifiableMap(parameters); + } else { + Map combinedParameters = new HashMap<>(parameters); + combinedParameters.putAll(options.getParameters()); + return unmodifiableMap(combinedParameters); + } + } + + /** + * Body publisher + * @return body publisher + */ + public Publisher body() { + return publisher; + } + + public static final class Builder { + private final String method; + private final String endpoint; + private final Map parameters = new HashMap<>(); + private final Publisher publisher; + private RequestOptions options = RequestOptions.DEFAULT; + + private Builder(String method, String endpoint, Publisher publisher) { + this.method = Objects.requireNonNull(method, "method cannot be null"); + this.endpoint = Objects.requireNonNull(endpoint, "endpoint cannot be null"); + this.publisher = Objects.requireNonNull(publisher, "publisher cannot be null"); + } + + /** + * Set the portion of an HTTP request to OpenSearch that can be + * manipulated without changing OpenSearch's behavior. + * + * @param options the options to be set. + * @throws NullPointerException if {@code options} is null. + */ + public void setOptions(RequestOptions options) { + Objects.requireNonNull(options, "options cannot be null"); + this.options = options; + } + + /** + * Add a query string parameter. + * @param name the name of the url parameter. Must not be null. + * @param value the value of the url url parameter. If {@code null} then + * the parameter is sent as {@code name} rather than {@code name=value} + * @throws IllegalArgumentException if a parameter with that name has + * already been set + */ + public Builder withParameter(String name, String value) { + Objects.requireNonNull(name, "url parameter name cannot be null"); + if (parameters.containsKey(name)) { + throw new IllegalArgumentException("url parameter [" + name + "] has already been set to [" + parameters.get(name) + "]"); + } else { + parameters.put(name, value); + } + return this; + } + + /** + * Add query parameters using the provided map of key value pairs. + * + * @param paramSource a map of key value pairs where the key is the url parameter. + * @throws IllegalArgumentException if a parameter with that name has already been set. + */ + public Builder withParameters(Map paramSource) { + paramSource.forEach(this::withParameter); + return this; + } + + /** + * Set the portion of an HTTP request to OpenSearch that can be + * manipulated without changing OpenSearch's behavior. + * + * @param options the options to be set. + * @throws NullPointerException if {@code options} is null. + */ + public Builder withOptions(RequestOptions.Builder options) { + Objects.requireNonNull(options, "options cannot be null"); + this.options = options.build(); + return this; + } + + /** + * Set the portion of an HTTP request to OpenSearch that can be + * manipulated without changing OpenSearch's behavior. + * + * @param options the options to be set. + * @throws NullPointerException if {@code options} is null. + */ + public Builder withOptions(RequestOptions options) { + Objects.requireNonNull(options, "options cannot be null"); + this.options = options; + return this; + } + + public StreamingRequest build() { + return new StreamingRequest(method, endpoint, parameters, publisher, options); + } + } +} diff --git a/client/rest-http-client/src/main/java/org/opensearch/internal/httpclient/StreamingResponse.java b/client/rest-http-client/src/main/java/org/opensearch/internal/httpclient/StreamingResponse.java new file mode 100644 index 0000000000000..629adf2a4b532 --- /dev/null +++ b/client/rest-http-client/src/main/java/org/opensearch/internal/httpclient/StreamingResponse.java @@ -0,0 +1,121 @@ +/* + * SPDX-License-Identifier: Apache-2.0 + * + * The OpenSearch Contributors require contributions made to + * this file be licensed under the Apache-2.0 license or a + * compatible open source license. + */ + +package org.opensearch.internal.httpclient; + +import java.net.http.HttpHeaders; +import java.net.http.HttpResponse; +import java.nio.ByteBuffer; +import java.util.List; +import java.util.concurrent.Flow; + +import org.reactivestreams.Publisher; +import reactor.adapter.JdkFlowAdapter; +import reactor.core.publisher.Mono; + +/** + * HTTP Streaming Response from OpenSearch. + * Note: This is an experimental API. + */ +public final class StreamingResponse { + private final RequestLine requestLine; + private final Mono>>> publisher; + + /** + * Constructor + * @param requestLine request line + * @param publisher message publisher(response with a body) + */ + public StreamingResponse(RequestLine requestLine, Publisher>>> publisher) { + this.requestLine = requestLine; + // We cache the publisher here so the body or / and HttpResponse could + // be consumed independently or/and more than once. + this.publisher = Mono.from(publisher).cache(); + } + + /** + * Get request line + * @return request line + */ + public RequestLine requestLine() { + return requestLine; + } + + /** + * Get response boby {@link Publisher} + * @return response boby {@link Publisher} + */ + public Publisher body() { + return publisher.flatMapMany(m -> { + final boolean compressed = m.headers() + .firstValue("Content-Encoding") + .filter("gzip"::equalsIgnoreCase) + .map(h -> true) + .orElse(false); + return JdkFlowAdapter.flowPublisherToFlux(m.body()).flatMapIterable(t -> t).map(b -> { + if (compressed) { + return BodyUtils.decompress(b); + } else { + return b; + } + }); + }); + } + + /** + * Returns the status line of the current response + */ + @SuppressWarnings("unchecked") + public StatusLine statusLine() { + return new StatusLine( + publisher.onErrorResume( + ResponseException.class, + e -> Mono.just((HttpResponse>>) e.getResponse().httpResponse()) + ).block() + ); + } + + /** + * Returns a list of all warning headers returned in the response. + */ + @SuppressWarnings("unchecked") + public List warnings() { + return ResponseWarningsExtractor.getWarnings( + publisher.onErrorResume( + ResponseException.class, + e -> Mono.just((HttpResponse>>) e.getResponse().httpResponse()) + ).block() + ); + } + + /** + * Returns a list of all headers returned in the response. + */ + @SuppressWarnings("unchecked") + public HttpHeaders headers() { + return publisher.onErrorResume( + ResponseException.class, + e -> Mono.just((HttpResponse>>) e.getResponse().httpResponse()) + ).map(HttpResponse::headers).block(); + } + + /** + * Returns the value of the first header with a specified name of this message. + * If there is more than one matching header in the message the first element is returned. + * If there is no matching header in the message null is returned. + * + * @param name header name + */ + @SuppressWarnings("unchecked") + public String header(String name) { + return publisher.onErrorResume( + ResponseException.class, + e -> Mono.just((HttpResponse>>) e.getResponse().httpResponse()) + ).mapNotNull(response -> response.headers().firstValue(name).orElse(null)).block(); + } +} diff --git a/client/rest-http-client/src/main/java/org/opensearch/internal/httpclient/WarningFailureException.java b/client/rest-http-client/src/main/java/org/opensearch/internal/httpclient/WarningFailureException.java new file mode 100644 index 0000000000000..53a0a1f97fe4b --- /dev/null +++ b/client/rest-http-client/src/main/java/org/opensearch/internal/httpclient/WarningFailureException.java @@ -0,0 +1,79 @@ +/* + * SPDX-License-Identifier: Apache-2.0 + * + * The OpenSearch Contributors require contributions made to + * this file be licensed under the Apache-2.0 license or a + * compatible open source license. + */ + +/* + * Licensed to Elasticsearch under one or more contributor + * license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright + * ownership. Elasticsearch 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. + */ + +/* + * Modifications Copyright OpenSearch Contributors. See + * GitHub history for details. + */ + +package org.opensearch.internal.httpclient; + +import java.io.IOException; + +import static org.opensearch.internal.httpclient.ResponseException.buildMessage; + +/** + * This exception is used to indicate that one or more {@link Response#warnings()} exist + * and is typically used when the {@link RestHttpClient} is set to fail by setting + * {@link RestHttpClientBuilder#setStrictDeprecationMode(boolean)} to `true`. + */ +// This class extends RuntimeException in order to deal with wrapping that is done in FutureUtils on exception. +// if the exception is not of type OpenSearchException or RuntimeException it will be wrapped in a UncategorizedExecutionException +public final class WarningFailureException extends RuntimeException { + + private final Response response; + + /** + * Creates a {@link WarningFailureException} instance. + * + * @param response the response that contains warnings. + * @throws IOException if there is a problem building the exception message. + */ + public WarningFailureException(Response response) throws IOException { + super(buildMessage(response)); + this.response = response; + } + + /** + * Wrap a {@linkplain WarningFailureException} with another one with the current + * stack trace. This is used during synchronous calls so that the caller + * ends up in the stack trace of the exception thrown. + * + * @param e the exception to be wrapped. + */ + WarningFailureException(WarningFailureException e) { + super(e.getMessage(), e); + this.response = e.getResponse(); + } + + /** + * Returns the {@link Response} that caused this exception to be thrown. + */ + public Response getResponse() { + return response; + } +} diff --git a/client/rest-http-client/src/main/java/org/opensearch/internal/httpclient/WarningsHandler.java b/client/rest-http-client/src/main/java/org/opensearch/internal/httpclient/WarningsHandler.java new file mode 100644 index 0000000000000..1d59a185a6a3e --- /dev/null +++ b/client/rest-http-client/src/main/java/org/opensearch/internal/httpclient/WarningsHandler.java @@ -0,0 +1,80 @@ +/* + * SPDX-License-Identifier: Apache-2.0 + * + * The OpenSearch Contributors require contributions made to + * this file be licensed under the Apache-2.0 license or a + * compatible open source license. + */ + +/* + * Licensed to Elasticsearch under one or more contributor + * license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright + * ownership. Elasticsearch 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. + */ + +/* + * Modifications Copyright OpenSearch Contributors. See + * GitHub history for details. + */ + +package org.opensearch.internal.httpclient; + +import java.util.List; + +/** + * Called if there are warnings to determine if those warnings should fail the + * request. + */ +public interface WarningsHandler { + + /** + * Determines whether the given list of warnings should fail the request. + * + * @param warnings a list of warnings. + * @return boolean indicating if the request should fail. + */ + boolean warningsShouldFailRequest(List warnings); + + /** + * The permissive warnings handler. Warnings will not fail the request. + */ + WarningsHandler PERMISSIVE = new WarningsHandler() { + @Override + public boolean warningsShouldFailRequest(List warnings) { + return false; + } + + @Override + public String toString() { + return "permissive"; + } + }; + + /** + * The strict warnings handler. Warnings will fail the request. + */ + WarningsHandler STRICT = new WarningsHandler() { + @Override + public boolean warningsShouldFailRequest(List warnings) { + return false == warnings.isEmpty(); + } + + @Override + public String toString() { + return "strict"; + } + }; +} diff --git a/client/rest-http-client/src/test/java/org/opensearch/internal/httpclient/BouncyCastleThreadFilter.java b/client/rest-http-client/src/test/java/org/opensearch/internal/httpclient/BouncyCastleThreadFilter.java new file mode 100644 index 0000000000000..814cf4bfaff11 --- /dev/null +++ b/client/rest-http-client/src/test/java/org/opensearch/internal/httpclient/BouncyCastleThreadFilter.java @@ -0,0 +1,25 @@ +/* + * SPDX-License-Identifier: Apache-2.0 + * + * The OpenSearch Contributors require contributions made to + * this file be licensed under the Apache-2.0 license or a + * compatible open source license. + */ + +package org.opensearch.internal.httpclient; + +import com.carrotsearch.randomizedtesting.ThreadFilter; + +/** + * ThreadFilter to exclude ThreadLeak checks for BC’s global background threads + * + *

    clone from the original, which is located in ':test:framework'

    + */ +public class BouncyCastleThreadFilter implements ThreadFilter { + @Override + public boolean reject(Thread t) { + String n = t.getName(); + // Ignore BC’s global background threads + return "BC Disposal Daemon".equals(n) || "BC Cleanup Executor".equals(n); + } +} diff --git a/client/rest-http-client/src/test/java/org/opensearch/internal/httpclient/HostsTrackingFailureListener.java b/client/rest-http-client/src/test/java/org/opensearch/internal/httpclient/HostsTrackingFailureListener.java new file mode 100644 index 0000000000000..c18c6539f8bc4 --- /dev/null +++ b/client/rest-http-client/src/test/java/org/opensearch/internal/httpclient/HostsTrackingFailureListener.java @@ -0,0 +1,71 @@ +/* + * SPDX-License-Identifier: Apache-2.0 + * + * The OpenSearch Contributors require contributions made to + * this file be licensed under the Apache-2.0 license or a + * compatible open source license. + */ + +/* + * Licensed to Elasticsearch under one or more contributor + * license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright + * ownership. Elasticsearch 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. + */ + +/* + * Modifications Copyright OpenSearch Contributors. See + * GitHub history for details. + */ + +package org.opensearch.internal.httpclient; + +import java.util.HashSet; +import java.util.List; +import java.util.Set; + +import static org.hamcrest.Matchers.containsInAnyOrder; +import static org.junit.Assert.assertEquals; +import static org.junit.Assert.assertThat; + +/** + * {@link RestHttpClient.FailureListener} impl that allows to track when it gets called for which host. + */ +class HostsTrackingFailureListener extends RestHttpClient.FailureListener { + private volatile Set hosts = new HashSet<>(); + + @Override + public void onFailure(Node node) { + hosts.add(node.getHost()); + } + + void assertCalled(List nodes) { + HttpHost[] hosts = new HttpHost[nodes.size()]; + for (int i = 0; i < nodes.size(); i++) { + hosts[i] = nodes.get(i).getHost(); + } + assertCalled(hosts); + } + + void assertCalled(HttpHost... hosts) { + assertEquals(hosts.length, this.hosts.size()); + assertThat(this.hosts, containsInAnyOrder(hosts)); + this.hosts.clear(); + } + + void assertNotCalled() { + assertEquals(0, hosts.size()); + } +} diff --git a/client/rest-http-client/src/test/java/org/opensearch/internal/httpclient/HttpClientThreadLeakFilter.java b/client/rest-http-client/src/test/java/org/opensearch/internal/httpclient/HttpClientThreadLeakFilter.java new file mode 100644 index 0000000000000..9ca72d417775c --- /dev/null +++ b/client/rest-http-client/src/test/java/org/opensearch/internal/httpclient/HttpClientThreadLeakFilter.java @@ -0,0 +1,24 @@ +/* + * SPDX-License-Identifier: Apache-2.0 + * + * The OpenSearch Contributors require contributions made to + * this file be licensed under the Apache-2.0 license or a + * compatible open source license. + */ + +package org.opensearch.internal.httpclient; + +import com.carrotsearch.randomizedtesting.ThreadFilter; + +import java.net.http.HttpClient; + +/** + * The {@link HttpClient} creates own ASYNC pool based of {@code ForkJoin.commonPool()} which + * is impossible to supress or override. + */ +public final class HttpClientThreadLeakFilter implements ThreadFilter { + @Override + public boolean reject(Thread t) { + return t.getName().startsWith("ForkJoinPool.commonPool-"); + } +} diff --git a/client/rest-http-client/src/test/java/org/opensearch/internal/httpclient/RestClientTestUtil.java b/client/rest-http-client/src/test/java/org/opensearch/internal/httpclient/RestClientTestUtil.java new file mode 100644 index 0000000000000..ac1dbcb4b8a2e --- /dev/null +++ b/client/rest-http-client/src/test/java/org/opensearch/internal/httpclient/RestClientTestUtil.java @@ -0,0 +1,120 @@ +/* + * SPDX-License-Identifier: Apache-2.0 + * + * The OpenSearch Contributors require contributions made to + * this file be licensed under the Apache-2.0 license or a + * compatible open source license. + */ + +/* + * Licensed to Elasticsearch under one or more contributor + * license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright + * ownership. Elasticsearch 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. + */ + +/* + * Modifications Copyright OpenSearch Contributors. See + * GitHub history for details. + */ + +package org.opensearch.internal.httpclient; + +import com.carrotsearch.randomizedtesting.generators.RandomNumbers; +import com.carrotsearch.randomizedtesting.generators.RandomPicks; +import com.carrotsearch.randomizedtesting.generators.RandomStrings; + +import java.util.ArrayList; +import java.util.Arrays; +import java.util.HashMap; +import java.util.List; +import java.util.Map; +import java.util.Random; + +final class RestClientTestUtil { + + private static final String[] HTTP_METHODS = new String[] { "DELETE", "HEAD", "GET", "OPTIONS", "PATCH", "POST", "PUT", "TRACE" }; + private static final List ALL_STATUS_CODES; + private static final List OK_STATUS_CODES = Arrays.asList(200, 201); + private static final List ALL_ERROR_STATUS_CODES; + private static List ERROR_NO_RETRY_STATUS_CODES = Arrays.asList(400, 401, 403, 404, 405, 500); + private static List ERROR_RETRY_STATUS_CODES = Arrays.asList(502, 503, 504); + + static { + ALL_ERROR_STATUS_CODES = new ArrayList<>(ERROR_RETRY_STATUS_CODES); + ALL_ERROR_STATUS_CODES.addAll(ERROR_NO_RETRY_STATUS_CODES); + ALL_STATUS_CODES = new ArrayList<>(ALL_ERROR_STATUS_CODES); + ALL_STATUS_CODES.addAll(OK_STATUS_CODES); + } + + private RestClientTestUtil() { + + } + + static String[] getHttpMethods() { + return HTTP_METHODS; + } + + static String randomHttpMethod(Random random) { + return RandomPicks.randomFrom(random, HTTP_METHODS); + } + + static int randomStatusCode(Random random) { + return RandomPicks.randomFrom(random, ALL_STATUS_CODES); + } + + static int randomOkStatusCode(Random random) { + return RandomPicks.randomFrom(random, OK_STATUS_CODES); + } + + static int randomErrorNoRetryStatusCode(Random random) { + return RandomPicks.randomFrom(random, ERROR_NO_RETRY_STATUS_CODES); + } + + static int randomErrorRetryStatusCode(Random random) { + return RandomPicks.randomFrom(random, ERROR_RETRY_STATUS_CODES); + } + + static List getOkStatusCodes() { + return OK_STATUS_CODES; + } + + static List getAllErrorStatusCodes() { + return ALL_ERROR_STATUS_CODES; + } + + static List getAllStatusCodes() { + return ALL_STATUS_CODES; + } + + /** + * Create a random number of HTTP headers. + * Generated header names will either be the {@code baseName} plus its index, or exactly the provided {@code baseName} so that the + * we test also support for multiple headers with same key and different values. + */ + static Map> randomHeaders(Random random, final String baseName) { + final Map> headers = new HashMap<>(); + int numHeaders = RandomNumbers.randomIntBetween(random, 0, 5); + for (int i = 0; i < numHeaders; i++) { + String headerName = baseName; + // randomly exercise the code path that supports multiple headers with same key + if (random.nextBoolean()) { + headerName = headerName + i; + } + headers.put(headerName, List.of(RandomStrings.randomAsciiLettersOfLengthBetween(random, 3, 10))); + } + return headers; + } +} diff --git a/client/rest-http-client/src/test/java/org/opensearch/internal/httpclient/RestHttpClientCompressionTests.java b/client/rest-http-client/src/test/java/org/opensearch/internal/httpclient/RestHttpClientCompressionTests.java new file mode 100644 index 0000000000000..6a74a752fddd0 --- /dev/null +++ b/client/rest-http-client/src/test/java/org/opensearch/internal/httpclient/RestHttpClientCompressionTests.java @@ -0,0 +1,142 @@ +/* + * SPDX-License-Identifier: Apache-2.0 + * + * The OpenSearch Contributors require contributions made to + * this file be licensed under the Apache-2.0 license or a + * compatible open source license. + */ + +package org.opensearch.internal.httpclient; + +import com.sun.net.httpserver.HttpExchange; +import com.sun.net.httpserver.HttpHandler; +import com.sun.net.httpserver.HttpServer; + +import org.junit.AfterClass; +import org.junit.Assert; +import org.junit.BeforeClass; + +import java.io.IOException; +import java.io.InputStream; +import java.io.OutputStream; +import java.net.InetAddress; +import java.net.InetSocketAddress; +import java.net.http.HttpClient.Version; +import java.net.http.HttpRequest.BodyPublishers; +import java.nio.charset.StandardCharsets; +import java.util.concurrent.CompletableFuture; +import java.util.zip.GZIPInputStream; +import java.util.zip.GZIPOutputStream; + +public class RestHttpClientCompressionTests extends RestHttpClientTestCase { + + private static HttpServer httpServer; + + @BeforeClass + public static void startHttpServer() throws Exception { + httpServer = HttpServer.create(new InetSocketAddress(InetAddress.getLoopbackAddress(), 0), 0); + httpServer.createContext("/", new GzipResponseHandler()); + httpServer.start(); + } + + @AfterClass + public static void stopHttpServers() throws IOException { + httpServer.stop(0); + httpServer = null; + } + + /** + * A response handler that accepts gzip-encoded data and replies request and response encoding values + * followed by the request body. The response is compressed if "Accept-Encoding" is "gzip". + */ + private static class GzipResponseHandler implements HttpHandler { + @Override + public void handle(HttpExchange exchange) throws IOException { + + // Decode body (if any) + String contentEncoding = exchange.getRequestHeaders().getFirst("Content-Encoding"); + String contentLength = exchange.getRequestHeaders().getFirst("Content-Length"); + InputStream body = exchange.getRequestBody(); + boolean compressedRequest = false; + if ("gzip".equals(contentEncoding)) { + body = new GZIPInputStream(body); + compressedRequest = true; + } + byte[] bytes = body.readAllBytes(); + boolean compress = "gzip".equals(exchange.getRequestHeaders().getFirst("Accept-Encoding")); + if (compress) { + exchange.getResponseHeaders().add("Content-Encoding", "gzip"); + } + + exchange.sendResponseHeaders(200, 0); + + // Encode response if needed + OutputStream out = exchange.getResponseBody(); + if (compress) { + out = new GZIPOutputStream(out); + } + + // Outputs ## + out.write(String.valueOf(contentEncoding).getBytes(StandardCharsets.UTF_8)); + out.write('#'); + out.write((compress ? "gzip" : "null").getBytes(StandardCharsets.UTF_8)); + out.write('#'); + out.write(((compressedRequest == true && contentLength != null) ? contentLength : "null").getBytes(StandardCharsets.UTF_8)); + out.write('#'); + out.write(bytes); + out.close(); + + exchange.close(); + } + } + + private RestHttpClient createClient(boolean enableCompression) { + InetSocketAddress address = httpServer.getAddress(); + return RestHttpClient.builder(new HttpHost("http", address.getHostString(), address.getPort())) + .setCompressionEnabled(enableCompression) + .setHttpClientConfigCallback(builder -> builder.version(Version.HTTP_1_1)) + .build(); + } + + public void testCompressingClientWithContentLengthSync() throws Exception { + try (RestHttpClient restClient = createClient(true)) { + Request request = Request.newRequest("POST", "/").withEntity(BodyPublishers.ofString("compressing client")).build(); + + Response response = restClient.performRequest(request); + + String content = BodyUtils.getBodyAsString(response.entity()); + // Content-Encoding#Accept-Encoding#Content-Length#Content + // With HttpClient, we don't sent Content-Length so it is always null + Assert.assertEquals("gzip#gzip#null#compressing client", content); + } + } + + public void testCompressingClientContentLengthAsync() throws Exception { + try (RestHttpClient restClient = createClient(true)) { + Request request = Request.newRequest("POST", "/").withEntity(BodyPublishers.ofString("compressing client")).build(); + + FutureResponse futureResponse = new FutureResponse(); + restClient.performRequestAsync(request, futureResponse); + Response response = futureResponse.get(); + + // Server should report it had a compressed request and sent back a compressed response + String content = BodyUtils.getBodyAsString(response.entity()); + + // Content-Encoding#Accept-Encoding#Content-Length#Content + // With HttpClient, we don't sent Content-Length so it is always null + Assert.assertEquals("gzip#gzip#null#compressing client", content); + } + } + + public static class FutureResponse extends CompletableFuture implements ResponseListener { + @Override + public void onSuccess(Response response) { + this.complete(response); + } + + @Override + public void onFailure(Exception exception) { + this.completeExceptionally(exception); + } + } +} diff --git a/client/rest-http-client/src/test/java/org/opensearch/internal/httpclient/RestHttpClientGzipCompressionTests.java b/client/rest-http-client/src/test/java/org/opensearch/internal/httpclient/RestHttpClientGzipCompressionTests.java new file mode 100644 index 0000000000000..64a7bb01f2b7a --- /dev/null +++ b/client/rest-http-client/src/test/java/org/opensearch/internal/httpclient/RestHttpClientGzipCompressionTests.java @@ -0,0 +1,191 @@ +/* + * SPDX-License-Identifier: Apache-2.0 + * + * The OpenSearch Contributors require contributions made to + * this file be licensed under the Apache-2.0 license or a + * compatible open source license. + */ + +/* + * Licensed to Elasticsearch under one or more contributor + * license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright + * ownership. Elasticsearch 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. + */ + +/* + * Modifications Copyright OpenSearch Contributors. See + * GitHub history for details. + */ + +package org.opensearch.internal.httpclient; + +import com.sun.net.httpserver.HttpExchange; +import com.sun.net.httpserver.HttpHandler; +import com.sun.net.httpserver.HttpServer; + +import org.junit.AfterClass; +import org.junit.Assert; +import org.junit.BeforeClass; + +import java.io.IOException; +import java.io.InputStream; +import java.io.OutputStream; +import java.net.InetAddress; +import java.net.InetSocketAddress; +import java.net.http.HttpRequest.BodyPublishers; +import java.nio.charset.StandardCharsets; +import java.util.concurrent.CompletableFuture; +import java.util.zip.GZIPInputStream; +import java.util.zip.GZIPOutputStream; + +public class RestHttpClientGzipCompressionTests extends RestHttpClientTestCase { + + private static HttpServer httpServer; + + @BeforeClass + public static void startHttpServer() throws Exception { + httpServer = HttpServer.create(new InetSocketAddress(InetAddress.getLoopbackAddress(), 0), 0); + httpServer.createContext("/", new GzipResponseHandler()); + httpServer.start(); + } + + @AfterClass + public static void stopHttpServers() throws IOException { + httpServer.stop(0); + httpServer = null; + } + + /** + * A response handler that accepts gzip-encoded data and replies request and response encoding values + * followed by the request body. The response is compressed if "Accept-Encoding" is "gzip". + */ + private static class GzipResponseHandler implements HttpHandler { + @Override + public void handle(HttpExchange exchange) throws IOException { + + // Decode body (if any) + String contentEncoding = exchange.getRequestHeaders().getFirst("Content-Encoding"); + InputStream body = exchange.getRequestBody(); + if ("gzip".equals(contentEncoding)) { + body = new GZIPInputStream(body); + } + byte[] bytes = body.readAllBytes(); + boolean compress = "gzip".equals(exchange.getRequestHeaders().getFirst("Accept-Encoding")); + if (compress) { + exchange.getResponseHeaders().add("Content-Encoding", "gzip"); + } + + exchange.sendResponseHeaders(200, 0); + + // Encode response if needed + OutputStream out = exchange.getResponseBody(); + if (compress) { + out = new GZIPOutputStream(out); + } + + // Outputs ## + out.write(String.valueOf(contentEncoding).getBytes(StandardCharsets.UTF_8)); + out.write('#'); + out.write((compress ? "gzip" : "null").getBytes(StandardCharsets.UTF_8)); + out.write('#'); + out.write(bytes); + out.close(); + + exchange.close(); + } + } + + private RestHttpClient createClient(boolean enableCompression) { + InetSocketAddress address = httpServer.getAddress(); + return RestHttpClient.builder(new HttpHost("http", address.getHostString(), address.getPort())) + .setCompressionEnabled(enableCompression) + .build(); + } + + public void testGzipHeaderSync() throws Exception { + try (RestHttpClient restClient = createClient(false)) { + // Send non-compressed request, expect compressed response + Request request = Request.newRequest("POST", "/") + .withEntity(BodyPublishers.ofString("plain request, gzip response")) + .withOptions(RequestOptions.DEFAULT.toBuilder().addHeader("Accept-Encoding", "gzip")) + .build(); + + Response response = restClient.performRequest(request); + + String content = BodyUtils.getBodyAsString(response.entity()); + Assert.assertEquals("null#gzip#plain request, gzip response", content); + } + } + + public void testGzipHeaderAsync() throws Exception { + try (RestHttpClient restClient = createClient(false)) { + // Send non-compressed request, expect compressed response + Request request = Request.newRequest("POST", "/") + .withEntity(BodyPublishers.ofString("plain request, gzip response")) + .withOptions(RequestOptions.DEFAULT.toBuilder().addHeader("Accept-Encoding", "gzip")) + .build(); + + FutureResponse futureResponse = new FutureResponse(); + restClient.performRequestAsync(request, futureResponse); + Response response = futureResponse.get(); + + String content = BodyUtils.getBodyAsString(response.entity()); + Assert.assertEquals("null#gzip#plain request, gzip response", content); + } + } + + public void testCompressingClientSync() throws Exception { + try (RestHttpClient restClient = createClient(true)) { + Request request = Request.newRequest("POST", "/").withEntity(BodyPublishers.ofString("compressing client")).build(); + + Response response = restClient.performRequest(request); + + String content = BodyUtils.getBodyAsString(response.entity()); + Assert.assertEquals("gzip#gzip#compressing client", content); + } + } + + public void testCompressingClientAsync() throws Exception { + InetSocketAddress address = httpServer.getAddress(); + try ( + RestHttpClient restClient = RestHttpClient.builder(new HttpHost("http", address.getHostString(), address.getPort())) + .setCompressionEnabled(true) + .build() + ) { + Request request = Request.newRequest("POST", "/").withEntity(BodyPublishers.ofString("compressing client")).build(); + + FutureResponse futureResponse = new FutureResponse(); + restClient.performRequestAsync(request, futureResponse); + Response response = futureResponse.get(); + + // Server should report it had a compressed request and sent back a compressed response + String content = BodyUtils.getBodyAsString(response.entity()); + Assert.assertEquals("gzip#gzip#compressing client", content); + } + } + + public static class FutureResponse extends CompletableFuture implements ResponseListener { + @Override + public void onSuccess(Response response) { + this.complete(response); + } + + @Override + public void onFailure(Exception exception) { + this.completeExceptionally(exception); + } + } +} diff --git a/client/rest-http-client/src/test/java/org/opensearch/internal/httpclient/RestHttpClientMultipleHostsIntegTests.java b/client/rest-http-client/src/test/java/org/opensearch/internal/httpclient/RestHttpClientMultipleHostsIntegTests.java new file mode 100644 index 0000000000000..74799a6c5159b --- /dev/null +++ b/client/rest-http-client/src/test/java/org/opensearch/internal/httpclient/RestHttpClientMultipleHostsIntegTests.java @@ -0,0 +1,346 @@ +/* + * SPDX-License-Identifier: Apache-2.0 + * + * The OpenSearch Contributors require contributions made to + * this file be licensed under the Apache-2.0 license or a + * compatible open source license. + */ + +/* + * Licensed to Elasticsearch under one or more contributor + * license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright + * ownership. Elasticsearch 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. + */ + +/* + * Modifications Copyright OpenSearch Contributors. See + * GitHub history for details. + */ + +package org.opensearch.internal.httpclient; + +import com.sun.net.httpserver.HttpExchange; +import com.sun.net.httpserver.HttpHandler; +import com.sun.net.httpserver.HttpServer; + +import org.junit.AfterClass; +import org.junit.Before; +import org.junit.BeforeClass; +import org.junit.Ignore; + +import java.io.IOException; +import java.net.ConnectException; +import java.net.InetAddress; +import java.net.InetSocketAddress; +import java.net.http.HttpClient; +import java.util.ArrayList; +import java.util.Iterator; +import java.util.List; +import java.util.concurrent.CancellationException; +import java.util.concurrent.CopyOnWriteArrayList; +import java.util.concurrent.CountDownLatch; +import java.util.concurrent.TimeUnit; + +import static org.opensearch.internal.httpclient.RestClientTestUtil.getAllStatusCodes; +import static org.opensearch.internal.httpclient.RestClientTestUtil.randomErrorNoRetryStatusCode; +import static org.opensearch.internal.httpclient.RestClientTestUtil.randomOkStatusCode; +import static org.hamcrest.CoreMatchers.containsString; +import static org.hamcrest.CoreMatchers.instanceOf; +import static org.junit.Assert.assertEquals; +import static org.junit.Assert.assertThat; +import static org.junit.Assert.assertTrue; +import static org.junit.Assert.fail; + +/** + * Integration test to check interaction between {@link RestHttpClient} and {@link HttpClient}. + * Works against real http servers, multiple hosts. Also tests failover by randomly shutting down hosts. + */ +public class RestHttpClientMultipleHostsIntegTests extends RestHttpClientTestCase { + + private static WaitForCancelHandler waitForCancelHandler; + private static HttpServer[] httpServers; + private static HttpHost[] httpHosts; + private static boolean stoppedFirstHost = false; + private static String pathPrefixWithoutLeadingSlash; + private static String pathPrefix; + private static RestHttpClient restClient; + + @BeforeClass + public static void startHttpServer() throws Exception { + if (randomBoolean()) { + pathPrefixWithoutLeadingSlash = "testPathPrefix/" + randomAsciiLettersOfLengthBetween(1, 5); + pathPrefix = "/" + pathPrefixWithoutLeadingSlash; + } else { + pathPrefix = pathPrefixWithoutLeadingSlash = ""; + } + int numHttpServers = randomIntBetween(2, 4); + httpServers = new HttpServer[numHttpServers]; + httpHosts = new HttpHost[numHttpServers]; + waitForCancelHandler = new WaitForCancelHandler(); + for (int i = 0; i < numHttpServers; i++) { + HttpServer httpServer = createHttpServer(); + httpServers[i] = httpServer; + httpHosts[i] = new HttpHost("http", httpServer.getAddress().getHostString(), httpServer.getAddress().getPort()); + } + restClient = buildRestClient(NodeSelector.ANY); + } + + private static RestHttpClient buildRestClient(NodeSelector nodeSelector) { + RestHttpClientBuilder restClientBuilder = RestHttpClient.builder(httpHosts); + if (pathPrefix.length() > 0) { + restClientBuilder.setPathPrefix((randomBoolean() ? "/" : "") + pathPrefixWithoutLeadingSlash); + } + restClientBuilder.setNodeSelector(nodeSelector); + return restClientBuilder.build(); + } + + private static HttpServer createHttpServer() throws Exception { + HttpServer httpServer = HttpServer.create(new InetSocketAddress(InetAddress.getLoopbackAddress(), 0), 0); + httpServer.start(); + // returns a different status code depending on the path + for (int statusCode : getAllStatusCodes()) { + httpServer.createContext(pathPrefix + "/" + statusCode, new ResponseHandler(statusCode)); + } + httpServer.createContext(pathPrefix + "/wait", waitForCancelHandler); + return httpServer; + } + + private static class WaitForCancelHandler implements HttpHandler { + private volatile CountDownLatch requestCameInLatch; + private volatile CountDownLatch cancelHandlerLatch; + + void reset() { + cancelHandlerLatch = new CountDownLatch(1); + requestCameInLatch = new CountDownLatch(1); + } + + void cancelDone() { + cancelHandlerLatch.countDown(); + } + + void awaitRequest() throws InterruptedException { + requestCameInLatch.await(); + } + + @Override + public void handle(HttpExchange exchange) throws IOException { + requestCameInLatch.countDown(); + try { + cancelHandlerLatch.await(); + } catch (InterruptedException ignore) {} finally { + exchange.sendResponseHeaders(200, 0); + exchange.close(); + } + } + } + + private static class ResponseHandler implements HttpHandler { + private final int statusCode; + + ResponseHandler(int statusCode) { + this.statusCode = statusCode; + } + + @Override + public void handle(HttpExchange httpExchange) throws IOException { + httpExchange.getRequestBody().close(); + httpExchange.sendResponseHeaders(statusCode, -1); + httpExchange.close(); + } + } + + @AfterClass + public static void stopHttpServers() throws IOException { + restClient.close(); + restClient = null; + for (HttpServer httpServer : httpServers) { + httpServer.stop(0); + } + httpServers = null; + } + + @Before + public void stopRandomHost() { + // verify that shutting down some hosts doesn't matter as long as one working host is left behind + if (httpServers.length > 1 && randomBoolean()) { + List updatedHttpServers = new ArrayList<>(httpServers.length - 1); + int nodeIndex = randomIntBetween(0, httpServers.length - 1); + if (0 == nodeIndex) { + stoppedFirstHost = true; + } + for (int i = 0; i < httpServers.length; i++) { + HttpServer httpServer = httpServers[i]; + if (i == nodeIndex) { + httpServer.stop(0); + } else { + updatedHttpServers.add(httpServer); + } + } + httpServers = updatedHttpServers.toArray(new HttpServer[0]); + } + } + + public void testSyncRequests() throws IOException { + int numRequests = randomIntBetween(5, 20); + for (int i = 0; i < numRequests; i++) { + final String method = RestClientTestUtil.randomHttpMethod(getRandom()); + // we don't test status codes that are subject to retries as they interfere with hosts being stopped + final int statusCode = randomBoolean() ? randomOkStatusCode(getRandom()) : randomErrorNoRetryStatusCode(getRandom()); + Response response; + try { + response = restClient.performRequest(Request.newRequest(method, "/" + statusCode).build()); + } catch (ResponseException responseException) { + response = responseException.getResponse(); + } + assertEquals(method, response.requestLine().method()); + assertEquals(statusCode, response.statusLine().statusCode()); + assertEquals((pathPrefix.length() > 0 ? pathPrefix : "") + "/" + statusCode, response.requestLine().uri()); + } + } + + public void testAsyncRequests() throws Exception { + int numRequests = randomIntBetween(5, 20); + final CountDownLatch latch = new CountDownLatch(numRequests); + final List responses = new CopyOnWriteArrayList<>(); + for (int i = 0; i < numRequests; i++) { + final String method = RestClientTestUtil.randomHttpMethod(getRandom()); + // we don't test status codes that are subject to retries as they interfere with hosts being stopped + final int statusCode = randomBoolean() ? randomOkStatusCode(getRandom()) : randomErrorNoRetryStatusCode(getRandom()); + restClient.performRequestAsync(Request.newRequest(method, "/" + statusCode).build(), new ResponseListener() { + @Override + public void onSuccess(Response response) { + responses.add(new TestResponse(method, statusCode, response)); + latch.countDown(); + } + + @Override + public void onFailure(Exception exception) { + responses.add(new TestResponse(method, statusCode, exception)); + latch.countDown(); + } + }); + } + assertTrue(latch.await(5, TimeUnit.SECONDS)); + + assertEquals(numRequests, responses.size()); + for (TestResponse testResponse : responses) { + Response response = testResponse.getResponse(); + assertEquals(testResponse.method, response.requestLine().method()); + assertEquals(testResponse.statusCode, response.statusLine().statusCode()); + assertEquals((pathPrefix.length() > 0 ? pathPrefix : "") + "/" + testResponse.statusCode, response.requestLine().uri()); + } + } + + @Ignore("https://github.com/elastic/elasticsearch/issues/45577") + public void testCancelAsyncRequests() throws Exception { + int numRequests = randomIntBetween(5, 20); + final List responses = new CopyOnWriteArrayList<>(); + final List exceptions = new CopyOnWriteArrayList<>(); + for (int i = 0; i < numRequests; i++) { + CountDownLatch latch = new CountDownLatch(1); + waitForCancelHandler.reset(); + Cancellable cancellable = restClient.performRequestAsync(Request.newRequest("GET", "/wait").build(), new ResponseListener() { + @Override + public void onSuccess(Response response) { + responses.add(response); + latch.countDown(); + } + + @Override + public void onFailure(Exception exception) { + exceptions.add(exception); + latch.countDown(); + } + }); + if (randomBoolean()) { + // we wait for the request to get to the server-side otherwise we almost always cancel + // the request artificially on the client-side before even sending it + waitForCancelHandler.awaitRequest(); + } + cancellable.cancel(); + waitForCancelHandler.cancelDone(); + assertTrue(latch.await(5, TimeUnit.SECONDS)); + } + assertEquals(0, responses.size()); + assertEquals(numRequests, exceptions.size()); + for (Exception exception : exceptions) { + assertThat(exception, instanceOf(CancellationException.class)); + } + } + + /** + * Test host selector against a real server and + * test what happens after calling + */ + public void testNodeSelector() throws Exception { + try (RestHttpClient restClient = buildRestClient(firstPositionNodeSelector())) { + Request request = Request.newRequest("GET", "/200").build(); + int rounds = between(1, 10); + for (int i = 0; i < rounds; i++) { + /* + * Run the request more than once to verify that the + * NodeSelector overrides the round robin behavior. + */ + if (stoppedFirstHost) { + try { + RestHttpClientSingleHostTests.performRequestSyncOrAsync(restClient, request); + fail("expected to fail to connect"); + } catch (ConnectException e) { + // HttpClient isn't consistent here. Sometimes the message is even null! + if (e.getMessage() != null) { + assertThat(e.getMessage(), containsString("Connection refused")); + } + } + } else { + Response response = RestHttpClientSingleHostTests.performRequestSyncOrAsync(restClient, request); + assertEquals(httpHosts[0], response.host()); + } + } + } + } + + private static class TestResponse { + private final String method; + private final int statusCode; + private final Object response; + + TestResponse(String method, int statusCode, Object response) { + this.method = method; + this.statusCode = statusCode; + this.response = response; + } + + Response getResponse() { + if (response instanceof Response) { + return (Response) response; + } + if (response instanceof ResponseException) { + return ((ResponseException) response).getResponse(); + } + throw new AssertionError("unexpected response " + response.getClass()); + } + } + + private NodeSelector firstPositionNodeSelector() { + return nodes -> { + for (Iterator itr = nodes.iterator(); itr.hasNext();) { + if (httpHosts[0] != itr.next().getHost()) { + itr.remove(); + } + } + }; + } +} diff --git a/client/rest-http-client/src/test/java/org/opensearch/internal/httpclient/RestHttpClientMultipleHostsTests.java b/client/rest-http-client/src/test/java/org/opensearch/internal/httpclient/RestHttpClientMultipleHostsTests.java new file mode 100644 index 0000000000000..2e644f427c4ee --- /dev/null +++ b/client/rest-http-client/src/test/java/org/opensearch/internal/httpclient/RestHttpClientMultipleHostsTests.java @@ -0,0 +1,346 @@ +/* + * SPDX-License-Identifier: Apache-2.0 + * + * The OpenSearch Contributors require contributions made to + * this file be licensed under the Apache-2.0 license or a + * compatible open source license. + */ + +/* + * Licensed to Elasticsearch under one or more contributor + * license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright + * ownership. Elasticsearch 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. + */ + +/* + * Modifications Copyright OpenSearch Contributors. See + * GitHub history for details. + */ + +package org.opensearch.internal.httpclient; + +import com.carrotsearch.randomizedtesting.generators.RandomNumbers; + +import org.junit.After; + +import java.io.IOException; +import java.net.http.HttpClient; +import java.util.ArrayList; +import java.util.Arrays; +import java.util.Collections; +import java.util.HashSet; +import java.util.Iterator; +import java.util.List; +import java.util.Map; +import java.util.Set; +import java.util.TreeSet; +import java.util.concurrent.ExecutorService; +import java.util.concurrent.Executors; + +import static org.opensearch.internal.httpclient.RestClientTestUtil.randomErrorNoRetryStatusCode; +import static org.opensearch.internal.httpclient.RestClientTestUtil.randomErrorRetryStatusCode; +import static org.opensearch.internal.httpclient.RestClientTestUtil.randomHttpMethod; +import static org.opensearch.internal.httpclient.RestClientTestUtil.randomOkStatusCode; +import static org.hamcrest.CoreMatchers.equalTo; +import static org.hamcrest.CoreMatchers.instanceOf; +import static org.junit.Assert.assertEquals; +import static org.junit.Assert.assertThat; +import static org.junit.Assert.assertTrue; +import static org.junit.Assert.fail; + +/** + * Tests for {@link RestHttpClient} behaviour against multiple hosts: fail-over, denylisting etc. + * Relies on a mock http client to intercept requests and return desired responses based on request path. + */ +public class RestHttpClientMultipleHostsTests extends RestHttpClientTestCase { + + private ExecutorService exec = Executors.newFixedThreadPool(1); + private List nodes; + private HostsTrackingFailureListener failureListener; + + public RestHttpClient createRestClient(NodeSelector nodeSelector) { + HttpClient httpClient = RestHttpClientSingleHostTests.mockHttpClient(exec); + int numNodes = RandomNumbers.randomIntBetween(getRandom(), 2, 5); + nodes = new ArrayList<>(numNodes); + for (int i = 0; i < numNodes; i++) { + nodes.add(new Node(new HttpHost("http", "localhost", 9200 + i))); + } + nodes = Collections.unmodifiableList(nodes); + failureListener = new HostsTrackingFailureListener(); + return new RestHttpClient(httpClient, Map.of(), nodes, null, failureListener, nodeSelector, false, false); + } + + /** + * Shutdown the executor so we don't leak threads into other test runs. + */ + @After + public void shutdownExec() { + exec.shutdown(); + } + + public void testRoundRobinOkStatusCodes() throws Exception { + RestHttpClient restClient = createRestClient(NodeSelector.ANY); + int numIters = RandomNumbers.randomIntBetween(getRandom(), 1, 5); + for (int i = 0; i < numIters; i++) { + Set hostsSet = hostsSet(); + for (int j = 0; j < nodes.size(); j++) { + int statusCode = randomOkStatusCode(getRandom()); + Response response = RestHttpClientSingleHostTests.performRequestSyncOrAsync( + restClient, + Request.newRequest(randomHttpMethod(getRandom()), "/" + statusCode).build() + ); + assertEquals(statusCode, response.statusLine().statusCode()); + assertTrue("host not found: " + response.host(), hostsSet.remove(response.host())); + } + assertEquals("every host should have been used but some weren't: " + hostsSet, 0, hostsSet.size()); + } + failureListener.assertNotCalled(); + } + + public void testRoundRobinNoRetryErrors() throws Exception { + RestHttpClient restClient = createRestClient(NodeSelector.ANY); + int numIters = RandomNumbers.randomIntBetween(getRandom(), 1, 5); + for (int i = 0; i < numIters; i++) { + Set hostsSet = hostsSet(); + for (int j = 0; j < nodes.size(); j++) { + String method = randomHttpMethod(getRandom()); + int statusCode = randomErrorNoRetryStatusCode(getRandom()); + try { + Response response = RestHttpClientSingleHostTests.performRequestSyncOrAsync( + restClient, + Request.newRequest(method, "/" + statusCode).build() + ); + if (method.equals("HEAD") && statusCode == 404) { + // no exception gets thrown although we got a 404 + assertEquals(404, response.statusLine().statusCode()); + assertEquals(statusCode, response.statusLine().statusCode()); + assertTrue("host not found: " + response.host(), hostsSet.remove(response.host())); + } else { + fail("request should have failed"); + } + } catch (ResponseException e) { + if (method.equals("HEAD") && statusCode == 404) { + throw e; + } + Response response = e.getResponse(); + assertEquals(statusCode, response.statusLine().statusCode()); + assertTrue("host not found: " + response.host(), hostsSet.remove(response.host())); + assertEquals(0, e.getSuppressed().length); + } + } + assertEquals("every host should have been used but some weren't: " + hostsSet, 0, hostsSet.size()); + } + failureListener.assertNotCalled(); + } + + public void testRoundRobinRetryErrors() throws Exception { + RestHttpClient restClient = createRestClient(NodeSelector.ANY); + String retryEndpoint = randomErrorRetryEndpoint(); + try { + RestHttpClientSingleHostTests.performRequestSyncOrAsync( + restClient, + Request.newRequest(randomHttpMethod(getRandom()), retryEndpoint).build() + ); + fail("request should have failed"); + } catch (ResponseException e) { + Set hostsSet = hostsSet(); + // first request causes all the hosts to be denylisted, the returned exception holds one suppressed exception each + failureListener.assertCalled(nodes); + do { + Response response = e.getResponse(); + assertEquals(Integer.parseInt(retryEndpoint.substring(1)), response.statusLine().statusCode()); + assertTrue("host [" + response.host() + "] not found, most likely used multiple times", hostsSet.remove(response.host())); + if (e.getSuppressed().length > 0) { + assertEquals(1, e.getSuppressed().length); + Throwable suppressed = e.getSuppressed()[0]; + assertThat(suppressed, instanceOf(ResponseException.class)); + e = (ResponseException) suppressed; + } else { + e = null; + } + } while (e != null); + assertEquals("every host should have been used but some weren't: " + hostsSet, 0, hostsSet.size()); + } catch (IOException e) { + Set hostsSet = hostsSet(); + // first request causes all the hosts to be denylisted, the returned exception holds one suppressed exception each + failureListener.assertCalled(nodes); + do { + HttpHost httpHost = HttpHost.create(e.getMessage()); + assertTrue("host [" + httpHost + "] not found, most likely used multiple times", hostsSet.remove(httpHost)); + if (e.getSuppressed().length > 0) { + assertEquals(1, e.getSuppressed().length); + Throwable suppressed = e.getSuppressed()[0]; + assertThat(suppressed, instanceOf(IOException.class)); + e = (IOException) suppressed; + } else { + e = null; + } + } while (e != null); + assertEquals("every host should have been used but some weren't: " + hostsSet, 0, hostsSet.size()); + } + + int numIters = RandomNumbers.randomIntBetween(getRandom(), 2, 5); + for (int i = 1; i <= numIters; i++) { + // check that one different host is resurrected at each new attempt + Set hostsSet = hostsSet(); + for (int j = 0; j < nodes.size(); j++) { + retryEndpoint = randomErrorRetryEndpoint(); + try { + RestHttpClientSingleHostTests.performRequestSyncOrAsync( + restClient, + Request.newRequest(randomHttpMethod(getRandom()), retryEndpoint).build() + ); + fail("request should have failed"); + } catch (ResponseException e) { + Response response = e.getResponse(); + assertThat(response.statusLine().statusCode(), equalTo(Integer.parseInt(retryEndpoint.substring(1)))); + assertTrue( + "host [" + response.host() + "] not found, most likely used multiple times", + hostsSet.remove(response.host()) + ); + // after the first request, all hosts are denylisted, a single one gets resurrected each time + failureListener.assertCalled(response.host()); + assertEquals(0, e.getSuppressed().length); + } catch (IOException e) { + HttpHost httpHost = HttpHost.create(e.getMessage()); + assertTrue("host [" + httpHost + "] not found, most likely used multiple times", hostsSet.remove(httpHost)); + // after the first request, all hosts are denylisted, a single one gets resurrected each time + failureListener.assertCalled(httpHost); + assertEquals(0, e.getSuppressed().length); + } + } + assertEquals("every host should have been used but some weren't: " + hostsSet, 0, hostsSet.size()); + if (getRandom().nextBoolean()) { + // mark one host back alive through a successful request and check that all requests after that are sent to it + HttpHost selectedHost = null; + int iters = RandomNumbers.randomIntBetween(getRandom(), 2, 10); + for (int y = 0; y < iters; y++) { + int statusCode = randomErrorNoRetryStatusCode(getRandom()); + Response response; + try { + response = RestHttpClientSingleHostTests.performRequestSyncOrAsync( + restClient, + Request.newRequest(randomHttpMethod(getRandom()), "/" + statusCode).build() + ); + } catch (ResponseException e) { + response = e.getResponse(); + } + assertThat(response.statusLine().statusCode(), equalTo(statusCode)); + if (selectedHost == null) { + selectedHost = response.host(); + } else { + assertThat(response.host(), equalTo(selectedHost)); + } + } + failureListener.assertNotCalled(); + // let the selected host catch up on number of failures, it gets selected a consecutive number of times as it's the one + // selected to be retried earlier (due to lower number of failures) till all the hosts have the same number of failures + for (int y = 0; y < i + 1; y++) { + retryEndpoint = randomErrorRetryEndpoint(); + try { + RestHttpClientSingleHostTests.performRequestSyncOrAsync( + restClient, + Request.newRequest(randomHttpMethod(getRandom()), retryEndpoint).build() + ); + fail("request should have failed"); + } catch (ResponseException e) { + Response response = e.getResponse(); + assertThat(response.statusLine().statusCode(), equalTo(Integer.parseInt(retryEndpoint.substring(1)))); + assertThat(response.host(), equalTo(selectedHost)); + failureListener.assertCalled(selectedHost); + } catch (IOException e) { + HttpHost httpHost = HttpHost.create(e.getMessage()); + assertThat(httpHost, equalTo(selectedHost)); + failureListener.assertCalled(selectedHost); + } + } + } + } + } + + public void testNodeSelector() throws Exception { + NodeSelector firstPositionOnly = restClientNodes -> { + boolean found = false; + for (Iterator itr = restClientNodes.iterator(); itr.hasNext();) { + if (nodes.get(0) == itr.next()) { + found = true; + } else { + itr.remove(); + } + } + assertTrue(found); + }; + RestHttpClient restClient = createRestClient(firstPositionOnly); + int rounds = between(1, 10); + for (int i = 0; i < rounds; i++) { + /* + * Run the request more than once to verify that the + * NodeSelector overrides the round robin behavior. + */ + Request request = Request.newRequest("GET", "/200").build(); + Response response = RestHttpClientSingleHostTests.performRequestSyncOrAsync(restClient, request); + assertEquals(nodes.get(0).getHost(), response.host()); + } + } + + public void testSetNodes() throws Exception { + RestHttpClient restClient = createRestClient(NodeSelector.SKIP_DEDICATED_CLUSTER_MANAGERS); + List newNodes = new ArrayList<>(nodes.size()); + for (int i = 0; i < nodes.size(); i++) { + Node.Roles roles = i == 0 + ? new Node.Roles(new TreeSet<>(Arrays.asList("data", "ingest"))) + : new Node.Roles(new TreeSet<>(Arrays.asList("master"))); + newNodes.add(new Node(nodes.get(i).getHost(), null, null, null, roles, null)); + } + restClient.setNodes(newNodes); + int rounds = between(1, 10); + for (int i = 0; i < rounds; i++) { + /* + * Run the request more than once to verify that the + * NodeSelector overrides the round robin behavior. + */ + Request request = Request.newRequest("GET", "/200").build(); + Response response = RestHttpClientSingleHostTests.performRequestSyncOrAsync(restClient, request); + assertEquals(newNodes.get(0).getHost(), response.host()); + } + } + + private static String randomErrorRetryEndpoint() { + switch (RandomNumbers.randomIntBetween(getRandom(), 0, 3)) { + case 0: + return "/" + randomErrorRetryStatusCode(getRandom()); + case 1: + return "/coe"; + case 2: + return "/soe"; + case 3: + return "/ioe"; + } + throw new UnsupportedOperationException(); + } + + /** + * Build a mutable {@link Set} containing all the {@link Node#getHost() hosts} + * in use by the test. + */ + private Set hostsSet() { + Set hosts = new HashSet<>(); + for (Node node : nodes) { + hosts.add(node.getHost()); + } + return hosts; + } +} diff --git a/client/rest-http-client/src/test/java/org/opensearch/internal/httpclient/RestHttpClientSingleHostIntegTests.java b/client/rest-http-client/src/test/java/org/opensearch/internal/httpclient/RestHttpClientSingleHostIntegTests.java new file mode 100644 index 0000000000000..b76081bd684e4 --- /dev/null +++ b/client/rest-http-client/src/test/java/org/opensearch/internal/httpclient/RestHttpClientSingleHostIntegTests.java @@ -0,0 +1,501 @@ +/* + * SPDX-License-Identifier: Apache-2.0 + * + * The OpenSearch Contributors require contributions made to + * this file be licensed under the Apache-2.0 license or a + * compatible open source license. + */ + +package org.opensearch.internal.httpclient; + +import com.sun.net.httpserver.Headers; +import com.sun.net.httpserver.HttpExchange; +import com.sun.net.httpserver.HttpHandler; +import com.sun.net.httpserver.HttpServer; + +import org.junit.After; +import org.junit.Before; + +import java.io.IOException; +import java.io.InputStreamReader; +import java.io.OutputStream; +import java.net.Authenticator; +import java.net.InetAddress; +import java.net.InetSocketAddress; +import java.net.PasswordAuthentication; +import java.net.http.HttpClient; +import java.net.http.HttpRequest.BodyPublishers; +import java.nio.charset.StandardCharsets; +import java.time.Duration; +import java.util.Arrays; +import java.util.Base64; +import java.util.HashMap; +import java.util.List; +import java.util.Locale; +import java.util.Map; +import java.util.Set; +import java.util.TreeSet; +import java.util.concurrent.CancellationException; +import java.util.concurrent.CopyOnWriteArrayList; +import java.util.concurrent.CountDownLatch; +import java.util.concurrent.ExecutorService; +import java.util.concurrent.Executors; +import java.util.concurrent.TimeUnit; +import java.util.concurrent.atomic.AtomicReference; +import java.util.concurrent.locks.LockSupport; + +import static org.hamcrest.CoreMatchers.equalTo; +import static org.hamcrest.CoreMatchers.nullValue; +import static org.hamcrest.CoreMatchers.startsWith; +import static org.hamcrest.Matchers.instanceOf; +import static org.junit.Assert.assertEquals; +import static org.junit.Assert.assertThat; +import static org.junit.Assert.assertThrows; +import static org.junit.Assert.assertTrue; + +/** + * Integration test to check interaction between {@link RestHttpClient} and {@link HttpClient}. + * Works against a real http server, one single host. + */ +public class RestHttpClientSingleHostIntegTests extends RestHttpClientTestCase { + + private HttpServer httpServer; + private RestHttpClient restClient; + private String pathPrefix; + private Map> defaultHeaders; + private WaitForCancelHandler waitForCancelHandler; + private ExecutorService httpClientExecutor; + private ExecutorService httpServerExecutor; + + @Before + public void startHttpServer() throws Exception { + pathPrefix = randomBoolean() ? "/testPathPrefix/" + randomAsciiLettersOfLengthBetween(1, 5) : ""; + httpServer = createHttpServer(); + httpClientExecutor = Executors.newFixedThreadPool(5); + httpServerExecutor = Executors.newFixedThreadPool(10); + defaultHeaders = RestClientTestUtil.randomHeaders(getRandom(), "Header-default"); + restClient = createRestClient(false, true); + } + + private HttpServer createHttpServer() throws Exception { + HttpServer httpServer = HttpServer.create(new InetSocketAddress(InetAddress.getLoopbackAddress(), 0), 0); + httpServer.setExecutor(httpServerExecutor); + httpServer.start(); + // returns a different status code depending on the path + for (int statusCode : RestClientTestUtil.getAllStatusCodes()) { + httpServer.createContext(pathPrefix + "/" + statusCode, new ResponseHandler(statusCode)); + } + waitForCancelHandler = new WaitForCancelHandler(); + httpServer.createContext(pathPrefix + "/wait", waitForCancelHandler); + return httpServer; + } + + private static class WaitForCancelHandler implements HttpHandler { + + private final CountDownLatch cancelHandlerLatch = new CountDownLatch(1); + + void cancelDone() { + cancelHandlerLatch.countDown(); + } + + @Override + public void handle(HttpExchange exchange) throws IOException { + try { + cancelHandlerLatch.await(); + } catch (InterruptedException ignore) {} finally { + exchange.sendResponseHeaders(200, 0); + exchange.close(); + } + } + } + + private static class ResponseHandler implements HttpHandler { + private final int statusCode; + + ResponseHandler(int statusCode) { + this.statusCode = statusCode; + } + + @Override + public void handle(HttpExchange httpExchange) throws IOException { + // copy request body to response body so we can verify it was sent + StringBuilder body = new StringBuilder(); + try (InputStreamReader reader = new InputStreamReader(httpExchange.getRequestBody(), StandardCharsets.UTF_8)) { + char[] buffer = new char[256]; + int read; + while ((read = reader.read(buffer)) != -1) { + body.append(buffer, 0, read); + } + } + // copy request headers to response headers so we can verify they were sent + Headers requestHeaders = httpExchange.getRequestHeaders(); + Headers responseHeaders = httpExchange.getResponseHeaders(); + for (Map.Entry> header : requestHeaders.entrySet()) { + responseHeaders.put(header.getKey(), header.getValue()); + } + httpExchange.getRequestBody().close(); + httpExchange.sendResponseHeaders(statusCode, body.length() == 0 ? -1 : body.length()); + if (body.length() > 0) { + try (OutputStream out = httpExchange.getResponseBody()) { + out.write(body.toString().getBytes(StandardCharsets.UTF_8)); + } + } + httpExchange.close(); + } + } + + private RestHttpClient createRestClient(final boolean useAuth, final boolean usePreemptiveAuth) { + final HttpHost httpHost = new HttpHost("http", httpServer.getAddress().getHostString(), httpServer.getAddress().getPort()); + final RestHttpClientBuilder restClientBuilder = RestHttpClient.builder(httpHost).setDefaultHeaders(defaultHeaders); + if (pathPrefix.length() > 0) { + restClientBuilder.setPathPrefix(pathPrefix); + } + + if (useAuth) { + if (usePreemptiveAuth == false) { + restClientBuilder.setHttpClientConfigCallback(new RestHttpClientBuilder.HttpClientConfigCallback() { + @Override + public HttpClient.Builder customizeHttpClient(final HttpClient.Builder httpClientBuilder) { + return httpClientBuilder.executor(httpClientExecutor).authenticator(new Authenticator() { + @Override + protected PasswordAuthentication getPasswordAuthentication() { + return new PasswordAuthentication("user", "pass".toCharArray()); + } + }); + } + }); + } else { + final String auth = Base64.getEncoder().encodeToString(("user" + ":" + "pass").getBytes(StandardCharsets.UTF_8)); + final Map> headers = new HashMap<>(defaultHeaders); + headers.put("Authorization", List.of("Basic " + auth)); + restClientBuilder.setDefaultHeaders(headers); + + restClientBuilder.setHttpClientConfigCallback(new RestHttpClientBuilder.HttpClientConfigCallback() { + @Override + public HttpClient.Builder customizeHttpClient(final HttpClient.Builder httpClientBuilder) { + return httpClientBuilder.executor(httpClientExecutor); + } + }); + } + } else { + restClientBuilder.setHttpClientConfigCallback(new RestHttpClientBuilder.HttpClientConfigCallback() { + @Override + public HttpClient.Builder customizeHttpClient(final HttpClient.Builder httpClientBuilder) { + return httpClientBuilder.executor(httpClientExecutor).connectTimeout(Duration.ofSeconds(15)); + } + }); + } + + return restClientBuilder.build(); + } + + @After + public void stopHttpServers() throws IOException, InterruptedException { + restClient.close(); + restClient = null; + httpServer.stop(0); + httpServer = null; + + httpServerExecutor.shutdown(); + if (httpServerExecutor.awaitTermination(30, TimeUnit.SECONDS) == false) { + httpServerExecutor.shutdownNow(); + } + + httpClientExecutor.shutdown(); + if (httpClientExecutor.awaitTermination(30, TimeUnit.SECONDS) == false) { + httpClientExecutor.shutdownNow(); + } + } + + /** + * Tests sending a bunch of async requests works well (e.g. no TimeoutException from the leased pool) + * See https://github.com/elastic/elasticsearch/issues/24069 + */ + public void testManyAsyncRequests() throws Exception { + int iters = randomIntBetween(500, 1000); + final CountDownLatch latch = new CountDownLatch(iters); + final List exceptions = new CopyOnWriteArrayList<>(); + for (int i = 0; i < iters; i++) { + Request request = Request.newRequest("PUT", "/200").withEntity(BodyPublishers.ofString("{}")).build(); + // Add random jitter so HttpServer will not refuse the connections + LockSupport.parkNanos(TimeUnit.MILLISECONDS.toNanos(randomLongBetween(1, 5))); + restClient.performRequestAsync(request, new ResponseListener() { + @Override + public void onSuccess(Response response) { + latch.countDown(); + } + + @Override + public void onFailure(Exception exception) { + exceptions.add(exception); + latch.countDown(); + } + }); + } + + assertTrue("timeout waiting for requests to be sent", latch.await(10, TimeUnit.SECONDS)); + if (exceptions.isEmpty() == false) { + AssertionError error = new AssertionError( + "expected no failures but got some. see suppressed for first 10 of [" + exceptions.size() + "] failures" + ); + for (Exception exception : exceptions.subList(0, Math.min(10, exceptions.size()))) { + error.addSuppressed(exception); + } + throw error; + } + } + + public void testCancelAsyncRequest() throws Exception { + Request request = Request.newRequest(RestClientTestUtil.randomHttpMethod(getRandom()), "/wait").build(); + CountDownLatch requestLatch = new CountDownLatch(1); + AtomicReference error = new AtomicReference<>(); + Cancellable cancellable = restClient.performRequestAsync(request, new ResponseListener() { + @Override + public void onSuccess(Response response) { + throw new AssertionError("onResponse called unexpectedly"); + } + + @Override + public void onFailure(Exception exception) { + error.set(exception); + requestLatch.countDown(); + } + }); + cancellable.cancel(); + waitForCancelHandler.cancelDone(); + assertTrue(requestLatch.await(5, TimeUnit.SECONDS)); + assertThat(error.get(), instanceOf(CancellationException.class)); + } + + /** + * End to end test for headers. We test it explicitly against a real http client as there are different ways + * to set/add headers to the {@link HttpClient}. + * Exercises the test http server ability to send back whatever headers it received. + */ + public void testHeaders() throws Exception { + for (String method : RestClientTestUtil.getHttpMethods()) { + final Set standardHeaders = new TreeSet<>(String.CASE_INSENSITIVE_ORDER); + standardHeaders.addAll( + Arrays.asList("Connection", "Host", "User-agent", "Date", "Upgrade", "HTTP2-Settings", "Content-Length") + ); + + final Map> requestHeaders = RestClientTestUtil.randomHeaders(getRandom(), "Header"); + final int statusCode = RestClientTestUtil.randomStatusCode(getRandom()); + Request request = Request.newRequest(method, "/" + statusCode) + .withOptions(RequestOptions.DEFAULT.toBuilder().addHeaders(requestHeaders)) + .build(); + Response esResponse; + try { + esResponse = RestHttpClientSingleHostTests.performRequestSyncOrAsync(restClient, request); + } catch (ResponseException e) { + esResponse = e.getResponse(); + } + + assertEquals(method, esResponse.requestLine().method()); + assertEquals(statusCode, esResponse.statusLine().statusCode()); + assertEquals(pathPrefix + "/" + statusCode, esResponse.requestLine().uri()); + + assertHeaders(defaultHeaders, requestHeaders, esResponse.headers(), standardHeaders); + final Set removedHeaders = new TreeSet<>(String.CASE_INSENSITIVE_ORDER); + for (final Map.Entry> responseHeader : esResponse.headers().map().entrySet()) { + String name = responseHeader.getKey().toLowerCase(Locale.ROOT); + // Some headers could be returned multiple times in response, like Connection fe. + if (name.startsWith("header") == false && removedHeaders.contains(name) == false) { + assertTrue("unknown header was returned " + name, standardHeaders.remove(name)); + removedHeaders.add(name); + } + } + assertTrue("some expected standard headers weren't returned: " + standardHeaders, standardHeaders.isEmpty()); + } + } + + /** + * End to end test for delete with body. We test it explicitly as it is not supported + * out of the box by {@link HttpClient}. + * Exercises the test http server ability to send back whatever body it received. + */ + public void testDeleteWithBody() throws Exception { + bodyTest("DELETE"); + } + + /** + * End to end test for get with body. We test it explicitly as it is not supported + * out of the box by {@link HttpClient}. + * Exercises the test http server ability to send back whatever body it received. + */ + public void testGetWithBody() throws Exception { + bodyTest("GET"); + } + + public void testEncodeParams() throws Exception { + { + Request request = Request.newRequest("PUT", "/200").withParameter("routing", "this/is/the/routing").build(); + Response response = RestHttpClientSingleHostTests.performRequestSyncOrAsync(restClient, request); + assertEquals(pathPrefix + "/200?routing=this%2Fis%2Fthe%2Frouting", response.requestLine().uri()); + } + { + Request request = Request.newRequest("PUT", "/200").withParameter("routing", "this|is|the|routing").build(); + Response response = RestHttpClientSingleHostTests.performRequestSyncOrAsync(restClient, request); + assertEquals(pathPrefix + "/200?routing=this%7Cis%7Cthe%7Crouting", response.requestLine().uri()); + } + { + Request request = Request.newRequest("PUT", "/200").withParameter("routing", "routing#1").build(); + Response response = RestHttpClientSingleHostTests.performRequestSyncOrAsync(restClient, request); + assertEquals(pathPrefix + "/200?routing=routing%231", response.requestLine().uri()); + } + { + Request request = Request.newRequest("PUT", "/200").withParameter("routing", "中文").build(); + Response response = RestHttpClientSingleHostTests.performRequestSyncOrAsync(restClient, request); + assertEquals(pathPrefix + "/200?routing=%E4%B8%AD%E6%96%87", response.requestLine().uri()); + } + { + Request request = Request.newRequest("PUT", "/200").withParameter("routing", "foo bar").build(); + Response response = RestHttpClientSingleHostTests.performRequestSyncOrAsync(restClient, request); + assertEquals(pathPrefix + "/200?routing=foo+bar", response.requestLine().uri()); + } + { + Request request = Request.newRequest("PUT", "/200").withParameter("routing", "foo+bar").build(); + Response response = RestHttpClientSingleHostTests.performRequestSyncOrAsync(restClient, request); + assertEquals(pathPrefix + "/200?routing=foo%2Bbar", response.requestLine().uri()); + } + { + Request request = Request.newRequest("PUT", "/200").withParameter("routing", "foo/bar").build(); + Response response = RestHttpClientSingleHostTests.performRequestSyncOrAsync(restClient, request); + assertEquals(pathPrefix + "/200?routing=foo%2Fbar", response.requestLine().uri()); + } + { + Request request = Request.newRequest("PUT", "/200").withParameter("routing", "foo^bar").build(); + Response response = RestHttpClientSingleHostTests.performRequestSyncOrAsync(restClient, request); + assertEquals(pathPrefix + "/200?routing=foo%5Ebar", response.requestLine().uri()); + } + { + Request request = Request.newRequest("PUT", "/200").withParameter("pretty", null).build(); + Response response = RestHttpClientSingleHostTests.performRequestSyncOrAsync(restClient, request); + assertEquals(pathPrefix + "/200?pretty", response.requestLine().uri()); + } + { + Request request = Request.newRequest("PUT", "/200").withParameter("pretty", "").build(); + Response response = RestHttpClientSingleHostTests.performRequestSyncOrAsync(restClient, request); + assertEquals(pathPrefix + "/200?pretty=", response.requestLine().uri()); + } + } + + /** + * Verify that credentials are sent on the first request with preemptive auth enabled (default when provided with credentials). + */ + public void testPreemptiveAuthEnabled() throws Exception { + final String[] methods = { "POST", "PUT", "GET", "DELETE" }; + + try (RestHttpClient restClient = createRestClient(true, true)) { + for (final String method : methods) { + final Response response = bodyTest(restClient, method); + assertThat(response.header("Authorization"), startsWith("Basic")); + } + } + } + + /** + * Verify that credentials are not sent on the first request with preemptive auth disabled. + */ + public void testPreemptiveAuthDisabled() throws Exception { + final String[] methods = { "POST", "PUT", "GET", "DELETE" }; + + try (RestHttpClient restClient = createRestClient(true, false)) { + for (final String method : methods) { + int statusCode = RestClientTestUtil.randomStatusCode(getRandom()); + if (statusCode == 401) { + final IOException ex = assertThrows(IOException.class, () -> bodyTest(restClient, method, statusCode, Map.of())); + assertThat(ex.getMessage(), equalTo("WWW-Authenticate header missing for response code 401")); + } else { + final Response response = bodyTest(restClient, method, statusCode, Map.of()); + assertThat(response.header("Authorization"), nullValue()); + } + } + } + } + + /** + * Verify that credentials continue to be sent even if a 401 (Unauthorized) response is received + */ + public void testAuthCredentialsAreNotClearedOnAuthChallenge() throws Exception { + final String[] methods = { "POST", "PUT", "GET", "DELETE" }; + + try (RestHttpClient restClient = createRestClient(true, true)) { + for (final String method : methods) { + Map> realmHeader = Map.of("WWW-Authenticate", List.of("Basic realm=\"test\"")); + final Response response401 = bodyTest(restClient, method, 401, realmHeader); + assertThat(response401.header("Authorization"), startsWith("Basic")); + + final Response response200 = bodyTest(restClient, method, 200, Map.of()); + assertThat(response200.header("Authorization"), startsWith("Basic")); + } + } + } + + public void testUrlWithoutLeadingSlash() throws Exception { + if (pathPrefix.length() == 0) { + Response response = RestHttpClientSingleHostTests.performRequestSyncOrAsync( + restClient, + Request.newRequest("GET", "200").build() + ); + // a trailing slash gets automatically added even if a pathPrefix is not configured (HttpClient uses full URI) + assertEquals(200, response.statusLine().statusCode()); + } else { + { + Response response = RestHttpClientSingleHostTests.performRequestSyncOrAsync( + restClient, + Request.newRequest("GET", "200").build() + ); + // a trailing slash gets automatically added if a pathPrefix is configured + assertEquals(200, response.statusLine().statusCode()); + } + { + // pathPrefix is not required to start with '/', will be added automatically + try ( + RestHttpClient restClient = RestHttpClient.builder( + new HttpHost("http", httpServer.getAddress().getHostString(), httpServer.getAddress().getPort()) + ).setPathPrefix(pathPrefix.substring(1)).build() + ) { + Response response = RestHttpClientSingleHostTests.performRequestSyncOrAsync( + restClient, + Request.newRequest("GET", "200").build() + ); + // a trailing slash gets automatically added if a pathPrefix is configured + assertEquals(200, response.statusLine().statusCode()); + } + } + } + } + + private Response bodyTest(final String method) throws Exception { + return bodyTest(restClient, method); + } + + private Response bodyTest(final RestHttpClient restClient, final String method) throws Exception { + int statusCode = RestClientTestUtil.randomStatusCode(getRandom()); + return bodyTest(restClient, method, statusCode, Map.of()); + } + + private Response bodyTest(RestHttpClient restClient, String method, int statusCode, Map> headers) + throws Exception { + String requestBody = "{ \"field\": \"value\" }"; + Request.Builder builder = Request.newRequest(method, "/" + statusCode).withEntity(requestBody); + RequestOptions.Builder options = RequestOptions.DEFAULT.toBuilder(); + for (Map.Entry> header : headers.entrySet()) { + header.getValue().forEach(v -> options.addHeader(header.getKey(), v)); + } + builder = builder.withOptions(options); + Response esResponse; + try { + esResponse = RestHttpClientSingleHostTests.performRequestSyncOrAsync(restClient, builder.build()); + } catch (ResponseException e) { + esResponse = e.getResponse(); + } + assertEquals(method, esResponse.requestLine().method()); + assertEquals(statusCode, esResponse.statusLine().statusCode()); + assertEquals(pathPrefix + "/" + statusCode, esResponse.requestLine().uri()); + assertEquals(requestBody, BodyUtils.getBodyAsString(esResponse)); + + return esResponse; + } +} diff --git a/client/rest-http-client/src/test/java/org/opensearch/internal/httpclient/RestHttpClientSingleHostStreamingIntegTests.java b/client/rest-http-client/src/test/java/org/opensearch/internal/httpclient/RestHttpClientSingleHostStreamingIntegTests.java new file mode 100644 index 0000000000000..6439826eac37b --- /dev/null +++ b/client/rest-http-client/src/test/java/org/opensearch/internal/httpclient/RestHttpClientSingleHostStreamingIntegTests.java @@ -0,0 +1,180 @@ +/* + * SPDX-License-Identifier: Apache-2.0 + * + * The OpenSearch Contributors require contributions made to + * this file be licensed under the Apache-2.0 license or a + * compatible open source license. + */ + +package org.opensearch.internal.httpclient; + +import com.sun.net.httpserver.Headers; +import com.sun.net.httpserver.HttpExchange; +import com.sun.net.httpserver.HttpHandler; +import com.sun.net.httpserver.HttpServer; + +import org.junit.After; +import org.junit.Before; + +import java.io.IOException; +import java.net.InetAddress; +import java.net.InetSocketAddress; +import java.net.http.HttpClient; +import java.nio.CharBuffer; +import java.nio.charset.StandardCharsets; +import java.time.Duration; +import java.util.List; +import java.util.Map; +import java.util.concurrent.ExecutorService; +import java.util.concurrent.Executors; +import java.util.concurrent.TimeUnit; + +import reactor.core.publisher.Flux; +import reactor.core.publisher.Mono; +import reactor.test.StepVerifier; + +import static org.junit.Assert.assertEquals; + +/** + * Integration test to check interaction between {@link RestHttpClient} and {@link HttpClient}. + * Works against a real http server, one single host, use streaming. + */ +public class RestHttpClientSingleHostStreamingIntegTests extends RestHttpClientTestCase { + + private HttpServer httpServer; + private RestHttpClient restClient; + private String pathPrefix; + private Map> defaultHeaders; + private ExecutorService httpClientExecutor; + + @Before + public void startHttpServer() throws Exception { + pathPrefix = randomBoolean() ? "/testPathPrefix/" + randomAsciiLettersOfLengthBetween(1, 5) : ""; + httpServer = createHttpServer(); + httpClientExecutor = Executors.newWorkStealingPool(); + defaultHeaders = RestClientTestUtil.randomHeaders(getRandom(), "Header-default"); + restClient = createRestClient(); + } + + private HttpServer createHttpServer() throws Exception { + HttpServer httpServer = HttpServer.create(new InetSocketAddress(InetAddress.getLoopbackAddress(), 0), 0); + httpServer.start(); + // returns a different status code depending on the path + for (int statusCode : RestClientTestUtil.getAllStatusCodes()) { + httpServer.createContext(pathPrefix + "/" + statusCode, new ResponseHandler(statusCode)); + } + return httpServer; + } + + private static class ResponseHandler implements HttpHandler { + private final int statusCode; + + ResponseHandler(int statusCode) { + this.statusCode = statusCode; + } + + @Override + public void handle(HttpExchange httpExchange) throws IOException { + Headers requestHeaders = httpExchange.getRequestHeaders(); + Headers responseHeaders = httpExchange.getResponseHeaders(); + for (Map.Entry> header : requestHeaders.entrySet()) { + responseHeaders.put(header.getKey(), header.getValue()); + } + try { + httpExchange.sendResponseHeaders(statusCode, 0); + httpExchange.getRequestBody().transferTo(httpExchange.getResponseBody()); + httpExchange.getResponseBody().flush(); + } finally { + httpExchange.getRequestBody().close(); + httpExchange.getResponseBody().close(); + httpExchange.close(); + } + } + } + + private RestHttpClient createRestClient() { + final HttpHost httpHost = new HttpHost("http", httpServer.getAddress().getHostString(), httpServer.getAddress().getPort()); + final RestHttpClientBuilder restClientBuilder = RestHttpClient.builder(httpHost) + .setDefaultHeaders(defaultHeaders) + .setCompressionEnabled(randomBoolean()); + if (pathPrefix.length() > 0) { + restClientBuilder.setPathPrefix(pathPrefix); + } + + restClientBuilder.setHttpClientConfigCallback(new RestHttpClientBuilder.HttpClientConfigCallback() { + @Override + public HttpClient.Builder customizeHttpClient(final HttpClient.Builder httpClientBuilder) { + return httpClientBuilder.executor(httpClientExecutor).connectTimeout(Duration.ofSeconds(15)); + } + }); + + return restClientBuilder.build(); + } + + @After + public void stopHttpServers() throws IOException, InterruptedException { + restClient.close(); + restClient = null; + httpServer.stop(0); + httpServer = null; + + httpClientExecutor.shutdown(); + if (httpClientExecutor.awaitTermination(30, TimeUnit.SECONDS) == false) { + httpClientExecutor.shutdownNow(); + } + } + + /** + * End to end test for delete with body. We test it explicitly as it is not supported + * out of the box by {@link HttpClient}. + * Exercises the test http server ability to send back whatever body it received. + */ + public void testDeleteWithBody() throws Exception { + bodyTest("DELETE"); + } + + /** + * End to end test for get with body. We test it explicitly as it is not supported + * out of the box by {@link HttpClient}. + * Exercises the test http server ability to send back whatever body it received. + */ + public void testGetWithBody() throws Exception { + bodyTest("GET"); + } + + private StreamingResponse bodyTest(final String method) throws Exception { + int statusCode = RestClientTestUtil.randomStatusCode(getRandom()); + return bodyTest(restClient, method, statusCode, Map.of()); + } + + private StreamingResponse bodyTest(RestHttpClient restClient, String method, int statusCode, Map> headers) + throws Exception { + String requestBody = "{ \"field\": \"value\" }"; + RequestOptions.Builder options = RequestOptions.DEFAULT.toBuilder(); + for (Map.Entry> header : headers.entrySet()) { + header.getValue().forEach(v -> options.addHeader(header.getKey(), v)); + } + + final StreamingRequest request = StreamingRequest.newRequest( + method, + "/" + statusCode, + Mono.just(StandardCharsets.UTF_8.encode(requestBody)) + ).withOptions(options).build(); + StreamingResponse esResponse = restClient.streamRequest(request); + + assertEquals(method, esResponse.requestLine().method()); + assertEquals(statusCode, esResponse.statusLine().statusCode()); + assertEquals(pathPrefix + "/" + statusCode, esResponse.requestLine().uri()); + + if (statusCode >= 200 && statusCode < 400) { + StepVerifier.create(Flux.from(esResponse.body()).map(StandardCharsets.UTF_8::decode).map(CharBuffer::toString)) + .expectNextMatches(s -> s.equals(requestBody)) + .expectComplete() + .verify(Duration.ofSeconds(5)); + } else { + StepVerifier.create(Flux.from(esResponse.body())).expectError(ResponseException.class).verify(Duration.ofSeconds(5)); + } + + return esResponse; + } +} diff --git a/client/rest-http-client/src/test/java/org/opensearch/internal/httpclient/RestHttpClientSingleHostTests.java b/client/rest-http-client/src/test/java/org/opensearch/internal/httpclient/RestHttpClientSingleHostTests.java new file mode 100644 index 0000000000000..eb8a16d257cc0 --- /dev/null +++ b/client/rest-http-client/src/test/java/org/opensearch/internal/httpclient/RestHttpClientSingleHostTests.java @@ -0,0 +1,728 @@ +/* + * SPDX-License-Identifier: Apache-2.0 + * + * The OpenSearch Contributors require contributions made to + * this file be licensed under the Apache-2.0 license or a + * compatible open source license. + */ + +package org.opensearch.internal.httpclient; + +import org.apache.commons.logging.Log; +import org.apache.commons.logging.LogFactory; +import org.junit.After; +import org.junit.Before; + +import javax.net.ssl.SSLContext; +import javax.net.ssl.SSLHandshakeException; +import javax.net.ssl.SSLParameters; +import javax.net.ssl.SSLSession; + +import java.io.IOException; +import java.io.PrintWriter; +import java.io.StringWriter; +import java.net.Authenticator; +import java.net.CookieHandler; +import java.net.ProxySelector; +import java.net.SocketTimeoutException; +import java.net.URI; +import java.net.URISyntaxException; +import java.net.URLEncoder; +import java.net.http.HttpClient; +import java.net.http.HttpClient.Version; +import java.net.http.HttpHeaders; +import java.net.http.HttpRequest; +import java.net.http.HttpRequest.BodyPublisher; +import java.net.http.HttpRequest.BodyPublishers; +import java.net.http.HttpResponse; +import java.net.http.HttpResponse.BodyHandler; +import java.net.http.HttpResponse.PushPromiseHandler; +import java.net.http.HttpTimeoutException; +import java.nio.channels.ClosedChannelException; +import java.nio.charset.StandardCharsets; +import java.time.Duration; +import java.util.Arrays; +import java.util.Collections; +import java.util.HashMap; +import java.util.HashSet; +import java.util.List; +import java.util.Map; +import java.util.Optional; +import java.util.Set; +import java.util.TreeSet; +import java.util.concurrent.CompletableFuture; +import java.util.concurrent.CountDownLatch; +import java.util.concurrent.Executor; +import java.util.concurrent.ExecutorService; +import java.util.concurrent.Executors; +import java.util.concurrent.TimeUnit; +import java.util.concurrent.atomic.AtomicReference; +import java.util.concurrent.atomic.LongAdder; +import java.util.stream.Collectors; + +import static java.util.Collections.singletonList; +import static org.opensearch.internal.httpclient.RestClientTestUtil.getAllErrorStatusCodes; +import static org.opensearch.internal.httpclient.RestClientTestUtil.getHttpMethods; +import static org.opensearch.internal.httpclient.RestClientTestUtil.getOkStatusCodes; +import static org.opensearch.internal.httpclient.RestClientTestUtil.randomStatusCode; +import static org.hamcrest.CoreMatchers.containsString; +import static org.hamcrest.CoreMatchers.equalTo; +import static org.hamcrest.CoreMatchers.instanceOf; +import static org.junit.Assert.assertEquals; +import static org.junit.Assert.assertFalse; +import static org.junit.Assert.assertThat; +import static org.junit.Assert.fail; + +/** + * Tests for basic functionality of {@link RestHttpClient} against one single host: tests http requests being sent, headers, + * body, different status codes and corresponding responses/exceptions. + * Relies on a mock http client to intercept requests and return desired responses based on request path. + */ +public class RestHttpClientSingleHostTests extends RestHttpClientTestCase { + private static final Log logger = LogFactory.getLog(RestHttpClientSingleHostTests.class); + + private ExecutorService exec = Executors.newFixedThreadPool(1); + private RestHttpClient restClient; + private Map> defaultHeaders; + private Node node; + private HttpClient httpClient; + private HostsTrackingFailureListener failureListener; + private boolean strictDeprecationMode; + private LongAdder requests; + private AtomicReference requestProducerCapture; + + @Before + public void createRestClient() { + requests = new LongAdder(); + requestProducerCapture = new AtomicReference<>(); + httpClient = mockHttpClient(exec, (requestProducer, bodyHandler) -> { + requests.increment(); + requestProducerCapture.set(requestProducer); + }); + defaultHeaders = RestClientTestUtil.randomHeaders(getRandom(), "Header-default"); + node = new Node(new HttpHost("http", "localhost", 9200)); + failureListener = new HostsTrackingFailureListener(); + strictDeprecationMode = randomBoolean(); + restClient = new RestHttpClient( + this.httpClient, + defaultHeaders, + singletonList(node), + null, + failureListener, + NodeSelector.ANY, + strictDeprecationMode, + false + ); + } + + interface HttpClientListener { + void onExecute(HttpRequest requestProducer, BodyHandler bodyHandler); + } + + @SuppressWarnings("unchecked") + static HttpClient mockHttpClient(final ExecutorService exec, final HttpClientListener... listeners) { + HttpClient httpClient = new HttpClient() { + @Override + public Optional authenticator() { + return Optional.empty(); + } + + @Override + public Optional connectTimeout() { + return Optional.empty(); + } + + @Override + public Optional cookieHandler() { + return Optional.empty(); + } + + @Override + public Optional executor() { + return Optional.of(exec); + } + + @Override + public Redirect followRedirects() { + return null; + } + + @Override + public Optional proxy() { + return Optional.empty(); + } + + @Override + public SSLContext sslContext() { + return null; + } + + @Override + public SSLParameters sslParameters() { + return null; + } + + @Override + public Version version() { + return Version.HTTP_1_1; + } + + @Override + public boolean awaitTermination(Duration duration) throws InterruptedException { + return exec.awaitTermination(duration.toMillis(), TimeUnit.MILLISECONDS); + } + + @Override + public HttpResponse send(java.net.http.HttpRequest request, BodyHandler responseBodyHandler) throws IOException, + InterruptedException { + return sendAsync(request, responseBodyHandler).join(); + } + + @Override + public CompletableFuture> sendAsync( + java.net.http.HttpRequest request, + BodyHandler responseBodyHandler, + PushPromiseHandler pushPromiseHandler + ) { + return sendAsync(request, responseBodyHandler); + } + + @Override + public CompletableFuture> sendAsync(HttpRequest request, BodyHandler responseBodyHandler) { + Arrays.stream(listeners).forEach(l -> l.onExecute(request, responseBodyHandler)); + + final CompletableFuture> callback = new CompletableFuture<>(); + exec.submit(() -> { + try { + HttpResponse httpResponse = responseOrException(request); + callback.complete((HttpResponse) httpResponse); + return (T) httpResponse; + } catch (Exception e) { + callback.completeExceptionally(e); + return null; + } + }); + + return callback; + } + }; + + return httpClient; + } + + private static HttpResponse responseOrException(HttpRequest request) throws Exception { + final HttpHost httpHost = new HttpHost(request.uri().getScheme(), request.uri().getHost(), request.uri().getPort()); + // return the desired status code or exception depending on the path + switch (request.uri().getPath()) { + case "/soe": + throw new SocketTimeoutException(httpHost.toString()); + case "/coe": + throw new HttpTimeoutException(httpHost.toString()); + case "/ioe": + throw new IOException(httpHost.toString()); + case "/closed": + throw new ClosedChannelException(); + case "/handshake": + throw new SSLHandshakeException(""); + case "/uri": + throw new URISyntaxException("", ""); + case "/runtime": + throw new RuntimeException(); + default: + int statusCode = Integer.parseInt(request.uri().getPath().substring(1)); + + // return the same body that was sent + final Object entity = BodyUtils.getBody(request).map(List::of).cache().block(); + final HttpResponse httpResponse = new HttpResponse<>() { + @Override + public int statusCode() { + return statusCode; + } + + @Override + public HttpRequest request() { + return request; + } + + @Override + public Optional> previousResponse() { + return Optional.empty(); + } + + @Override + public HttpHeaders headers() { + return request.headers(); + } + + @Override + public Object body() { + return entity; + } + + @Override + public Optional sslSession() { + return Optional.empty(); + } + + @Override + public URI uri() { + return request.uri(); + } + + @Override + public Version version() { + return request.version().orElse(Version.HTTP_1_1); + } + }; + + return httpResponse; + } + } + + /** + * Shutdown the executor so we don't leak threads into other test runs. + */ + @After + public void shutdownExec() { + exec.shutdown(); + } + + /** + * Verifies the content of the {@link HttpRequest} that's internally created and passed through to the http client + */ + @SuppressWarnings("unchecked") + public void testInternalHttpRequest() throws Exception { + int times = 0; + for (String httpMethod : getHttpMethods()) { + HttpRequest expectedRequest = performRandomRequest(httpMethod); + assertThat(requests.intValue(), equalTo(++times)); + + HttpRequest actualRequest = requestProducerCapture.get(); + assertEquals(expectedRequest.uri(), actualRequest.uri()); + assertEquals(expectedRequest.method(), actualRequest.method()); + assertEquals(expectedRequest.headers(), actualRequest.headers()); + + Object expectedEntity = BodyUtils.getBody(expectedRequest).block(); + if (expectedEntity != null) { + Object actualEntity = BodyUtils.getBody(actualRequest).block(); + assertEquals(expectedEntity, actualEntity); + } + } + } + + /** + * End to end test for ok status codes + */ + public void testOkStatusCodes() throws Exception { + for (String method : getHttpMethods()) { + for (int okStatusCode : getOkStatusCodes()) { + Response response = performRequestSyncOrAsync(restClient, Request.newRequest(method, "/" + okStatusCode).build()); + assertThat(response.statusLine().statusCode(), equalTo(okStatusCode)); + } + } + failureListener.assertNotCalled(); + } + + /** + * End to end test for error status codes: they should cause an exception to be thrown, apart from 404 with HEAD requests + */ + public void testErrorStatusCodes() throws Exception { + for (String method : getHttpMethods()) { + Set expectedIgnores = new HashSet<>(); + String ignoreParam = ""; + if ("HEAD".equalsIgnoreCase(method)) { + expectedIgnores.add(404); + } + if (randomBoolean()) { + int numIgnores = randomIntBetween(1, 3); + for (int i = 0; i < numIgnores; i++) { + Integer code = randomFrom(getAllErrorStatusCodes()); + expectedIgnores.add(code); + ignoreParam += code; + if (i < numIgnores - 1) { + ignoreParam += ","; + } + } + } + // error status codes should cause an exception to be thrown + for (int errorStatusCode : getAllErrorStatusCodes()) { + try { + Request.Builder builder = Request.newRequest(method, "/" + errorStatusCode); + if (false == ignoreParam.isEmpty()) { + builder = builder.withParameter("ignore", ignoreParam); + } + Response response = restClient.performRequest(builder.build()); + if (expectedIgnores.contains(errorStatusCode)) { + // no exception gets thrown although we got an error status code, as it was configured to be ignored + assertEquals(errorStatusCode, response.statusLine().statusCode()); + } else { + fail("request should have failed"); + } + } catch (ResponseException e) { + if (expectedIgnores.contains(errorStatusCode)) { + throw e; + } + assertEquals(errorStatusCode, e.getResponse().statusLine().statusCode()); + assertExceptionStackContainsCallingMethod(e); + } + if (errorStatusCode <= 500 || expectedIgnores.contains(errorStatusCode)) { + failureListener.assertNotCalled(); + } else { + failureListener.assertCalled(singletonList(node)); + } + } + } + } + + public void testPerformRequestIOExceptions() throws Exception { + for (String method : getHttpMethods()) { + // IOExceptions should be let bubble up + try { + restClient.performRequest(Request.newRequest(method, "/ioe").build()); + fail("request should have failed"); + } catch (IOException e) { + // And we do all that so the thrown exception has our method in the stacktrace + assertExceptionStackContainsCallingMethod(e); + } + failureListener.assertCalled(singletonList(node)); + try { + restClient.performRequest(Request.newRequest(method, "/coe").build()); + fail("request should have failed"); + } catch (HttpTimeoutException e) { + // And we do all that so the thrown exception has our method in the stacktrace + assertExceptionStackContainsCallingMethod(e); + } + failureListener.assertCalled(singletonList(node)); + try { + restClient.performRequest(Request.newRequest(method, "/soe").build()); + fail("request should have failed"); + } catch (SocketTimeoutException e) { + // And we do all that so the thrown exception has our method in the stacktrace + assertExceptionStackContainsCallingMethod(e); + } + failureListener.assertCalled(singletonList(node)); + try { + restClient.performRequest(Request.newRequest(method, "/closed").build()); + fail("request should have failed"); + } catch (ClosedChannelException e) { + // And we do all that so the thrown exception has our method in the stacktrace + assertExceptionStackContainsCallingMethod(e); + } + failureListener.assertCalled(singletonList(node)); + try { + restClient.performRequest(Request.newRequest(method, "/handshake").build()); + fail("request should have failed"); + } catch (SSLHandshakeException e) { + // And we do all that so the thrown exception has our method in the stacktrace + assertExceptionStackContainsCallingMethod(e); + } + failureListener.assertCalled(singletonList(node)); + } + } + + public void testPerformRequestRuntimeExceptions() throws Exception { + for (String method : getHttpMethods()) { + try { + restClient.performRequest(Request.newRequest(method, "/runtime").build()); + fail("request should have failed"); + } catch (RuntimeException e) { + // And we do all that so the thrown exception has our method in the stacktrace + assertExceptionStackContainsCallingMethod(e); + } + failureListener.assertCalled(singletonList(node)); + } + } + + public void testPerformRequestExceptions() throws Exception { + for (String method : getHttpMethods()) { + try { + restClient.performRequest(Request.newRequest(method, "/uri").build()); + fail("request should have failed"); + } catch (RuntimeException e) { + assertThat(e.getCause(), instanceOf(URISyntaxException.class)); + // And we do all that so the thrown exception has our method in the stacktrace + assertExceptionStackContainsCallingMethod(e); + } + failureListener.assertCalled(singletonList(node)); + } + } + + /** + * End to end test for request and response body. Exercises the mock http client ability to send back + * whatever body it has received. + */ + public void testBody() throws Exception { + String body = "{ \"field\": \"value\" }"; + BodyPublisher entity = BodyPublishers.ofString(body); + for (String method : Arrays.asList("DELETE", "GET", "PATCH", "POST", "PUT", "TRACE")) { + for (int okStatusCode : getOkStatusCodes()) { + Request request = Request.newRequest(method, "/" + okStatusCode).withEntity(entity).build(); + Response response = restClient.performRequest(request); + assertThat(response.statusLine().statusCode(), equalTo(okStatusCode)); + assertThat(BodyUtils.getBodyAsString(response), equalTo(body)); + } + for (int errorStatusCode : getAllErrorStatusCodes()) { + Request request = Request.newRequest(method, "/" + errorStatusCode).withEntity(entity).build(); + try { + restClient.performRequest(request); + fail("request should have failed"); + } catch (ResponseException e) { + Response response = e.getResponse(); + assertThat(response.statusLine().statusCode(), equalTo(errorStatusCode)); + assertThat(BodyUtils.getBodyAsString(response), equalTo(body)); + assertExceptionStackContainsCallingMethod(e); + } + } + } + } + + /** + * End to end test for request and response headers. Exercises the mock http client ability to send back + * whatever headers it has received. + */ + public void testHeaders() throws Exception { + for (String method : getHttpMethods()) { + final Map> requestHeaders = RestClientTestUtil.randomHeaders(getRandom(), "Header"); + final int statusCode = randomStatusCode(getRandom()); + Request.Builder builder = Request.newRequest(method, "/" + statusCode); + RequestOptions.Builder options = RequestOptions.DEFAULT.toBuilder(); + for (Map.Entry> requestHeader : requestHeaders.entrySet()) { + requestHeader.getValue().forEach(v -> options.addHeader(requestHeader.getKey(), v)); + } + builder = builder.withOptions(options); + Response esResponse; + try { + esResponse = performRequestSyncOrAsync(restClient, builder.build()); + } catch (ResponseException e) { + esResponse = e.getResponse(); + } + assertThat(esResponse.statusLine().statusCode(), equalTo(statusCode)); + assertHeaders(defaultHeaders, requestHeaders, esResponse.headers(), Collections.emptySet()); + assertFalse(esResponse.hasWarnings()); + } + } + + public void testDeprecationWarnings() throws Exception { + String chars = randomAsciiAlphanumOfLength(5); + assertDeprecationWarnings(singletonList("poorly formatted " + chars), singletonList("poorly formatted " + chars)); + assertDeprecationWarnings(singletonList(formatWarningWithoutDate(chars)), singletonList(chars)); + assertDeprecationWarnings(singletonList(formatWarning(chars)), singletonList(chars)); + assertDeprecationWarnings( + Arrays.asList(formatWarning(chars), "another one", "and another"), + Arrays.asList(chars, "another one", "and another") + ); + assertDeprecationWarnings(Arrays.asList("ignorable one", "and another"), Arrays.asList("ignorable one", "and another")); + assertDeprecationWarnings(singletonList("exact"), singletonList("exact")); + assertDeprecationWarnings(Collections.emptyList(), Collections.emptyList()); + + String proxyWarning = "112 - \"network down\" \"Sat, 25 Aug 2012 23:34:45 GMT\""; + assertDeprecationWarnings(singletonList(proxyWarning), singletonList(proxyWarning)); + } + + private enum DeprecationWarningOption { + PERMISSIVE { + protected WarningsHandler warningsHandler() { + return WarningsHandler.PERMISSIVE; + } + }, + STRICT { + protected WarningsHandler warningsHandler() { + return WarningsHandler.STRICT; + } + }, + FILTERED { + protected WarningsHandler warningsHandler() { + return new WarningsHandler() { + @Override + public boolean warningsShouldFailRequest(List warnings) { + for (String warning : warnings) { + if (false == warning.startsWith("ignorable")) { + return true; + } + } + return false; + } + }; + } + }, + EXACT { + protected WarningsHandler warningsHandler() { + return new WarningsHandler() { + @Override + public boolean warningsShouldFailRequest(List warnings) { + return false == warnings.equals(Arrays.asList("exact")); + } + }; + } + }; + + protected abstract WarningsHandler warningsHandler(); + } + + private void assertDeprecationWarnings(List warningHeaderTexts, List warningBodyTexts) throws Exception { + String method = randomFrom(getHttpMethods()); + Request.Builder builder = Request.newRequest(method, "/200"); + RequestOptions.Builder options = RequestOptions.DEFAULT.toBuilder(); + for (String warningHeaderText : warningHeaderTexts) { + options.addHeader("Warning", warningHeaderText); + } + + final boolean expectFailure; + if (randomBoolean()) { + logger.info("checking strictWarningsMode=[" + strictDeprecationMode + "] and warnings=" + warningBodyTexts); + expectFailure = strictDeprecationMode && false == warningBodyTexts.isEmpty(); + } else { + DeprecationWarningOption warningOption = randomFrom(DeprecationWarningOption.values()); + logger.info("checking warningOption=" + warningOption + " and warnings=" + warningBodyTexts); + options.setWarningsHandler(warningOption.warningsHandler()); + expectFailure = warningOption.warningsHandler().warningsShouldFailRequest(warningBodyTexts); + } + builder = builder.withOptions(options); + + Response response; + if (expectFailure) { + try { + performRequestSyncOrAsync(restClient, builder.build()); + fail("expected WarningFailureException from warnings"); + return; + } catch (WarningFailureException e) { + if (false == warningBodyTexts.isEmpty()) { + assertThat(e.getMessage(), containsString("\nWarnings: " + warningBodyTexts)); + } + response = e.getResponse(); + } + } else { + response = performRequestSyncOrAsync(restClient, builder.build()); + } + assertEquals(false == warningBodyTexts.isEmpty(), response.hasWarnings()); + assertEquals(warningBodyTexts, response.warnings()); + } + + /** + * Emulates OpenSearch's HeaderWarningLogger.formatWarning in simple + * cases. We don't have that available because we're testing against 1.7. + */ + private static String formatWarningWithoutDate(String warningBody) { + final String hash = new String(new byte[40], StandardCharsets.UTF_8).replace('\0', 'e'); + return "299 OpenSearch-1.2.2-SNAPSHOT-" + hash + " \"" + warningBody + "\""; + } + + private static String formatWarning(String warningBody) { + return formatWarningWithoutDate(warningBody) + " \"Mon, 01 Jan 2001 00:00:00 GMT\""; + } + + private HttpRequest performRandomRequest(String method) throws Exception { + String uriAsString = "/" + randomStatusCode(getRandom()); + Request.Builder builder = Request.newRequest(method, uriAsString); + + Map params = new HashMap<>(); + if (randomBoolean()) { + int numParams = randomIntBetween(1, 3); + for (int i = 0; i < numParams; i++) { + String name = "param-" + i; + String value = randomAsciiAlphanumOfLengthBetween(3, 10); + builder = builder.withParameter(name, value); + params.put(name, value); + } + } + if (randomBoolean()) { + // randomly add some ignore parameter, which doesn't get sent as part of the request + String ignore = Integer.toString(randomFrom(RestClientTestUtil.getAllErrorStatusCodes())); + if (randomBoolean()) { + ignore += "," + randomFrom(RestClientTestUtil.getAllErrorStatusCodes()); + } + builder = builder.withParameter("ignore", ignore); + } + + final String additionalQuery = params.entrySet() + .stream() + .map(e -> e.getKey() + "=" + URLEncoder.encode(e.getValue(), StandardCharsets.UTF_8)) + .collect(Collectors.joining("&")); + + URI uri = new URI("http", null, "localhost", 9200, uriAsString, additionalQuery, null); + BodyPublisher bodyPublisher = BodyPublishers.noBody(); + if (getRandom().nextBoolean()) { + bodyPublisher = BodyPublishers.ofString(randomAsciiAlphanumOfLengthBetween(10, 100)); + builder = builder.withEntity(bodyPublisher); + } + + HttpRequest.Builder expectedRequest = HttpRequest.newBuilder(uri).method(method, bodyPublisher); + final Set uniqueNames = new TreeSet<>(String.CASE_INSENSITIVE_ORDER); + if (randomBoolean()) { + Map> headers = RestClientTestUtil.randomHeaders(getRandom(), "Header"); + RequestOptions.Builder options = RequestOptions.DEFAULT.toBuilder(); + options.addHeaders(headers); + for (Map.Entry> header : headers.entrySet()) { + header.getValue().forEach(v -> expectedRequest.header(header.getKey(), v)); + uniqueNames.add(header.getKey()); + } + builder = builder.withOptions(options); + } + for (Map.Entry> defaultHeader : defaultHeaders.entrySet()) { + // request level headers override default headers + if (uniqueNames.contains(defaultHeader.getKey()) == false) { + defaultHeader.getValue().forEach(v -> expectedRequest.header(defaultHeader.getKey(), v)); + } + } + try { + performRequestSyncOrAsync(restClient, builder.build()); + } catch (Exception e) { + // all good + } + return expectedRequest.build(); + } + + static Response performRequestSyncOrAsync(RestHttpClient restClient, Request request) throws Exception { + // randomize between sync and async methods + if (randomBoolean()) { + return restClient.performRequest(request); + } else { + final AtomicReference exceptionRef = new AtomicReference<>(); + final AtomicReference responseRef = new AtomicReference<>(); + final CountDownLatch latch = new CountDownLatch(1); + restClient.performRequestAsync(request, new ResponseListener() { + @Override + public void onSuccess(Response response) { + responseRef.set(response); + latch.countDown(); + + } + + @Override + public void onFailure(Exception exception) { + exceptionRef.set(exception); + latch.countDown(); + } + }); + latch.await(); + if (exceptionRef.get() != null) { + throw exceptionRef.get(); + } + return responseRef.get(); + } + } + + /** + * Asserts that the provided {@linkplain Exception} contains the method + * that called this somewhere on its stack. This is + * normally the case for synchronous calls but {@link RestHttpClient} performs + * synchronous calls by performing asynchronous calls and blocking the + * current thread until the call returns so it has to take special care + * to make sure that the caller shows up in the exception. We use this + * assertion to make sure that we don't break that "special care". + */ + private static void assertExceptionStackContainsCallingMethod(Throwable t) { + // 0 is getStackTrace + // 1 is this method + // 2 is the caller, what we want + StackTraceElement myMethod = Thread.currentThread().getStackTrace()[2]; + for (StackTraceElement se : t.getStackTrace()) { + if (se.getClassName().equals(myMethod.getClassName()) && se.getMethodName().equals(myMethod.getMethodName())) { + return; + } + } + StringWriter stack = new StringWriter(); + t.printStackTrace(new PrintWriter(stack)); + fail("didn't find the calling method (looks like " + myMethod + ") in:\n" + stack); + } +} diff --git a/client/rest-http-client/src/test/java/org/opensearch/internal/httpclient/RestHttpClientTestCase.java b/client/rest-http-client/src/test/java/org/opensearch/internal/httpclient/RestHttpClientTestCase.java new file mode 100644 index 0000000000000..e7ddb131b4205 --- /dev/null +++ b/client/rest-http-client/src/test/java/org/opensearch/internal/httpclient/RestHttpClientTestCase.java @@ -0,0 +1,99 @@ +/* + * SPDX-License-Identifier: Apache-2.0 + * + * The OpenSearch Contributors require contributions made to + * this file be licensed under the Apache-2.0 license or a + * compatible open source license. + */ + +package org.opensearch.internal.httpclient; + +import com.carrotsearch.randomizedtesting.JUnit3MethodProvider; +import com.carrotsearch.randomizedtesting.MixWithSuiteName; +import com.carrotsearch.randomizedtesting.RandomizedTest; +import com.carrotsearch.randomizedtesting.annotations.SeedDecorators; +import com.carrotsearch.randomizedtesting.annotations.TestMethodProviders; +import com.carrotsearch.randomizedtesting.annotations.ThreadLeakAction; +import com.carrotsearch.randomizedtesting.annotations.ThreadLeakFilters; +import com.carrotsearch.randomizedtesting.annotations.ThreadLeakGroup; +import com.carrotsearch.randomizedtesting.annotations.ThreadLeakLingering; +import com.carrotsearch.randomizedtesting.annotations.ThreadLeakScope; +import com.carrotsearch.randomizedtesting.annotations.ThreadLeakZombies; +import com.carrotsearch.randomizedtesting.annotations.TimeoutSuite; + +import java.net.http.HttpHeaders; +import java.util.ArrayList; +import java.util.HashSet; +import java.util.List; +import java.util.Map; +import java.util.Set; +import java.util.TreeMap; +import java.util.TreeSet; + +import static org.junit.Assert.assertEquals; +import static org.junit.Assert.assertNotNull; +import static org.junit.Assert.assertTrue; + +@TestMethodProviders({ JUnit3MethodProvider.class }) +@SeedDecorators({ MixWithSuiteName.class }) // See LUCENE-3995 for rationale. +@ThreadLeakScope(ThreadLeakScope.Scope.SUITE) +@ThreadLeakGroup(ThreadLeakGroup.Group.MAIN) +@ThreadLeakAction({ ThreadLeakAction.Action.WARN, ThreadLeakAction.Action.INTERRUPT }) +@ThreadLeakZombies(ThreadLeakZombies.Consequence.IGNORE_REMAINING_TESTS) +@ThreadLeakLingering(linger = 5000) // 5 sec lingering +@ThreadLeakFilters(filters = { HttpClientThreadLeakFilter.class, BouncyCastleThreadFilter.class }) +@TimeoutSuite(millis = 2 * 60 * 60 * 1000) +public abstract class RestHttpClientTestCase extends RandomizedTest { + /** + * Assert that the actual headers are the expected ones given the original default and request headers. Some headers can be ignored, + * for instance in case the http client is adding its own automatically. + * + * @param defaultHeaders the default headers set to the REST client instance + * @param requestHeaders the request headers sent with a particular request + * @param actualHeaders the actual headers as a result of the provided default and request headers + * @param ignoreHeaders header keys to be ignored as they are not part of default nor request headers, yet they + * will be part of the actual ones + */ + protected static void assertHeaders( + final Map> defaultHeaders, + final Map> requestHeaders, + final HttpHeaders actualHeaders, + final Set ignoreHeaders + ) { + final Map> expectedHeaders = new TreeMap<>(String.CASE_INSENSITIVE_ORDER); + final Set requestHeaderKeys = new HashSet<>(); + for (final Map.Entry> header : requestHeaders.entrySet()) { + final String name = header.getKey(); + addValueToListEntry(expectedHeaders, name, header.getValue()); + requestHeaderKeys.add(name); + } + for (final Map.Entry> defaultHeader : defaultHeaders.entrySet()) { + final String name = defaultHeader.getKey(); + if (requestHeaderKeys.contains(name) == false) { + addValueToListEntry(expectedHeaders, name, defaultHeader.getValue()); + } + } + Set actualIgnoredHeaders = new TreeSet<>(String.CASE_INSENSITIVE_ORDER); + for (final Map.Entry> responseHeader : actualHeaders.map().entrySet()) { + final String name = responseHeader.getKey(); + if (ignoreHeaders.contains(name)) { + expectedHeaders.remove(name); + actualIgnoredHeaders.add(name); + continue; + } + final String value = responseHeader.getValue().getFirst(); + final List values = expectedHeaders.get(name); + assertNotNull("found response header [" + name + "] that wasn't originally sent: " + value, values); + assertTrue("found incorrect response header [" + name + "]: " + value, values.remove(value)); + if (values.isEmpty()) { + expectedHeaders.remove(name); + } + } + assertEquals("some headers meant to be ignored were not part of the actual headers", ignoreHeaders, actualIgnoredHeaders); + assertTrue("some headers that were sent weren't returned " + expectedHeaders, expectedHeaders.isEmpty()); + } + + private static void addValueToListEntry(final Map> map, final String name, final List values) { + map.computeIfAbsent(name, k -> new ArrayList<>()).addAll(values); + } +} diff --git a/client/rest/src/test/java/org/opensearch/client/RestClientSingleHostIntegTests.java b/client/rest/src/test/java/org/opensearch/client/RestClientSingleHostIntegTests.java index fc73e03201e84..92b7e740da0e1 100644 --- a/client/rest/src/test/java/org/opensearch/client/RestClientSingleHostIntegTests.java +++ b/client/rest/src/test/java/org/opensearch/client/RestClientSingleHostIntegTests.java @@ -483,6 +483,18 @@ public void testEncodeParams() throws Exception { Response response = RestClientSingleHostTests.performRequestSyncOrAsync(restClient, request); assertEquals(pathPrefix + "/200?routing=foo%5Ebar", response.getRequestLine().getUri()); } + { + Request request = new Request("PUT", "/200"); + request.addParameter("pretty", null); + Response response = RestClientSingleHostTests.performRequestSyncOrAsync(restClient, request); + assertEquals(pathPrefix + "/200?pretty", response.getRequestLine().getUri()); + } + { + Request request = new Request("PUT", "/200"); + request.addParameter("pretty", ""); + Response response = RestClientSingleHostTests.performRequestSyncOrAsync(restClient, request); + assertEquals(pathPrefix + "/200?pretty=", response.getRequestLine().getUri()); + } } /** diff --git a/settings.gradle b/settings.gradle index 56051b83cc3e6..2d4b3d4db6755 100644 --- a/settings.gradle +++ b/settings.gradle @@ -49,6 +49,7 @@ List projects = [ 'rest-api-spec', 'docs', 'client:rest', + 'client:rest-http-client', 'client:rest-high-level', 'client:sniffer', 'client:test', From 1fea922af302b4c015a339b68835a6eb7e175888 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=F0=9D=90=87=F0=9D=90=9A=F0=9D=90=AB=F0=9D=90=A2?= =?UTF-8?q?=F0=9D=90=AC=F0=9D=90=A1=20=F0=9D=90=8D=F0=9D=90=9A=F0=9D=90=AB?= =?UTF-8?q?=F0=9D=90=9A=F0=9D=90=AC=F0=9D=90=A2=F0=9D=90=A6=F0=9D=90=A1?= =?UTF-8?q?=F0=9D=90=9A=F0=9D=90=A7=20=F0=9D=90=8A?= <163456392+HarishNarasimhanK@users.noreply.github.com> Date: Fri, 12 Jun 2026 05:46:54 +0530 Subject: [PATCH 10/14] [Sandbox] feat(datafusion): expose parquet metadata / statistics cache stats and search stats on the analytics backend stats endpoint (#21854) * [Sandbox] feat(datafusion): expose cache stats and search execution stats on the analytics backend stats endpoint Signed-off-by: Harish Narasimhan * [Sandbox] feat(datafusion): add object_store_read_time_ms to search execution stats Adds the object_store_read_time_ms search stat: cumulative wall-clock time spent in object-store reads (get_range/get_ranges) on the indexed scan path, accumulated per-partition to avoid over-counting. Signed-off-by: Harish Narasimhan * [Sandbox] feat(datafusion): add parquet_bytes_scanned to search execution stats Signed-off-by: Harish Narasimhan --------- Signed-off-by: Harish Narasimhan Co-authored-by: Harish Narasimhan --- .../rust/src/cache.rs | 28 +- .../rust/src/custom_cache_manager.rs | 49 ++++ .../rust/src/ffm.rs | 32 ++- .../rust/src/indexed_table/metrics.rs | 32 ++- .../rust/src/indexed_table/parquet_bridge.rs | 44 +++- .../rust/src/indexed_table/stream.rs | 26 +- .../rust/src/indexed_table/table_provider.rs | 43 +-- .../rust/src/lib.rs | 1 + .../rust/src/search_stats.rs | 249 ++++++++++++++++++ .../rust/src/statistics_cache.rs | 25 +- .../rust/src/stats.rs | 182 ++++++++++++- .../be/datafusion/DataFusionService.java | 6 +- .../be/datafusion/nativelib/NativeBridge.java | 27 +- .../be/datafusion/nativelib/StatsLayout.java | 165 +++++++++++- .../be/datafusion/stats/CacheGroupStats.java | 118 +++++++++ .../be/datafusion/stats/CacheStats.java | 96 +++++++ .../be/datafusion/stats/DataFusionStats.java | 64 ++++- .../be/datafusion/stats/SearchStats.java | 195 ++++++++++++++ .../nativelib/StatsLayoutPropertyTests.java | 105 +++++++- .../nativelib/StatsLayoutTests.java | 6 +- .../be/datafusion/stats/CacheStatsTests.java | 161 +++++++++++ .../stats/DataFusionStatsPropertyTests.java | 70 ++++- .../stats/DataFusionStatsTests.java | 45 ++++ 23 files changed, 1683 insertions(+), 86 deletions(-) create mode 100644 sandbox/plugins/analytics-backend-datafusion/rust/src/search_stats.rs create mode 100644 sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/stats/CacheGroupStats.java create mode 100644 sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/stats/CacheStats.java create mode 100644 sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/stats/SearchStats.java create mode 100644 sandbox/plugins/analytics-backend-datafusion/src/test/java/org/opensearch/be/datafusion/stats/CacheStatsTests.java diff --git a/sandbox/plugins/analytics-backend-datafusion/rust/src/cache.rs b/sandbox/plugins/analytics-backend-datafusion/rust/src/cache.rs index 5ef074d0beff7..d5fb186acbc51 100644 --- a/sandbox/plugins/analytics-backend-datafusion/rust/src/cache.rs +++ b/sandbox/plugins/analytics-backend-datafusion/rust/src/cache.rs @@ -6,6 +6,7 @@ * compatible open source license. */ +use std::sync::atomic::{AtomicUsize, Ordering}; use std::sync::{Arc, Mutex}; use datafusion::execution::cache::cache_manager::{ @@ -28,15 +29,32 @@ fn log_cache_error(operation: &str, error: &str) { // Wrapper to make Mutex implement FileMetadataCache pub struct MutexFileMetadataCache { pub inner: Mutex, + hit_count: AtomicUsize, + miss_count: AtomicUsize, } impl MutexFileMetadataCache { pub fn new(cache: DefaultFilesMetadataCache) -> Self { Self { inner: Mutex::new(cache), + hit_count: AtomicUsize::new(0), + miss_count: AtomicUsize::new(0), } } + pub fn hit_count(&self) -> usize { + self.hit_count.load(Ordering::Relaxed) + } + + pub fn miss_count(&self) -> usize { + self.miss_count.load(Ordering::Relaxed) + } + + pub fn reset_stats(&self) { + self.hit_count.store(0, Ordering::Relaxed); + self.miss_count.store(0, Ordering::Relaxed); + } + pub fn clear_cache(&self) { if let Ok(cache) = self.inner.lock() { cache.clear(); @@ -61,7 +79,15 @@ impl MutexFileMetadataCache { impl CacheAccessor for MutexFileMetadataCache { fn get(&self, k: &Path) -> Option { match self.inner.lock() { - Ok(cache) => cache.get(k), + Ok(cache) => { + let result = cache.get(k); + if result.is_some() { + self.hit_count.fetch_add(1, Ordering::Relaxed); + } else { + self.miss_count.fetch_add(1, Ordering::Relaxed); + } + result + } Err(e) => { log_cache_error("get", &e.to_string()); None diff --git a/sandbox/plugins/analytics-backend-datafusion/rust/src/custom_cache_manager.rs b/sandbox/plugins/analytics-backend-datafusion/rust/src/custom_cache_manager.rs index a40bf8f1c0df5..097d3657b9e8e 100644 --- a/sandbox/plugins/analytics-backend-datafusion/rust/src/custom_cache_manager.rs +++ b/sandbox/plugins/analytics-backend-datafusion/rust/src/custom_cache_manager.rs @@ -504,10 +504,59 @@ impl CustomCacheManager { .unwrap_or(0.0) } + /// Get statistics cache entry count + pub fn statistics_cache_entry_count(&self) -> usize { + self.statistics_cache.as_ref() + .map(|cache| >::len(cache)) + .unwrap_or(0) + } + + /// Get statistics cache size limit in bytes + pub fn statistics_cache_size_limit(&self) -> usize { + self.statistics_cache.as_ref() + .map(|cache| cache.current_size_limit()) + .unwrap_or(0) + } + /// Reset statistics cache stats pub fn statistics_cache_reset_stats(&self) { if let Some(cache) = &self.statistics_cache { cache.reset_stats(); } } + + /// Get metadata cache hit count + pub fn metadata_cache_hit_count(&self) -> usize { + self.file_metadata_cache.as_ref() + .map(|cache| cache.hit_count()) + .unwrap_or(0) + } + + /// Get metadata cache miss count + pub fn metadata_cache_miss_count(&self) -> usize { + self.file_metadata_cache.as_ref() + .map(|cache| cache.miss_count()) + .unwrap_or(0) + } + + /// Get metadata cache entry count + pub fn metadata_cache_entry_count(&self) -> usize { + self.file_metadata_cache.as_ref() + .map(|cache| >::len(cache)) + .unwrap_or(0) + } + + /// Get metadata cache size limit in bytes + pub fn metadata_cache_size_limit(&self) -> usize { + self.file_metadata_cache.as_ref() + .map(|cache| cache.get_cache_limit()) + .unwrap_or(0) + } + + /// Reset metadata cache stats + pub fn metadata_cache_reset_stats(&self) { + if let Some(cache) = &self.file_metadata_cache { + cache.reset_stats(); + } + } } diff --git a/sandbox/plugins/analytics-backend-datafusion/rust/src/ffm.rs b/sandbox/plugins/analytics-backend-datafusion/rust/src/ffm.rs index edb43e957f60c..2307b0b4356b8 100644 --- a/sandbox/plugins/analytics-backend-datafusion/rust/src/ffm.rs +++ b/sandbox/plugins/analytics-backend-datafusion/rust/src/ffm.rs @@ -854,6 +854,7 @@ pub unsafe extern "C" fn df_create_session_context( plan_ptr: *const u8, plan_len: i64, ) -> i64 { + crate::search_stats::inc_listing_table_scan(); let table_name = str_from_raw(table_name_ptr, table_name_len) .map_err(|e| format!("df_create_session_context: {}", e))?; let query_config = @@ -893,6 +894,11 @@ pub unsafe extern "C" fn df_create_session_context_indexed( plan_ptr: *const u8, plan_len: i64, ) -> i64 { + match tree_shape { + 1 => crate::search_stats::inc_single_collector_scan(), + 2 => crate::search_stats::inc_bitmap_tree_scan(), + _ => {} + } let table_name = str_from_raw(table_name_ptr, table_name_len) .map_err(|e| format!("df_create_session_context_indexed: {}", e))?; let query_config = @@ -1156,12 +1162,18 @@ pub unsafe extern "C" fn df_execute_with_context( /// Collects all native executor metrics into a caller-provided byte buffer. /// -/// The buffer must have capacity for at least `size_of::()` bytes (344). +/// `runtime_ptr` may be `0` to skip cache-stats collection. When non-zero it +/// must be a valid pointer returned by [`df_create_global_runtime`]. +/// +/// The buffer must have capacity for at least `size_of::()` bytes (552). /// Returns 0 on success. #[ffm_safe] #[no_mangle] -pub unsafe extern "C" fn df_stats(out_ptr: *mut u8, out_cap: i64) -> i64 { - use crate::stats::{layout, pack_runtime_metrics, pack_task_monitor, pack_partition_gate, DfStatsBuffer, RuntimeMetricsRepr}; +pub unsafe extern "C" fn df_stats(runtime_ptr: i64, out_ptr: *mut u8, out_cap: i64) -> i64 { + use crate::stats::{ + layout, pack_cache_stats, pack_partition_gate, pack_runtime_metrics, pack_task_monitor, + CacheStatsRepr, DfStatsBuffer, RuntimeMetricsRepr, + }; use crate::task_monitors::{ coordinator_reduce_monitor, query_execution_monitor, stream_next_monitor, plan_setup_monitor, @@ -1190,6 +1202,18 @@ pub unsafe extern "C" fn df_stats(out_ptr: *mut u8, out_cap: i64) -> i64 { RuntimeMetricsRepr::zeroed() }; + // Cache stats (zeroed when no runtime pointer or no cache manager) + let cache_stats = if runtime_ptr != 0 { + let runtime = &*(runtime_ptr as *const DataFusionRuntime); + runtime + .custom_cache_manager + .as_ref() + .map(pack_cache_stats) + .unwrap_or_else(CacheStatsRepr::default) + } else { + CacheStatsRepr::default() + }; + let buf = DfStatsBuffer { io_runtime, cpu_runtime, @@ -1199,6 +1223,8 @@ pub unsafe extern "C" fn df_stats(out_ptr: *mut u8, out_cap: i64) -> i64 { plan_setup: pack_task_monitor(plan_setup_monitor()), fragment_executor_gate: pack_partition_gate(mgr.cpu_executor.concurrency_gate()), reduce_executor_gate: pack_partition_gate(mgr.coordinator_gate()), + cache_stats, + search_stats: crate::search_stats::snapshot(), }; // Copy struct bytes to caller buffer diff --git a/sandbox/plugins/analytics-backend-datafusion/rust/src/indexed_table/metrics.rs b/sandbox/plugins/analytics-backend-datafusion/rust/src/indexed_table/metrics.rs index fc2a6a1254cef..72834911d7057 100644 --- a/sandbox/plugins/analytics-backend-datafusion/rust/src/indexed_table/metrics.rs +++ b/sandbox/plugins/analytics-backend-datafusion/rust/src/indexed_table/metrics.rs @@ -18,6 +18,8 @@ use datafusion::physical_plan::metrics::{ Count, ExecutionPlanMetricsSet, MetricBuilder, MetricsSet, Time, }; +use crate::indexed_table::parquet_bridge::ReadIoStats; + /// Lightweight metric handles passed from `IndexedExec` to the streaming loop. /// /// All fields are `Option` because standalone uses of `IndexedExec` (i.e. not @@ -100,6 +102,14 @@ pub struct StreamMetrics { /// Time spent polling the inner parquet stream (pull decoded /// batch), isolating decode from our own processing. pub parquet_poll_time: Option