Share sparkey readers across DoFns with memory-aware heap fallback

[spkrka]

Summary

Sparkey side inputs on Dataflow suffer from two problems:

1. Per-DoFn reader duplication: Beam creates one DoFn per vCPU (e.g. 80 on n4-standard-80). Each independently opens the same sparkey files via mmap, wasting file descriptors and mmap regions.

2. Page cache thrashing: Dataflow hardcodes the JVM heap to ~70% of worker memory. On a 320 GB machine, only ~96 GB remains for the OS page cache — but sparkey side inputs can easily exceed 200 GB. The result is severe thrashing and I/O-bound jobs (10-14h runtimes).

This PR addresses both:

• Shared reader cache — A static ConcurrentHashMap<path, CompletableFuture<SparkeyReader>> deduplicates readers across all DoFn instances on a worker. The first thread to request a reader loads it; others block on future.join(). Failed futures stay cached so subsequent callers fast-fail (reader failures are deterministic and fatal).

• Memory-aware heap fallback — New HostMemoryTracker estimates two budgets at JVM startup:

  • Off-heap: totalPhysicalMemory - maxHeap - 2 GB reserve (for OS, kernel, Beam shuffle, etc.)
  • Heap: maxHeap - max(4 GB, 10% of maxHeap) (GC headroom and non-sparkey allocations)

  Each shard atomically claims budget via getAndAccumulate (lock-free CAS). Strategy per shard:

  1. Try off-heap → open via mmap (default, zero-copy)
  2. If off-heap exhausted → try heap → open via Sparkey.reader().useHeap(true) (reads into byte[])
  3. If neither fits → fall back to mmap with a warning (best-effort, may thrash)

  This lets the otherwise-wasted JVM heap absorb overflow, achieving ~100% memory utilization.

Design choices

• Dedicated 4-thread loader pool instead of ForkJoinPool.commonPool(): reader loading involves blocking I/O (GCS download, disk reads, heap copying). Using the common pool risks starving other JVM work. 4 threads balances parallelism for multiple side inputs without overwhelming disk/network. Threads are daemon (won't block JVM shutdown) and named sparkey-reader-loader-N for debuggability.

• No budget release: Readers are cached for the JVM lifetime and never closed (Beam side inputs have no teardown lifecycle). Budget is similarly never released. If a sharded reader fails partway through (some shards claimed), that budget is leaked — acceptable since the failure is fatal to the pipeline.

• No cache eviction on failure: If getReader throws, the failed CompletableFuture stays in the cache. Reader failures are deterministic (corrupt files, missing shards), so retrying would just repeat the same failure. Fast-failing is preferable.

• Heap reserve minimum of 4 GB: Workers with <4 GB heap get zero heap budget, effectively disabling the heap fallback. This is intentional — small heaps shouldn't absorb sparkey data.

• Sparkey 3.5.1 → 3.7.0: Required for Sparkey.reader().useHeap(true) builder API.

Impact

Tested on a pipeline processing 1.6B user entities with ~200 GB of sparkey side inputs on n4-standard-80 (320 GB):

• Runtime: 10-13h → 33 min (20x)
• Throughput: 8-15K → 1.4M elem/s (100x)
• Cost per run: ~$1,300 → ~$70
• CPU utilization: low (I/O bound) → 100% (compute bound)

Test plan

• Existing sparkey unit tests pass
• MiMa binary compatibility check passes
• Manual validation on Dataflow with sharded sparkey side inputs
• Verify log output shows budget claims and heap/mmap decisions

🤖 Generated with Claude Code

[codecov]

Codecov Report

✅ All modified and coverable lines are covered by tests.
✅ Project coverage is 61.73%. Comparing base (d59fbba) to head (47495bd).

Additional details and impacted files

@@            Coverage Diff             @@
##             main    #5903      +/-   ##
==========================================
+ Coverage   61.54%   61.73%   +0.19%     
==========================================
  Files         317      318       +1     
  Lines       11653    11709      +56     
  Branches      822      833      +11     
==========================================
+ Hits         7172     7229      +57     
+ Misses       4481     4480       -1     

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

[spkrka]

I guess this is technically not purely a sparkey thing, could also be used for other resources/sideinput types