Skip to content

[Java] Add headers provider - #524

Open
danilonajkov-db wants to merge 6 commits into
mainfrom
effort/java-headers-provider
Open

[Java] Add headers provider#524
danilonajkov-db wants to merge 6 commits into
mainfrom
effort/java-headers-provider

Conversation

@danilonajkov-db

@danilonajkov-db danilonajkov-db commented Jul 15, 2026

Copy link
Copy Markdown
Member

What changes are proposed in this pull request?

Adds custom HeadersProvider authentication to the Java SDK, matching the Rust stream-builder API.

  • Adds the HeadersProvider interface with getHeaders() and optional invalidate().
  • Adds StreamBuilder.headersProvider() as an alternative to .oauth().
  • Supports JSON, Protocol Buffer, and Arrow streams.
  • Retains the provider for stream recovery and explicit stream recreation.
  • Keeps the deprecated legacy create*Stream APIs unchanged.
  • Ensures JNI callbacks run on blocking threads and reclaim local references after every invocation.
  • Validates custom gRPC metadata keys fallibly instead of allowing invalid names to panic.
  • Documents that header names should come from a small, fixed set.

Authentication Behavior

.oauth() and .headersProvider() are mutually exclusive. Calling either replaces the previously configured authentication method.

When the server rejects credentials, the SDK calls HeadersProvider.invalidate() before retrying so providers can clear cached tokens.

Testing

  • cargo fmt --all --check
  • cargo check -p zerobus-jni
  • cargo test -p databricks-zerobus-ingest-sdk invalid_custom_header_name_returns_error
  • Java 8-compatible compilation of the SDK and provider harness
  • Staging E2E:
    • Provider fetched an OAuth token inside getHeaders()
    • Created a JSON stream through StreamBuilder.headersProvider()
    • Ingested a record
    • Flushed successfully

Signed-off-by: danilo-najkov-db <danilo.najkov@databricks.com>
@danilonajkov-db
danilonajkov-db force-pushed the effort/java-headers-provider branch from a2da75a to 9a9db35 Compare July 16, 2026 08:44
…rovider

Signed-off-by: danilo-najkov-db <danilo.najkov@databricks.com>

# Conflicts:
#	java/NEXT_CHANGELOG.md
#	rust/jni/Cargo.toml
@danilonajkov-db
danilonajkov-db marked this pull request as ready for review July 16, 2026 13:11
@danilonajkov-db
danilonajkov-db requested a review from a team July 16, 2026 15:47

@teodordelibasic-db teodordelibasic-db left a comment

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Mostly edge cases and races flagged by LLM.

Comment thread java/examples/json/README.md
Comment thread rust/jni/src/headers_provider.rs
Comment thread java/src/test/java/com/databricks/zerobus/StreamBuilderTest.java
Comment thread rust/sdk/src/stream/grpc/connection.rs Outdated
Comment thread rust/jni/src/headers_provider.rs
Comment thread rust/jni/src/headers_provider.rs Outdated
Comment thread rust/jni/src/headers_provider.rs Outdated
Signed-off-by: danilo-najkov-db <danilo.najkov@databricks.com>
Signed-off-by: danilo-najkov-db <danilo.najkov@databricks.com>
Signed-off-by: danilo-najkov-db <danilo.najkov@databricks.com>
Signed-off-by: danilo-najkov-db <danilo.najkov@databricks.com>

#[async_trait]
impl HeadersProvider for JavaHeadersProvider {
async fn get_headers(&self) -> ZerobusResult<HashMap<&'static str, String>> {

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

A callback that exceeds recoveryTimeoutMs keeps running after the stream attempt times out. Tokio cannot cancel a spawn_blocking task once it has started, so dropping the future in rust/sdk/src/stream/grpc/connection.rs:80 leaves the Java call holding its worker and cloned GlobalRef. The next retry can start another callback, and repeated stream attempts can accumulate blocked workers and retained provider objects.

Could we serialize calls per provider and hold the guard inside the blocking closure, so cancellation cannot release it while the Java call is still running?

pub struct JavaHeadersProvider {
    provider_ref: GlobalRef,
    callback_lock: Arc<tokio::sync::Mutex<()>>,
}

let guard = Arc::clone(&self.callback_lock).lock_owned().await;
let provider_ref = self.provider_ref.clone();
tokio::task::spawn_blocking(move || {
    let _guard = guard;
    get_headers_blocking(provider_ref)
})
.await

This would not make a synchronous Java call cancellable, but it would keep one stuck provider adapter from consuming a new worker on every retry. A process-wide bounded executor would give a stronger global bound.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants