diff --git a/.github/workflows/ci-dotnet.yml b/.github/workflows/ci-dotnet.yml new file mode 100644 index 00000000..0138ea94 --- /dev/null +++ b/.github/workflows/ci-dotnet.yml @@ -0,0 +1,132 @@ +name: .NET CI + +on: + workflow_call: + inputs: + use_local_sdk: + description: 'Override databricks-zerobus-ingest-sdk with the local path dep (for integration testing unreleased Rust changes)' + type: boolean + default: false + +jobs: + fmt: + runs-on: + group: databricks-protected-runner-group + labels: linux-ubuntu-latest + defaults: + run: + working-directory: dotnet + env: + SkipNativeBuild: true + steps: + - uses: actions/checkout@v6 + + - uses: actions/setup-dotnet@v5 + with: + dotnet-version: 10.0.x + + - name: Install dotnet-format + run: | + dotnet tool install -g dotnet-format + echo "$HOME/.dotnet/tools" >> $GITHUB_PATH + + - uses: dtolnay/rust-toolchain@stable + with: + components: rustfmt + + - name: Format all files + run: make fmt + + - name: Fail on differences + run: git diff --exit-code + + test: + strategy: + fail-fast: false + matrix: + include: + - os: linux-ubuntu-latest + runner: { group: databricks-protected-runner-group, labels: linux-ubuntu-latest } + - os: windows-server-latest + runner: { group: databricks-protected-runner-group, labels: windows-server-latest } + - os: macos-arm64 + runner: [self-hosted, macOS, ARM64, zerobus-sdk] + runs-on: ${{ matrix.runner }} + defaults: + run: + working-directory: dotnet + shell: bash + steps: + - uses: actions/checkout@v6 + + - name: Unshallow + run: git fetch --prune --unshallow + + - uses: actions/setup-dotnet@v5 + with: + dotnet-version: 10.0.x + + - uses: dtolnay/rust-toolchain@stable + with: + targets: x86_64-pc-windows-gnu + + - name: Patch Rust workspace to use local SDK + if: inputs.use_local_sdk + run: printf '\n[patch.crates-io]\ndatabricks-zerobus-ingest-sdk = { path = "sdk" }\n' >> ../rust/Cargo.toml + + - name: Build Rust FFI + run: make build-rust + + - name: Build .NET SDK + run: make build-dotnet + + - name: Run tests + run: make test + + - name: Build examples + run: make examples + + user-experience: + # Skip when testing local SDK — this job tests the published release download flow + if: ${{ !inputs.use_local_sdk }} + strategy: + fail-fast: false + matrix: + include: + - os: linux-ubuntu-latest + runner: { group: databricks-protected-runner-group, labels: linux-ubuntu-latest } + lib: lib/linux_amd64/libzerobus_ffi.a + - os: windows-server-latest + runner: { group: databricks-protected-runner-group, labels: windows-server-latest } + lib: lib/windows_amd64/libzerobus_ffi.a + - os: macos-arm64 + runner: [self-hosted, macOS, ARM64, zerobus-sdk] + lib: lib/darwin_arm64/libzerobus_ffi.a + runs-on: ${{ matrix.runner }} + defaults: + run: + working-directory: dotnet + shell: bash + steps: + - uses: actions/checkout@v6 + + - uses: actions/setup-dotnet@v5 + with: + dotnet-version: 10.0.x + + - uses: dtolnay/rust-toolchain@stable + + - name: Test .NET build works (user step 2) + run: dotnet build + + - name: Test JSON single example builds + working-directory: dotnet/examples/JsonSingle + run: dotnet build + + - name: Test JSON batch example builds + working-directory: dotnet/examples/JsonBatch + run: dotnet build + + - name: Test Proto single example builds + working-directory: dotnet/examples/ProtoSingle + run: dotnet build diff --git a/.github/workflows/push.yml b/.github/workflows/push.yml index 60a0ba35..8b497b6a 100644 --- a/.github/workflows/push.yml +++ b/.github/workflows/push.yml @@ -11,7 +11,7 @@ jobs: runs-on: group: databricks-protected-runner-group labels: linux-ubuntu-latest - needs: [changes, rust, java, typescript, python, go, purego, cpp] + needs: [changes, rust, java, typescript, python, go, purego, cpp, dotnet] if: always() steps: - run: | @@ -22,7 +22,8 @@ jobs: "${{ needs.python.result }}" \ "${{ needs.go.result }}" \ "${{ needs.purego.result }}" \ - "${{ needs.cpp.result }}"; do + "${{ needs.cpp.result }}" \ + "${{ needs.dotnet.result }}"; do if [[ "$result" != "success" && "$result" != "skipped" ]]; then echo "Job failed or was cancelled: $result" exit 1 @@ -41,6 +42,7 @@ jobs: go: ${{ steps.filter.outputs.go }} purego: ${{ steps.filter.outputs.purego }} cpp: ${{ steps.filter.outputs.cpp }} + dotnet: ${{ steps.filter.outputs.dotnet }} steps: - uses: actions/checkout@34e114876b0b11c390a56381ad16ebd13914f8d5 # v4 - uses: dorny/paths-filter@de90cc6fb38fc0963ad72b210f1f284cd68cea36 # v3 @@ -71,6 +73,9 @@ jobs: cpp: - 'cpp/**' - '.github/workflows/ci-cpp.yml' + dotnet: + - 'dotnet/**' + - '.github/workflows/ci-dotnet.yml' # Blocking jobs — required for merge rust: @@ -129,6 +134,11 @@ jobs: contents: read uses: ./.github/workflows/ci-cpp.yml + dotnet: + needs: changes + if: needs.changes.outputs.dotnet == 'true' + uses: ./.github/workflows/ci-dotnet.yml + # Non-blocking cross-SDK jobs — run all SDK tests against the local Rust # source when rust/** changes, to catch breakage before a new FFI release. # These are intentionally excluded from `gate` so they don't block merging. @@ -150,6 +160,13 @@ jobs: contents: read uses: ./.github/workflows/ci-cpp.yml + cross-sdk-dotnet: + needs: changes + if: needs.changes.outputs.rust == 'true' + uses: ./.github/workflows/ci-dotnet.yml + with: + use_local_sdk: true + cross-sdk-python: needs: changes if: needs.changes.outputs.rust == 'true' diff --git a/.github/workflows/release-dotnet.yml b/.github/workflows/release-dotnet.yml new file mode 100644 index 00000000..1efb1510 --- /dev/null +++ b/.github/workflows/release-dotnet.yml @@ -0,0 +1,194 @@ +name: Build & Release .NET SDK +run-name: "Release .NET SDK ${{ github.ref_name }}" + +on: + push: + tags: + - 'dotnet/v*' + workflow_dispatch: + +concurrency: + group: ${{ github.workflow }}-${{ github.ref }} + cancel-in-progress: true + +permissions: + contents: write + id-token: write # Required for attestation + attestations: write # Required for attestation + +jobs: + build-ffi: + strategy: + fail-fast: false + matrix: + include: + - runner: { group: databricks-protected-runner-group, labels: linux-ubuntu-latest } + target: x86_64-unknown-linux-gnu + static_lib: libzerobus_ffi.a + dynamic_lib: libzerobus_ffi.so + artifact_dir: linux-x64 + cross: false + - runner: { group: databricks-protected-runner-group, labels: linux-ubuntu-latest } + target: aarch64-unknown-linux-gnu + static_lib: libzerobus_ffi.a + dynamic_lib: libzerobus_ffi.so + artifact_dir: linux-arm64 + cross: true + - runner: { group: databricks-protected-runner-group, labels: windows-server-latest } + target: x86_64-pc-windows-msvc + static_lib: zerobus_ffi.lib + dynamic_lib: zerobus_ffi.dll + artifact_dir: win-x64 + cross: false + - runner: [self-hosted, macOS, ARM64, zerobus-sdk] + target: aarch64-apple-darwin + static_lib: libzerobus_ffi.a + dynamic_lib: libzerobus_ffi.dylib + artifact_dir: osx-arm64 + cross: false + - runner: [self-hosted, macOS, ARM64, zerobus-sdk] + target: x86_64-apple-darwin + static_lib: libzerobus_ffi.a + dynamic_lib: libzerobus_ffi.dylib + artifact_dir: osx-x64 + cross: false + + name: Build FFI - ${{ matrix.artifact_dir }} + runs-on: ${{ matrix.runner }} + + defaults: + run: + working-directory: rust + + steps: + - uses: actions/checkout@v4 + + - name: Install protoc + uses: arduino/setup-protoc@v3 + + - name: Install cross-compilation toolchain (Linux ARM64) + if: matrix.cross + run: | + sudo apt-get update -qq + sudo apt-get install -y -qq gcc-aarch64-linux-gnu + mkdir -p .cargo + echo '[target.aarch64-unknown-linux-gnu]' > .cargo/config.toml + echo 'linker = "aarch64-linux-gnu-gcc"' >> .cargo/config.toml + + - name: Install Rust toolchain + uses: dtolnay/rust-toolchain@stable + with: + targets: ${{ matrix.target }} + + - name: Build FFI library + run: cargo build --release -p zerobus-ffi --target ${{ matrix.target }} + + - name: Prepare artifacts + shell: bash + run: | + mkdir -p artifacts/${{ matrix.artifact_dir }}/native + cp target/${{ matrix.target }}/release/${{ matrix.static_lib }} artifacts/${{ matrix.artifact_dir }}/native/ + cp target/${{ matrix.target }}/release/${{ matrix.dynamic_lib }} artifacts/${{ matrix.artifact_dir }}/native/ + + - name: Upload artifact + uses: actions/upload-artifact@v4 + with: + name: ffi-${{ matrix.artifact_dir }} + path: rust/artifacts/${{ matrix.artifact_dir }} + + publish: + needs: build-ffi + runs-on: + group: databricks-protected-runner-group + labels: linux-ubuntu-latest + defaults: + run: + working-directory: dotnet + + steps: + - uses: actions/checkout@v4 + + - name: Set up .NET + uses: actions/setup-dotnet@v5 + with: + dotnet-version: 10.0.x + + - name: Download all native libraries + uses: actions/download-artifact@v8 + with: + path: native + pattern: ffi-* + merge-multiple: false + + - name: Consolidate native libraries + run: | + mkdir -p src/Zerobus/runtimes + cp -r native/ffi-linux-x64/linux-x64 src/Zerobus/runtimes/ + cp -r native/ffi-linux-arm64/linux-arm64 src/Zerobus/runtimes/ + cp -r native/ffi-osx-x64/osx-x64 src/Zerobus/runtimes/ + cp -r native/ffi-osx-arm64/osx-arm64 src/Zerobus/runtimes/ + cp -r native/ffi-win-x64/win-x64 src/Zerobus/runtimes/ + + - name: Extract version from tag + id: version + run: echo "version=${GITHUB_REF#refs/tags/dotnet/v}" >> "$GITHUB_OUTPUT" + + - name: Package NuGet + env: + CI: true + run: | + dotnet pack src/Zerobus/Zerobus.csproj \ + --configuration Release \ + --output ./packages \ + -p:SkipNativeBuild=true \ + -p:Version=${{ steps.version.outputs.version }} \ + -p:PackageVersion=${{ steps.version.outputs.version }} \ + -p:AssemblyVersion=${{ steps.version.outputs.version }} \ + -p:FileVersion=${{ steps.version.outputs.version }} \ + -p:InformationalVersion=${{ steps.version.outputs.version }} \ + -p:CI=true + + - name: Generate attestation + uses: actions/attest-build-provenance@v4 + with: + subject-path: "packages/*.nupkg" + + - name: NuGet login (OIDC → temp API key) + uses: NuGet/login@v1 + id: login + with: + user: ${{ secrets.NUGET_USERNAME }} + + - name: Publish to Nuget Packages + env: + NUGET_AUTH_TOKEN: ${{ steps.login.outputs.NUGET_API_KEY }} + run: | + dotnet nuget push packages/*.nupkg \ + --api-key "${{ steps.login.outputs.NUGET_API_KEY }}" \ + --skip-duplicate \ + --source https://api.nuget.org/v3/index.json + + release: + name: Create GitHub Release + runs-on: + group: databricks-protected-runner-group + labels: linux-ubuntu-latest + needs: publish + + steps: + - uses: actions/checkout@v4 + + - name: Extract version from tag + id: version + run: echo "version=${GITHUB_REF#refs/tags/dotnet/v}" >> "$GITHUB_OUTPUT" + + - name: Extract release notes from changelog + run: | + version="v${{ steps.version.outputs.version }}" + awk "/^## Release ${version}/{found=1; next} /^## /{if(found) exit} found{print}" dotnet/CHANGELOG.md > /tmp/release-notes.md + + - name: Create GitHub Release + uses: softprops/action-gh-release@v2 + with: + name: .NET SDK v${{ steps.version.outputs.version }} + body_path: /tmp/release-notes.md diff --git a/README.md b/README.md index 09d69b39..394f00cd 100644 --- a/README.md +++ b/README.md @@ -22,6 +22,7 @@ Zerobus is a high-throughput streaming service for direct data ingestion into Da | TypeScript | [`typescript/`](typescript/) | [`@databricks/zerobus-ingest-sdk`](https://www.npmjs.com/package/@databricks/zerobus-ingest-sdk) | | Java | [`java/`](java/) | [`com.databricks:zerobus-ingest-sdk`](https://central.sonatype.com/artifact/com.databricks/zerobus-ingest-sdk) | | C++ | [`cpp/`](cpp/) | Source / CMake (`zerobus::zerobus`) | +| C# | [`dotnet/`](dotnet/) | [`Databricks.Zerobus`](https://www.nuget.org/packages/Databricks.Zerobus) | ## Platform Support diff --git a/dotnet/.gitignore b/dotnet/.gitignore new file mode 100644 index 00000000..7026de3e --- /dev/null +++ b/dotnet/.gitignore @@ -0,0 +1,50 @@ +# .NET +bin/ +obj/ +[Bb]in/ +[Oo]bj/ + +# Visual Studio +.vs/ +*.user +*.suo +*.userosscache +*.sln.docstates + +# Rider +.idea/ + +# Build results +*.dll +*.pdb +*.exe +*.app +*.cache +*.log + +# NuGet +*.nupkg +*.snupkg +packages/ +**/packages/ + +# .NET Core +project.lock.json +project.fragment.lock.json + +# Coverage +coverage*/ +*.coverage +*.coveragexml + +# Rust +/target/ +**/target/ + +# Rust tools +**/*.rs.bk + +# OS +.DS_Store + +runtimes/ diff --git a/dotnet/CHANGELOG.md b/dotnet/CHANGELOG.md new file mode 100644 index 00000000..a577603a --- /dev/null +++ b/dotnet/CHANGELOG.md @@ -0,0 +1,5 @@ +# Version changelog + +## Release v0.1.0 + +Initial release of the Databricks Zerobus Ingest SDK for .NET. diff --git a/dotnet/CONTRIBUTING.md b/dotnet/CONTRIBUTING.md new file mode 100644 index 00000000..2979c9d4 --- /dev/null +++ b/dotnet/CONTRIBUTING.md @@ -0,0 +1,155 @@ +# Contributing to the Zerobus SDK for .NET + +Please read the [top-level CONTRIBUTING.md](https://github.com/databricks/zerobus-sdk/blob/main/CONTRIBUTING.md) first for general contribution guidelines, pull request process, commit requirements, and DCO/sign-off requirements. + +This document covers .NET-specific development setup and workflow. + +## Development Setup + +### Prerequisites + +- Git +- .NET SDK 8.0 or higher +- Rust toolchain (`cargo`) - [Install Rust](https://rustup.rs/) +- Bash shell (used by `build_native.sh`) + +### Setting Up Your Development Environment + +1. **Clone the repository:** + ```bash + git clone https://github.com/databricks/zerobus-sdk.git + cd zerobus-sdk/dotnet + ``` + +2. **Build the project:** + ```bash + make build + ``` + + This will: + - Build the Rust FFI layer (`zerobus_ffi`) + - Build the .NET SDK + - Place the native library in `src/Zerobus/runtimes//native/` + +3. **Run tests:** + ```bash + make test + ``` + +## Coding Style + +Code style is enforced by formatters in pull requests. We use `dotnet format` and `rustfmt`. + +### Running the Formatter + +Format your code before committing: + +```bash +make fmt +``` + +This runs: +- `dotnet format` for .NET code +- `cargo fmt --all` for Rust FFI code + +### Running Linters + +Check your code for issues: + +```bash +make lint +``` + +This runs: +- Rust linting via `cargo clippy --all -- -D warnings` +- .NET lint target currently prints a message and does not run standalone analyzers + +## Testing + +Run the full test suite: + +```bash +make test +``` + +This runs: +- Rust FFI tests via `cargo test` +- .NET tests via `dotnet test -p:SkipNativeBuild=true` + +### Running Tests Individually + +- Unit tests: + ```bash + dotnet test tests/Zerobus.Tests + ``` + +- Integration tests: + ```bash + dotnet test tests/Zerobus.IntegrationTests + ``` + +Integration tests use a mock gRPC server and exercise the SDK through the native FFI layer. + +## Working with the Native FFI Layer + +The .NET SDK relies on the Rust FFI library in `../rust/ffi`. + +When making FFI-related changes: + +1. Update Rust code in `../rust/ffi/src/` +2. Update exported C API in `../rust/ffi/zerobus.h` if needed +3. Update .NET interop bindings in `src/Zerobus/Interop/` +4. Rebuild native artifacts: + ```bash + ./build_native.sh + ``` +5. Run tests: + ```bash + make test + ``` + +### Native Build Notes + +- `dotnet build` automatically invokes `build_native.sh` from MSBuild. +- To skip automatic native build (for example, when a prebuilt library is already present): + ```bash + dotnet build -p:SkipNativeBuild=true + ``` +- To force a native rebuild: + ```bash + ./build_native.sh --force + ``` + +## Continuous Integration + +All pull requests must pass CI checks. + +Typical checks include: + +- **formatting**: `dotnet format` and `cargo fmt` +- **lint**: Rust lint (`cargo clippy`) +- **test**: .NET and Rust test suites + +You can view CI results in the GitHub Actions tab of your pull request. + +## Makefile Targets + +Available make targets: + +- `make build` - Build both Rust FFI and .NET SDK +- `make build-rust` - Build only the Rust FFI layer +- `make build-rust-force` - Force rebuild Rust FFI artifacts +- `make build-dotnet` - Build only the .NET SDK +- `make clean` - Remove build artifacts +- `make fmt` - Format all code (Rust and .NET) +- `make fmt-dotnet` - Format .NET code +- `make fmt-rust` - Format Rust code +- `make lint` - Run linters on all code +- `make lint-dotnet` - Run .NET lint target (informational currently) +- `make lint-rust` - Run Rust clippy +- `make check` - Run formatting and lint checks +- `make test` - Run all tests +- `make test-rust` - Run Rust tests +- `make test-dotnet` - Run .NET tests +- `make examples` - Build all .NET examples +- `make help` - Show available targets diff --git a/dotnet/Directory.Build.props b/dotnet/Directory.Build.props new file mode 100644 index 00000000..3f9e490a --- /dev/null +++ b/dotnet/Directory.Build.props @@ -0,0 +1,8 @@ + + + latest + enable + enable + true + + \ No newline at end of file diff --git a/dotnet/Directory.Packages.props b/dotnet/Directory.Packages.props new file mode 100644 index 00000000..21e280d9 --- /dev/null +++ b/dotnet/Directory.Packages.props @@ -0,0 +1,24 @@ + + + + + + true + + + + + + + + + + + + + + + + + + \ No newline at end of file diff --git a/dotnet/Makefile b/dotnet/Makefile new file mode 100644 index 00000000..b6795e5e --- /dev/null +++ b/dotnet/Makefile @@ -0,0 +1,81 @@ +# Makefile for zerobus-sdk-dotnet + +.PHONY: help build build-rust build-dotnet clean fmt fmt-dotnet fmt-rust lint lint-dotnet lint-rust check test test-dotnet test-rust examples release + +help: + @echo "Available targets:" + @echo " make build - Build both Rust FFI and .NET SDK" + @echo " make build-rust - Build only the Rust FFI layer" + @echo " make build-dotnet - Build only the .NET SDK" + @echo " make clean - Remove build artifacts" + @echo " make fmt - Format all code (Rust and .NET)" + @echo " make fmt-dotnet - Format .NET code" + @echo " make fmt-rust - Format Rust code" + @echo " make lint - Run linters on all code" + @echo " make lint-dotnet - Run .NET linters" + @echo " make lint-rust - Run Rust linters" + @echo " make check - Run all checks (fmt and lint)" + @echo " make test - Run all tests (Rust and .NET)" + @echo " make test-rust - Run Rust unit tests" + @echo " make test-dotnet - Run .NET unit tests" + @echo " make examples - Build all examples" + +build: build-rust build-dotnet + +build-rust: + ./build_native.sh +build-rust-force: + ./build_native.sh --force + +build-dotnet: build-rust + @echo "Building .NET SDK..." + dotnet build + @echo "✓ .NET SDK built successfully" + +clean: + @echo "Cleaning build artifacts..." + cd ../rust/ffi && cargo clean + rm -rf src/Zerobus/runtimes/* + rm -f libzerobus_ffi.a + rm -rf releases + @echo "✓ Clean complete" + +fmt: fmt-dotnet fmt-rust + +fmt-dotnet: + @echo "Formatting .NET code..." + dotnet format + +fmt-rust: + @echo "Formatting Rust code..." + cd ../rust/ffi && cargo fmt --all + +lint: lint-dotnet lint-rust + +lint-dotnet: + @echo "Linting .NET code..." + @echo ".NET linting not available as standalone yet - skipping lint-dotnet" + +lint-rust: + @echo "Linting Rust code..." + cd ../rust/ffi && cargo clippy --all -- -D warnings + +check: fmt lint + +test: test-rust test-dotnet + +test-rust: + @echo "Running Rust tests..." + cd ../rust/ffi && cargo test -- --test-threads=1 + +test-dotnet: + @echo "Running .NET unit tests..." + dotnet test -p:SkipNativeBuild=true + @echo "✓ All .NET tests passed" + +examples: build + @echo "Building examples..." + dotnet build examples/JsonSingle + dotnet build examples/JsonBatch + dotnet build examples/ProtoSingle + @echo "✓ Examples built successfully" diff --git a/dotnet/NEXT_CHANGELOG.md b/dotnet/NEXT_CHANGELOG.md new file mode 100644 index 00000000..3406a332 --- /dev/null +++ b/dotnet/NEXT_CHANGELOG.md @@ -0,0 +1,15 @@ +# NEXT CHANGELOG + +## Release v0.2.0 + +### New Features and Improvements + +### Deprecations + +### Bug Fixes + +### Documentation + +### Internal Changes + +### API Changes diff --git a/dotnet/NOTICE b/dotnet/NOTICE new file mode 100644 index 00000000..35abc4a8 --- /dev/null +++ b/dotnet/NOTICE @@ -0,0 +1,48 @@ +Copyright (2025) Databricks, Inc. + +This Software includes software developed at Databricks (https://www.databricks.com/) and its use is subject to the included LICENSE file. + +--- + +This Software includes the following open source components: + +## Rust Dependencies (via FFI) + +This .NET SDK wraps the Zerobus Rust SDK, which includes the following open source components subject to the MIT License and Apache License 2.0: + +### MIT License + +- **tokio**: Copyright (c) 2021 Tokio Contributors (https://github.com/tokio-rs/tokio) +- **tonic**: Copyright (c) 2019-2021, Luc Perkins and contributors (https://github.com/hyperium/tonic) +- **tokio-stream**: Copyright (c) 2019 Tokio Contributors (https://github.com/tokio-rs/tokio) +- **tracing**: Copyright (c) 2019-2021, The `tracing` project authors (https://github.com/tokio-rs/tracing) + +### Apache License 2.0 + +- **prost**: Copyright 2017, Dan Burkert (https://github.com/tokio-rs/prost) +- **prost-types**: Copyright 2017, Dan Burkert (https://github.com/tokio-rs/prost) +- **tokio-retry**: Copyright 2018, App-Tech-Crew (https://github.com/app-tech-crew/tokio-retry) +- **reqwest**: Copyright (c) 2016-2021, Sean McArthur (https://github.com/seanmonstar/reqwest) +- **serde_json**: Copyright (c) 2014-2021, David Tolnay (https://github.com/serde-rs/json) +- **thiserror**: Copyright (c) 2019-2021, David Tolnay (https://github.com/dtolnay/thiserror) + +## .NET Ecosystem Dependencies + +This SDK and its integration test infrastructure use the following open source components: + +### Apache License 2.0 + +- **Grpc.AspNetCore**: Copyright 2014 gRPC authors (https://github.com/grpc/grpc-dotnet) +- **Grpc.Tools**: Copyright 2014 gRPC authors (https://github.com/grpc/grpc) + +### BSD 3-Clause License + +- **Google.Protobuf**: Copyright 2008 Google Inc. (https://github.com/protocolbuffers/protobuf) +- **NSubstitute**: Copyright (c) 2010-2024, NSubstitute Contributors (https://github.com/nsubstitute/NSubstitute) + +### MIT License + +- **Microsoft.SourceLink.GitHub**: Copyright (c) .NET Foundation and Contributors (https://github.com/dotnet/sourcelink) +- **Microsoft.NET.Test.Sdk**: Copyright (c) .NET Foundation and Contributors (https://github.com/microsoft/vstest) +- **NUnit**: Copyright (c) NUnit contributors (https://github.com/nunit/nunit) +- **NUnit3TestAdapter**: Copyright (c) NUnit contributors (https://github.com/nunit/nunit3-vs-adapter) diff --git a/dotnet/README.md b/dotnet/README.md new file mode 100644 index 00000000..b50024ea --- /dev/null +++ b/dotnet/README.md @@ -0,0 +1,443 @@ +# Databricks.Zerobus — .NET SDK + +High-performance .NET SDK for streaming data ingestion into Databricks Delta tables using the Zerobus service. Built on the same Rust core as the Go SDK, exposed via P/Invoke (C FFI bindings). + +## Requirements + +- **.NET 8** or **.NET 10** +- **Rust toolchain** (for building the native `zerobus_ffi` library from source) + +## Quick Start + +```csharp +using Databricks.Zerobus; + +// 1. Create SDK instance. +using var sdk = ZerobusSdk.CreateBuilder() + .Endpoint("https://your-shard.zerobus.databricks.com") + .UnityCatalogUrl("https://your-workspace.databricks.com") + .Build(); + +// 2. Configure stream options. +var options = StreamConfigurationOptions.Default with +{ + MaxInflightRequests = 50_000, +}; + +// 3. Create stream. +using var stream = sdk.CreateJsonStream( + "catalog.schema.table", + clientId, + clientSecret, + options); + +// 4. Ingest records. +long offset = stream.IngestRecord("""{"id": 1, "message": "Hello"}"""); + +// 5. Wait for acknowledgment. +stream.WaitForOffset(offset); +``` + +## Installation + +### NuGet (when published) + +```bash +dotnet add package Databricks.Zerobus +``` + +### From Source + +```bash +cd dotnet +dotnet build +``` + +The build automatically invokes `build_native.sh` to compile the Rust FFI shared library and place it in the correct `runtimes//native/` directory. You need `cargo` on your `PATH` (or in `~/.cargo/bin/`). + +To skip the automatic native build (e.g. when the library is pre-built): + +```bash +dotnet build -p:SkipNativeBuild=true +``` + +## API Reference + +### `ZerobusSdk` + +The main entry point. Manages the connection to Zerobus and Unity Catalog. + +```csharp +using var sdk = ZerobusSdk.CreateBuilder() + .Endpoint(zerobusEndpoint) + .UnityCatalogUrl(unityCatalogUrl) + .Build(); +``` + +#### `CreateJsonStream` + +Creates a JSON-only stream with OAuth 2.0 client credentials authentication. + +```csharp +using var stream = sdk.CreateJsonStream( + "catalog.schema.table", + clientId, + clientSecret, + options); // optional, defaults if null +``` + +#### `CreateProtoStream` + +Creates a protobuf-only stream with OAuth 2.0 client credentials authentication. + +```csharp +using var stream = sdk.CreateProtoStream( + "catalog.schema.table", + descriptorProto, + clientId, + clientSecret, + options); // optional, defaults if null +``` + +#### `CreateStream` + +Creates the legacy untyped stream with OAuth 2.0 client credentials authentication. +Use this only when you intentionally need access to both JSON and protobuf overloads. +The SDK now validates the `RecordType`/`DescriptorProto` combination before calling Rust. + +```csharp +using var stream = sdk.CreateStream( + new TableProperties("catalog.schema.table"), + clientId, + clientSecret, + options); // optional, defaults if null +``` + +#### `CreateJsonStreamWithHeadersProvider` + +Creates a JSON-only stream with custom authentication headers. + +```csharp +using var stream = sdk.CreateJsonStreamWithHeadersProvider( + "catalog.schema.table", + new MyHeadersProvider(), + options); // optional +``` + +#### `CreateProtoStreamWithHeadersProvider` + +Creates a protobuf-only stream with custom authentication headers. + +```csharp +using var stream = sdk.CreateProtoStreamWithHeadersProvider( + "catalog.schema.table", + descriptorProto, + new MyHeadersProvider(), + options); // optional +``` + +#### `CreateStreamWithHeadersProvider` + +Creates a stream with custom authentication headers. + +```csharp +using var stream = sdk.CreateStreamWithHeadersProvider( + new TableProperties("catalog.schema.table"), + new MyHeadersProvider(), + options); // optional +``` + +#### `RecreateStream` + +Recreates a stream for recovery after the original stream has failed or closed. + +Important: `RecreateStream(stream)` transfers ownership from `stream` to the returned +stream. The input `stream` is disposed during recreation and must not be used afterward. +A later `Dispose()` on the original wrapper (for example at the end of a `using` scope) +is a no-op. + +### `JsonZerobusStream` and `ProtoZerobusStream` + +Typed stream wrappers expose only the matching ingest overloads while preserving the +same lifecycle APIs (`Flush`, `WaitForOffset`, `Close`, `Dispose`, `GetUnackedRecords`). + +#### `JsonZerobusStream.IngestRecord` + +```csharp +stream.IngestRecord("""{"field": "value"}"""); +stream.Flush(); +``` + +#### `ProtoZerobusStream.IngestRecord` + +```csharp +byte[] protoBytes = myMessage.ToByteArray(); +stream.IngestRecord(protoBytes); +stream.Flush(); +``` + +### `ZerobusStream` + +The untyped stream remains available for advanced callers. Thread-safe. + +#### `IngestRecord` + +Ingests a single record and returns its offset immediately (acknowledgment happens in background). +If you use this untyped API, JSON streams must set `RecordType.Json` and use the string overloads. +Proto streams must provide `DescriptorProto` and use the byte-oriented overloads. + +```csharp +// JSON +long offset = stream.IngestRecord("""{"field": "value"}"""); + +// Protobuf +byte[] protoBytes = myMessage.ToByteArray(); +long offset = stream.IngestRecord(protoBytes); +stream.WaitForOffset(offset); +``` + +#### `IngestRecords` + +Ingests a batch of records and returns one offset for the whole batch. + +```csharp +string[] records = [ + """{"device": "sensor-001", "temp": 20}""", + """{"device": "sensor-002", "temp": 21}""", +]; +long lastOffset = stream.IngestRecords(records); +stream.WaitForOffset(lastOffset); +``` + +#### `WaitForOffset` (sync) + +Blocks until a specific offset is acknowledged by the server. + +```csharp +stream.WaitForOffset(offset); +``` + +#### `Flush` (sync) + +Blocks until all pending records are acknowledged. + +```csharp +stream.Flush(); +``` + +#### `GetUnackedRecords` + +Retrieves unacknowledged records after stream failure (call after close/failure only). + +```csharp +ReadOnlyMemory[] unacked = stream.GetUnackedRecords(); +foreach (var payload in unacked) +{ + // Decode as UTF-8 if you know this stream ingests JSON. + Console.WriteLine($"record bytes: {payload.Length}"); +} +``` + +#### `Close` / `Dispose` / `DisposeAsync` + +`Close()` gracefully closes the stream (flushes first) but keeps the stream +readable for recovery (`GetUnackedRecords`, `RecreateStream`). +If you call `RecreateStream(stream)`, the original stream object is disposed and +no longer usable. +`Dispose()` frees native resources and should be called when recovery work is done. +`using` calls `Dispose()` automatically. + +```csharp +stream.Close(); +// or let `using` / `await using` handle cleanup +``` + +### `IHeadersProvider` + +Interface for custom authentication. + +```csharp +public class CustomHeadersProvider : IHeadersProvider +{ + public IDictionary GetHeaders() + { + return new Dictionary + { + ["authorization"] = "Bearer " + GetToken(), + ["x-databricks-zerobus-table-name"] = "catalog.schema.table", + }; + } +} +``` + +### `StreamConfigurationOptions` + +Use C# record `with` expressions to customise: + +```csharp +var options = StreamConfigurationOptions.Default with +{ + MaxInflightRequests = 50_000, + RecoveryRetries = 10, +}; +``` + +| Property | Default | Description | +| --------------------------- | --------- | ---------------------------- | +| `MaxInflightRequests` | 1,000,000 | Backpressure control | +| `Recovery` | `true` | Auto-recovery on failures | +| `RecoveryTimeoutMs` | 15,000 | Timeout per recovery attempt | +| `RecoveryBackoffMs` | 2,000 | Delay between retries | +| `RecoveryRetries` | 4 | Max recovery attempts | +| `ServerLackOfAckTimeoutMs` | 60,000 | Server ack timeout | +| `FlushTimeoutMs` | 300,000 | Flush timeout (5 min) | +| `RecordType` | `Proto` | Proto / Json / Unspecified | +| `StreamPausedMaxWaitTimeMs` | `null` | Graceful close wait time | + +Typed factories set `RecordType` automatically. You only need to set it manually when using +the untyped `CreateStream` or `CreateStreamWithHeadersProvider` APIs. + +### Error Handling + +Errors throw `ZerobusException` with an `IsRetryable` property: + +```csharp +try +{ + long offset = stream.IngestRecord(data); +} +catch (ZerobusException ex) when (ex.IsRetryable) +{ + // Transient error — SDK auto-recovers when Recovery is enabled. + Console.WriteLine($"Retryable error: {ex.RawMessage}"); +} +catch (ZerobusException ex) +{ + // Fatal error — manual intervention needed. + Console.WriteLine($"Fatal error: {ex.RawMessage}"); +} +``` + +## Native Library Setup + +The native `zerobus_ffi` shared library is built automatically when you run `dotnet build`. The MSBuild target invokes `build_native.sh`, which: + +1. Detects your OS and architecture +2. Runs `cargo build --release` in the `zerobus-ffi` crate +3. Copies the shared library (`.dylib` / `.so` / `.dll`) to `src/Zerobus/runtimes//native/` +4. Skips the rebuild if the library is already up to date + +### Manual Build + +You can also run the script directly: + +```bash +cd dotnet +./build_native.sh # Build for current platform +./build_native.sh --force # Force rebuild +``` + +### Runtime Directories + +The native library is placed in the standard .NET runtime identifier layout: + +| Platform | Path | +| ----------- | ------------------------------------------------ | +| Linux x64 | `runtimes/linux-x64/native/libzerobus_ffi.so` | +| Linux arm64 | `runtimes/linux-arm64/native/libzerobus_ffi.so` | +| macOS x64 | `runtimes/osx-x64/native/libzerobus_ffi.dylib` | +| macOS arm64 | `runtimes/osx-arm64/native/libzerobus_ffi.dylib` | +| Windows x64 | `runtimes/win-x64/native/zerobus_ffi.dll` | + +## Testing + +### Unit Tests + +Unit tests are isolated and do not require the native library: + +```bash +dotnet test tests/Zerobus.Tests +``` + +### Integration Tests + +Integration tests spin up a mock gRPC server (per test) and exercise the full SDK through the native FFI layer. They require the Rust toolchain to build the native library: + +```bash +dotnet test tests/Zerobus.IntegrationTests +``` + +The integration tests cover: + +| Test | Scenario | +| --------------------------------------------------- | ------------------------------------------ | +| `SuccessfulStreamCreation` | Stream creation succeeds | +| `TimeoutedStreamCreation` | Timeout when server responds slowly | +| `NonRetriableErrorDuringStreamCreation` | Non-retriable error (e.g. Unauthenticated) | +| `RetriableErrorWithoutRecoveryDuringStreamCreation` | Retriable error with recovery disabled | +| `GracefulClose` | Ingest record then close gracefully | +| `IdempotentClose` | Multiple `Close()` calls succeed | +| `IngestAfterClose` | Ingest after close throws | +| `IngestSingleRecord` | Single record ingest and ack | +| `IngestMultipleRecords` | Multiple sequential records with ack | +| `IngestBatchRecords` | Batch ingest of 5 records | +| `IngestRecordsAfterClose` | Batch ingest after close throws | + +Each test gets its own mock gRPC server on a unique port, so all tests run in parallel. + +### Running All Tests + +```bash +dotnet test +``` + +## Project Structure + +``` +dotnet/ +├── Zerobus.slnx # Solution file +├── Directory.Build.props # Shared build settings +├── build_native.sh # Rust FFI build script +├── README.md +├── src/ +│ └── Zerobus/ # Main SDK library +│ ├── Zerobus.csproj +│ ├── ZerobusSdk.cs # SDK entry point (IDisposable) +│ ├── ZerobusStream.cs # Stream for record ingestion (IDisposable) +│ ├── ZerobusException.cs # Error type with IsRetryable +│ ├── IHeadersProvider.cs # Custom auth interface +│ ├── RecordType.cs # Proto / Json / Unspecified enum +│ ├── StreamConfigurationOptions.cs # Config record with defaults +│ ├── TableProperties.cs # Table name + optional descriptor +│ ├── Properties/ +│ │ └── AssemblyInfo.cs +│ └── Native/ # P/Invoke layer (internal) +│ ├── NativeBindings.cs # Raw DllImport declarations +│ ├── NativeInterop.cs # Safe wrappers + marshalling +│ └── HeadersProviderBridge.cs # Managed→native callback bridge +├── tests/ +│ ├── Zerobus.Tests/ # Unit tests (NUnit) +│ └── Zerobus.IntegrationTests/ # Integration tests (NUnit + gRPC mock) +│ ├── Zerobus.IntegrationTests.csproj +│ ├── IntegrationTests.cs # 11 integration tests +│ ├── MockZerobusServer.cs # Mock gRPC server +│ ├── TestHelpers.cs # Fixtures, response builders, interceptor +│ └── Protos/ +│ └── zerobus_service.proto # Proto definition for gRPC stubs +└── examples/ + ├── JsonSingle/ # Single JSON record ingestion + ├── JsonBatch/ # Batch JSON record ingestion + ├── ProtoSingle/ # Single protobuf record ingestion +``` + +## Architecture + +``` +.NET SDK (Databricks.Zerobus) + ↓ P/Invoke +Rust FFI (zerobus-ffi / libzerobus_ffi) + ↓ +Rust Core (databricks-zerobus-ingest-sdk) + ↓ gRPC +Zerobus Service +``` diff --git a/dotnet/Zerobus.slnx b/dotnet/Zerobus.slnx new file mode 100644 index 00000000..9e7da260 --- /dev/null +++ b/dotnet/Zerobus.slnx @@ -0,0 +1,14 @@ + + + + + + + + + + + + + + diff --git a/dotnet/build_native.sh b/dotnet/build_native.sh new file mode 100755 index 00000000..7fbfce97 --- /dev/null +++ b/dotnet/build_native.sh @@ -0,0 +1,189 @@ +#!/bin/bash +set -e + +# ───────────────────────────────────────────────────────────────────────────── +# build_native.sh — Build the zerobus_ffi shared library for the .NET SDK +# +# This script builds the Rust FFI crate as a cdylib (shared library) and places +# the output in the dotnet/src/Zerobus/runtimes//native/ directory, which +# is the standard .NET layout for native P/Invoke libraries. +# +# Usage: +# ./build_native.sh # Build for the current platform +# ./build_native.sh --force # Rebuild even if up to date +# ./build_native.sh --all # Build all RIDs for the current OS +# +# The script is also invoked automatically by the .NET build (via MSBuild +# targets in Zerobus.csproj) so that `dotnet build` and `dotnet test` always +# have the native library available. +# ───────────────────────────────────────────────────────────────────────────── + +SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" +REPO_ROOT="$SCRIPT_DIR/.." +FFI_DIR="$REPO_ROOT/rust/ffi" +RUNTIMES_DIR="$SCRIPT_DIR/src/Zerobus/runtimes" + +FORCE=0 +BUILD_ALL=0 +for arg in "$@"; do + case "$arg" in + --force) FORCE=1 ;; + --all) BUILD_ALL=1 ;; + *) + echo "✗ Unknown argument: $arg" + exit 1 + ;; + esac +done + +# ── Detect platform ────────────────────────────────────────────────────────── + +OS="$(uname -s)" +ARCH="$(uname -m)" + +case "$OS" in + Darwin) + LIB_NAME="libzerobus_ffi.dylib" + if [[ $BUILD_ALL -eq 1 ]]; then + RIDS=("osx-arm64" "osx-x64") + RUST_TARGETS=("aarch64-apple-darwin" "x86_64-apple-darwin") + else + case "$ARCH" in + arm64) RIDS=("osx-arm64"); RUST_TARGETS=("aarch64-apple-darwin") ;; + x86_64) RIDS=("osx-x64"); RUST_TARGETS=("x86_64-apple-darwin") ;; + *) echo "✗ Unsupported macOS architecture: $ARCH"; exit 1 ;; + esac + fi + ;; + Linux) + LIB_NAME="libzerobus_ffi.so" + if [[ $BUILD_ALL -eq 1 ]]; then + RIDS=("linux-arm64" "linux-x64") + RUST_TARGETS=("aarch64-unknown-linux-gnu" "x86_64-unknown-linux-gnu") + else + case "$ARCH" in + aarch64|arm64) RIDS=("linux-arm64"); RUST_TARGETS=("aarch64-unknown-linux-gnu") ;; + x86_64) RIDS=("linux-x64"); RUST_TARGETS=("x86_64-unknown-linux-gnu") ;; + *) echo "✗ Unsupported Linux architecture: $ARCH"; exit 1 ;; + esac + fi + ;; + MINGW*|MSYS*|CYGWIN*) + if [[ $BUILD_ALL -eq 1 ]]; then + echo "✗ --all is not supported on Windows" + exit 1 + fi + RIDS=("win-x64") + RUST_TARGETS=("x86_64-pc-windows-gnu") + LIB_NAME="zerobus_ffi.dll" + ;; + *) + echo "✗ Unsupported OS: $OS" + exit 1 + ;; +esac + +echo "┌──────────────────────────────────────────────────────────────────┐" +echo "│ Zerobus .NET — Native Library Build │" +echo "├──────────────────────────────────────────────────────────────────┤" +echo "│ Platform: $OS / $ARCH" +if [[ ${#RIDS[@]} -eq 1 ]]; then + echo "│ RID: ${RIDS[0]}" +else + echo "│ RIDs: ${RIDS[*]}" +fi +echo "│ Library: $LIB_NAME" +echo "│ Output: $RUNTIMES_DIR//native/" +echo "└──────────────────────────────────────────────────────────────────┘" + +# ── Check for Rust toolchain ───────────────────────────────────────────────── + +if ! command -v cargo &>/dev/null; then + # Try sourcing the standard cargo env + if [[ -f "$HOME/.cargo/env" ]]; then + # shellcheck disable=SC1091 + source "$HOME/.cargo/env" + fi +fi + +if ! command -v cargo &>/dev/null; then + echo "" + echo "✗ cargo not found. Install the Rust toolchain:" + echo " curl --proto '=https' --tlsv1.2 -sSf https://sh.rustup.rs | sh" + exit 1 +fi + +# ── Check if rebuild is needed ─────────────────────────────────────────────── + +if [[ $FORCE -eq 0 ]]; then + ALL_UP_TO_DATE=1 + for rid in "${RIDS[@]}"; do + target_dir="$RUNTIMES_DIR/$rid/native" + target_path="$target_dir/$LIB_NAME" + + if [[ ! -f "$target_path" ]]; then + ALL_UP_TO_DATE=0 + break + fi + + if [[ "$FFI_DIR/Cargo.toml" -nt "$target_path" ]]; then + ALL_UP_TO_DATE=0 + break + fi + + while IFS= read -r -d '' file; do + if [[ "$file" -nt "$target_path" ]]; then + ALL_UP_TO_DATE=0 + break 2 + fi + done < <(find "$FFI_DIR/src" -name "*.rs" -print0 2>/dev/null) + done + + if [[ $ALL_UP_TO_DATE -eq 1 ]]; then + echo "" + echo "✓ Native libraries are up to date." + echo " (use --force to rebuild anyway)" + exit 0 + fi +fi + +# ── Build ──────────────────────────────────────────────────────────────────── + +echo "" +echo "Building zerobus-ffi (release)..." +echo "" + +cd "$FFI_DIR" + +build_target() { + rid="$1" + rust_target="$2" + target_dir="$RUNTIMES_DIR/$rid/native" + target_path="$target_dir/$LIB_NAME" + + echo "→ Building $rid ($rust_target)" + cargo build --release --target "$rust_target" + cargo_out="../target/$rust_target/release/$LIB_NAME" + + if [[ ! -f "$cargo_out" ]]; then + echo "" + echo "✗ Build succeeded but library not found at: $cargo_out" + echo " Expected shared library ($LIB_NAME) — check Cargo.toml has:" + echo ' crate-type = ["staticlib", "cdylib"]' + exit 1 + fi + + mkdir -p "$target_dir" + cp "$cargo_out" "$target_path" +} + +for i in "${!RIDS[@]}"; do + build_target "${RIDS[$i]}" "${RUST_TARGETS[$i]}" +done + +# ── Copy to runtimes ───────────────────────────────────────────────────────── + +echo "" +echo "✓ Native libraries built and placed under:" +echo " $RUNTIMES_DIR//native/" +echo "" diff --git a/dotnet/examples/JsonBatch/JsonBatch.csproj b/dotnet/examples/JsonBatch/JsonBatch.csproj new file mode 100644 index 00000000..ce333cf0 --- /dev/null +++ b/dotnet/examples/JsonBatch/JsonBatch.csproj @@ -0,0 +1,12 @@ + + + + Exe + net8.0 + + + + + + + diff --git a/dotnet/examples/JsonBatch/Program.cs b/dotnet/examples/JsonBatch/Program.cs new file mode 100644 index 00000000..8627f531 --- /dev/null +++ b/dotnet/examples/JsonBatch/Program.cs @@ -0,0 +1,52 @@ +using Databricks.Zerobus; + +// Get configuration from environment. +var zerobusEndpoint = Environment.GetEnvironmentVariable("ZEROBUS_SERVER_ENDPOINT") + ?? throw new InvalidOperationException("ZEROBUS_SERVER_ENDPOINT not set"); +var unityCatalogUrl = Environment.GetEnvironmentVariable("DATABRICKS_WORKSPACE_URL") + ?? throw new InvalidOperationException("DATABRICKS_WORKSPACE_URL not set"); +var clientId = Environment.GetEnvironmentVariable("DATABRICKS_CLIENT_ID") + ?? throw new InvalidOperationException("DATABRICKS_CLIENT_ID not set"); +var clientSecret = Environment.GetEnvironmentVariable("DATABRICKS_CLIENT_SECRET") + ?? throw new InvalidOperationException("DATABRICKS_CLIENT_SECRET not set"); +var tableName = Environment.GetEnvironmentVariable("ZEROBUS_TABLE_NAME") + ?? throw new InvalidOperationException("ZEROBUS_TABLE_NAME not set"); + +// Create SDK instance. +using var sdk = ZerobusSdk.CreateBuilder() + .Endpoint(zerobusEndpoint) + .UnityCatalogUrl(unityCatalogUrl) + .Build(); + +// Configure stream options (optional). +var options = StreamConfigurationOptions.Default with +{ + MaxInflightRequests = 50_000, +}; + +// Create JSON stream. +using var stream = sdk.CreateJsonStream( + tableName, + clientId, + clientSecret, + options); + +Console.WriteLine("Ingesting batch of records..."); + +string[] batchRecords = +[ + """{"device_name": "sensor-001", "temp": 20, "humidity": 60}""", + """{"device_name": "sensor-002", "temp": 21, "humidity": 61}""", + """{"device_name": "sensor-003", "temp": 22, "humidity": 62}""", + """{"device_name": "sensor-004", "temp": 23, "humidity": 63}""", + """{"device_name": "sensor-005", "temp": 24, "humidity": 64}""", +]; + +long lastOffset = stream.IngestRecords(batchRecords); +Console.WriteLine($"Batch of {batchRecords.Length} records ingested, last offset: {lastOffset}"); + +// Wait for the last offset to ensure the entire batch is acknowledged. +stream.WaitForOffset(lastOffset); +Console.WriteLine("Batch acknowledged!"); + +Console.WriteLine("All operations completed successfully!"); diff --git a/dotnet/examples/JsonSingle/JsonSingle.csproj b/dotnet/examples/JsonSingle/JsonSingle.csproj new file mode 100644 index 00000000..ce333cf0 --- /dev/null +++ b/dotnet/examples/JsonSingle/JsonSingle.csproj @@ -0,0 +1,12 @@ + + + + Exe + net8.0 + + + + + + + diff --git a/dotnet/examples/JsonSingle/Program.cs b/dotnet/examples/JsonSingle/Program.cs new file mode 100644 index 00000000..e349bec3 --- /dev/null +++ b/dotnet/examples/JsonSingle/Program.cs @@ -0,0 +1,72 @@ +using Databricks.Zerobus; + +// Get configuration from environment. +var zerobusEndpoint = Environment.GetEnvironmentVariable("ZEROBUS_SERVER_ENDPOINT") + ?? throw new InvalidOperationException("ZEROBUS_SERVER_ENDPOINT not set"); +var unityCatalogUrl = Environment.GetEnvironmentVariable("DATABRICKS_WORKSPACE_URL") + ?? throw new InvalidOperationException("DATABRICKS_WORKSPACE_URL not set"); +var clientId = Environment.GetEnvironmentVariable("DATABRICKS_CLIENT_ID") + ?? throw new InvalidOperationException("DATABRICKS_CLIENT_ID not set"); +var clientSecret = Environment.GetEnvironmentVariable("DATABRICKS_CLIENT_SECRET") + ?? throw new InvalidOperationException("DATABRICKS_CLIENT_SECRET not set"); +var tableName = Environment.GetEnvironmentVariable("ZEROBUS_TABLE_NAME") + ?? throw new InvalidOperationException("ZEROBUS_TABLE_NAME not set"); + +// Create SDK instance. +using var sdk = ZerobusSdk.CreateBuilder() + .Endpoint(zerobusEndpoint) + .UnityCatalogUrl(unityCatalogUrl) + .Build(); + +// Configure stream options (optional). +var options = StreamConfigurationOptions.Default with +{ + MaxInflightRequests = 50_000, +}; + +// Create JSON stream. +using var stream = sdk.CreateJsonStream( + tableName, + clientId, + clientSecret, + options); + +Console.WriteLine("Ingesting records..."); +var offsets = new List(); + +for (int i = 0; i < 5; i++) +{ + // Change this JSON to match the schema of your table. + var jsonRecord = """ + { + "device_name": "sensor-001", + "temp": 20, + "humidity": 60 + } + """; + + try + { + long offset = stream.IngestRecord(jsonRecord); + Console.WriteLine($"Ingested record {i} at offset {offset}"); + offsets.Add(offset); + } + catch (ZerobusException ex) when (ex.IsRetryable) + { + Console.WriteLine($"Failed to ingest record {i} (retryable): {ex.RawMessage}"); + } + catch (ZerobusException ex) + { + Console.WriteLine($"Failed to ingest record {i}: {ex.RawMessage}"); + } +} + +// Wait for specific offsets to be acknowledged. +Console.WriteLine("Waiting for acknowledgments..."); +foreach (var offset in offsets) +{ + stream.WaitForOffset(offset); + Console.WriteLine($"Record at offset {offset} acknowledged"); +} + +Console.WriteLine("All records successfully ingested and acknowledged!"); diff --git a/dotnet/examples/ProtoSingle/Program.cs b/dotnet/examples/ProtoSingle/Program.cs new file mode 100644 index 00000000..aa8d4353 --- /dev/null +++ b/dotnet/examples/ProtoSingle/Program.cs @@ -0,0 +1,54 @@ +using Databricks.Zerobus; + +// Get configuration from environment. +var zerobusEndpoint = Environment.GetEnvironmentVariable("ZEROBUS_SERVER_ENDPOINT") + ?? throw new InvalidOperationException("ZEROBUS_SERVER_ENDPOINT not set"); +var unityCatalogUrl = Environment.GetEnvironmentVariable("DATABRICKS_WORKSPACE_URL") + ?? throw new InvalidOperationException("DATABRICKS_WORKSPACE_URL not set"); +var clientId = Environment.GetEnvironmentVariable("DATABRICKS_CLIENT_ID") + ?? throw new InvalidOperationException("DATABRICKS_CLIENT_ID not set"); +var clientSecret = Environment.GetEnvironmentVariable("DATABRICKS_CLIENT_SECRET") + ?? throw new InvalidOperationException("DATABRICKS_CLIENT_SECRET not set"); +var tableName = Environment.GetEnvironmentVariable("ZEROBUS_TABLE_NAME") + ?? throw new InvalidOperationException("ZEROBUS_TABLE_NAME not set"); + +// In a real application, you would load the DescriptorProto from your compiled .proto file. +// For example, using Google.Protobuf.Reflection: +// var descriptor = MyMessage.Descriptor.File.SerializedData.ToByteArray(); +byte[] descriptorProto = []; // Replace with your actual descriptor bytes. + +// Create SDK instance. +using var sdk = ZerobusSdk.CreateBuilder() + .Endpoint(zerobusEndpoint) + .UnityCatalogUrl(unityCatalogUrl) + .Build(); + +// Configure stream options for protobuf. +var options = StreamConfigurationOptions.Default with +{ + MaxInflightRequests = 50_000, +}; + +// Create protobuf stream. +using var stream = sdk.CreateProtoStream( + tableName, + descriptorProto, + clientId, + clientSecret, + options); + +Console.WriteLine("Ingesting protobuf records..."); + +for (int i = 0; i < 5; i++) +{ + // In a real application, serialize your protobuf message: + // byte[] protoBytes = myMessage.ToByteArray(); + byte[] protoBytes = [0x08, 0x01]; // Placeholder — replace with real protobuf data. + + long offset = stream.IngestRecord(protoBytes); + Console.WriteLine($"Ingested protobuf record {i} at offset {offset}"); +} + +// Flush all pending records. +stream.Flush(); +Console.WriteLine("All protobuf records flushed and acknowledged!"); diff --git a/dotnet/examples/ProtoSingle/ProtoSingle.csproj b/dotnet/examples/ProtoSingle/ProtoSingle.csproj new file mode 100644 index 00000000..ce333cf0 --- /dev/null +++ b/dotnet/examples/ProtoSingle/ProtoSingle.csproj @@ -0,0 +1,12 @@ + + + + Exe + net8.0 + + + + + + + diff --git a/dotnet/src/Zerobus/IHeadersProvider.cs b/dotnet/src/Zerobus/IHeadersProvider.cs new file mode 100644 index 00000000..07fa0ec0 --- /dev/null +++ b/dotnet/src/Zerobus/IHeadersProvider.cs @@ -0,0 +1,32 @@ +namespace Databricks.Zerobus; + +/// +/// Interface for providing custom authentication headers. +/// Implement this interface to supply custom authentication logic +/// (e.g. fetching tokens from a vault, using managed identity, etc.). +/// +/// +/// +/// public class CustomHeadersProvider : IHeadersProvider +/// { +/// public IDictionary<string, string> GetHeaders() +/// { +/// return new Dictionary<string, string> +/// { +/// ["authorization"] = "Bearer " + GetToken(), +/// ["x-databricks-zerobus-table-name"] = "catalog.schema.table", +/// }; +/// } +/// } +/// +/// +public interface IHeadersProvider +{ + /// + /// Returns the headers to be used for authentication. + /// This method is called by the SDK whenever authentication is needed. + /// + /// A dictionary of header key-value pairs. + /// If the headers cannot be obtained. + IDictionary GetHeaders(); +} diff --git a/dotnet/src/Zerobus/JsonZerobusStream.cs b/dotnet/src/Zerobus/JsonZerobusStream.cs new file mode 100644 index 00000000..66d12b0e --- /dev/null +++ b/dotnet/src/Zerobus/JsonZerobusStream.cs @@ -0,0 +1,72 @@ +namespace Databricks.Zerobus; + +/// +/// A stream that accepts JSON payloads only. +/// +public sealed class JsonZerobusStream : TypedZerobusStream +{ + internal JsonZerobusStream(ZerobusStream innerStream) + : base(innerStream) + { + } + + /// + public long IngestRecord(string payload) + { + return InnerStream.IngestRecord(payload); + } + + /// + public Task IngestRecordAsync(string payload) + { + return InnerStream.IngestRecordAsync(payload); + } + + /// + public long IngestRecords(string[] records) + { + return InnerStream.IngestRecords(records); + } + + /// + public Task IngestRecordsAsync(string[] records) + { + return InnerStream.IngestRecordsAsync(records); + } + + /// + /// Retrieves all records that have not yet been acknowledged by the server. + /// + /// Important: This should only be called after the stream has + /// closed or failed. Calling it on an active stream will return an error. + /// + /// + /// + /// An array of JSON record payloads as strings, decoded from UTF-8. + /// + /// Thrown if retrieval fails. + /// Thrown if the stream has been disposed. + public string[] GetUnackedRecords() + { + var records = InnerStream.GetUnackedRecords(); + return Array.ConvertAll(records, memory => System.Text.Encoding.UTF8.GetString(memory.Span)); + } + + /// + /// Asynchronously retrieves all records that have not yet been acknowledged by the server. + /// + /// Important: This should only be called after the stream has + /// closed or failed. Calling it on an active stream will return an error. + /// + /// + /// + /// A task that resolves to an array of JSON record payloads as strings, decoded from UTF-8. + /// + /// Thrown if retrieval fails. + /// Thrown if the stream has been disposed. + public async Task GetUnackedRecordsAsync() + { + var records = await InnerStream.GetUnackedRecordsAsync(); + return Array.ConvertAll(records, memory => System.Text.Encoding.UTF8.GetString(memory.Span)); + } +} \ No newline at end of file diff --git a/dotnet/src/Zerobus/Native/HeadersProviderBridge.cs b/dotnet/src/Zerobus/Native/HeadersProviderBridge.cs new file mode 100644 index 00000000..557f867a --- /dev/null +++ b/dotnet/src/Zerobus/Native/HeadersProviderBridge.cs @@ -0,0 +1,117 @@ +// Bridges the managed IHeadersProvider interface to the native callback expected by Rust. +// This is the .NET equivalent of the goGetHeaders / cHeadersCallback pattern in ffi.go. + +using System.Text; + +namespace Databricks.Zerobus.Native; + +/// +/// Bridges a managed to the native +/// function pointer. +/// +internal sealed class HeadersProviderBridge +{ + private readonly IHeadersProvider _provider; + + public HeadersProviderBridge(IHeadersProvider provider) + { + _provider = provider; + } + + /// + /// The native callback that can be passed as a function pointer to Rust. + /// Called from native code on the Rust runtime thread. + /// + public unsafe CHeaders NativeCallback(IntPtr userData) + { + var result = new CHeaders(); + IntPtr arrayPtr = IntPtr.Zero; + nuint arrayCount = 0; + int filledCount = 0; + + try + { + var headers = _provider.GetHeaders(); + + if (headers is null || headers.Count == 0) + { + result.Headers = IntPtr.Zero; + result.Count = 0; + result.ErrorMessage = IntPtr.Zero; + return result; + } + + // Allocate via Rust FFI so allocation/free always happen on the same heap. + arrayCount = (nuint)headers.Count; + arrayPtr = NativeMethods.AllocHeaderArray(arrayCount); + if (arrayPtr == IntPtr.Zero) + { + throw new OutOfMemoryException("Failed to allocate headers array"); + } + + var headerArray = (CHeader*)arrayPtr; + foreach (var (key, value) in headers) + { + var cKey = AllocUtf8String(key); + if (cKey == IntPtr.Zero) + { + throw new OutOfMemoryException("Failed to allocate header key string"); + } + + var cValue = AllocUtf8String(value); + if (cValue == IntPtr.Zero) + { + throw new OutOfMemoryException("Failed to allocate header value string"); + } + + headerArray[filledCount] = new CHeader + { + Key = cKey, + Value = cValue, + }; + filledCount++; + } + + result.Headers = arrayPtr; + result.Count = arrayCount; + result.ErrorMessage = IntPtr.Zero; + arrayPtr = IntPtr.Zero; + arrayCount = 0; + } + catch (Exception ex) + { + if (arrayPtr != IntPtr.Zero) + { + // Only free the headers that were actually filled, not the entire allocated count. + // This prevents attempting to free uninitialized memory if an allocation failed mid-loop. + NativeMethods.FreeHeaders(new CHeaders + { + Headers = arrayPtr, + Count = (nuint)filledCount, + ErrorMessage = IntPtr.Zero, + }); + } + + result.Headers = IntPtr.Zero; + result.Count = 0; + result.ErrorMessage = AllocUtf8String(ex.Message); + } + + return result; + } + + private static unsafe IntPtr AllocUtf8String(string s) + { + var byteCount = Encoding.UTF8.GetByteCount(s); + if (byteCount == 0) + { + return NativeMethods.AllocCString(null, 0); + } + + var utf8 = Encoding.UTF8.GetBytes(s); + fixed (byte* ptr = utf8) + { + return NativeMethods.AllocCString(ptr, (nuint)utf8.Length); + } + } +} diff --git a/dotnet/src/Zerobus/Native/NativeBindings.cs b/dotnet/src/Zerobus/Native/NativeBindings.cs new file mode 100644 index 00000000..e445eb8b --- /dev/null +++ b/dotnet/src/Zerobus/Native/NativeBindings.cs @@ -0,0 +1,461 @@ +// P/Invoke bindings to the Rust FFI layer (zerobus-ffi). +// This is the .NET equivalent of ffi.go in the Go SDK. + +using System.Reflection; +using System.Runtime.InteropServices; + +namespace Databricks.Zerobus.Native; + +/// +/// A single header key-value pair for C FFI. +/// +[StructLayout(LayoutKind.Sequential)] +internal struct CHeader +{ + public IntPtr Key; // char* + public IntPtr Value; // char* +} + +/// +/// A collection of headers returned from a managed callback. +/// +[StructLayout(LayoutKind.Sequential)] +internal struct CHeaders +{ + public IntPtr Headers; // CHeader* + public nuint Count; + public IntPtr ErrorMessage; // char* +} + +/// +/// Opaque SDK handle. We only ever hold pointers to this. +/// +[StructLayout(LayoutKind.Sequential)] +internal struct CZerobusSdk +{ + // Opaque - zero-sized in C, only used via pointer. +} + +/// +/// Result struct returned by most FFI calls. +/// +[StructLayout(LayoutKind.Sequential)] +internal struct CResult +{ + [MarshalAs(UnmanagedType.U1)] + public bool Success; + + public IntPtr ErrorMessage; // char* — must be freed with zerobus_free_error_message + + [MarshalAs(UnmanagedType.U1)] + public bool IsRetryable; +} + +/// +/// Opaque stream handle. +/// +[StructLayout(LayoutKind.Sequential)] +internal struct CZerobusStream +{ + // Opaque. +} + +/// +/// Stream configuration options passed to the native layer. +/// +[StructLayout(LayoutKind.Sequential)] +internal struct CStreamConfigurationOptions +{ + public nuint MaxInflightRequests; + + [MarshalAs(UnmanagedType.U1)] + public bool Recovery; + + public ulong RecoveryTimeoutMs; + public ulong RecoveryBackoffMs; + public uint RecoveryRetries; + public ulong ServerLackOfAckTimeoutMs; + public ulong FlushTimeoutMs; + public int RecordType; + public ulong StreamPausedMaxWaitTimeMs; + + [MarshalAs(UnmanagedType.U1)] + public bool HasStreamPausedMaxWaitTimeMs; + + public ulong CallbackMaxWaitTimeMs; + + [MarshalAs(UnmanagedType.U1)] + public bool HasCallbackMaxWaitTimeMs; + + // Optional ack callback function pointers and user data. + // Keep null for .NET for now; this preserves ABI layout with Rust FFI. + public IntPtr AckOnAck; + public IntPtr AckOnError; + public IntPtr AckUserData; +} + +/// +/// Represents a single record (either Proto or JSON) returned by get_unacked_records. +/// +[StructLayout(LayoutKind.Sequential)] +internal struct CRecord +{ + [MarshalAs(UnmanagedType.U1)] + public bool IsJson; + + public IntPtr Data; // uint8_t* + public nuint DataLen; +} + +/// +/// An array of records returned by get_unacked_records. +/// +[StructLayout(LayoutKind.Sequential)] +internal struct CRecordArray +{ + public IntPtr Records; // CRecord* + public nuint Len; +} + +/// +/// Callback signature for the headers provider. +/// Matches: CHeaders (*HeadersProviderCallback)(void* user_data) +/// +[UnmanagedFunctionPointer(CallingConvention.Cdecl)] +internal delegate CHeaders HeadersProviderCallback(IntPtr userData); + +/// +/// Callback for async stream creation completion. +/// Matches: void (*CreateStreamAsyncCallback)(CZerobusStream* stream, const CResult* result, void* user_data) +/// +[UnmanagedFunctionPointer(CallingConvention.Cdecl)] +internal delegate void CreateStreamAsyncCallback(IntPtr stream, IntPtr result, IntPtr userData); + +/// +/// Callback for async offset-returning operations. +/// Matches: void (*OffsetAsyncCallback)(int64_t offset, const CResult* result, void* user_data) +/// +[UnmanagedFunctionPointer(CallingConvention.Cdecl)] +internal delegate void OffsetAsyncCallback(long offset, IntPtr result, IntPtr userData); + +/// +/// Callback for async bool-returning operations. +/// Matches: void (*BoolAsyncCallback)(bool value, const CResult* result, void* user_data) +/// +[UnmanagedFunctionPointer(CallingConvention.Cdecl)] +internal delegate void BoolAsyncCallback([MarshalAs(UnmanagedType.U1)] bool value, IntPtr result, IntPtr userData); + +/// +/// Callback for async record-array-returning operations. +/// Matches: void (*RecordArrayAsyncCallback)(CRecordArray records, const CResult* result, void* user_data) +/// +[UnmanagedFunctionPointer(CallingConvention.Cdecl)] +internal delegate void RecordArrayAsyncCallback(CRecordArray records, IntPtr result, IntPtr userData); + +/// +/// P/Invoke declarations for the zerobus_ffi native library. +/// +internal static partial class NativeMethods +{ + private const string LibName = "zerobus_ffi"; + + static NativeMethods() => NativeLibrary.SetDllImportResolver(typeof(NativeMethods).Assembly, ResolveLibrary); + + private static IntPtr ResolveLibrary( + string libraryName, + Assembly assembly, + DllImportSearchPath? searchPath) + { + if (!string.Equals(libraryName, LibName, StringComparison.Ordinal)) + { + return IntPtr.Zero; + } + + var fileName = GetLibraryFileName(); + + if (NativeLibrary.TryLoad(fileName, assembly, searchPath, out var handle)) + { + return handle; + } + + var baseDir = AppContext.BaseDirectory; + var rid = RuntimeInformation.RuntimeIdentifier; + var candidate = Path.Combine(baseDir, "runtimes", rid, "native", fileName); + + return NativeLibrary.TryLoad(candidate, out handle) ? handle : IntPtr.Zero; + } + + private static string GetLibraryFileName() + { + if (RuntimeInformation.IsOSPlatform(OSPlatform.Windows)) + { + return "zerobus_ffi.dll"; + } + + if (RuntimeInformation.IsOSPlatform(OSPlatform.OSX)) + { + return "libzerobus_ffi.dylib"; + } + + return "libzerobus_ffi.so"; + } + + // --- SDK lifecycle --- + + [DllImport(LibName, CallingConvention = CallingConvention.Cdecl, EntryPoint = "zerobus_sdk_free")] + public static extern void SdkFree(IntPtr sdk); + + // --- Stream creation --- + + [DllImport(LibName, CallingConvention = CallingConvention.Cdecl, EntryPoint = "zerobus_sdk_create_stream")] + public static extern unsafe IntPtr SdkCreateStream( + IntPtr sdk, + [MarshalAs(UnmanagedType.LPUTF8Str)] string tableName, + byte* descriptorProtoBytes, + nuint descriptorProtoLen, + [MarshalAs(UnmanagedType.LPUTF8Str)] string clientId, + [MarshalAs(UnmanagedType.LPUTF8Str)] string clientSecret, + ref CStreamConfigurationOptions options, + ref CResult result); + + [DllImport(LibName, CallingConvention = CallingConvention.Cdecl, EntryPoint = "zerobus_sdk_create_stream_with_headers_provider")] + public static extern unsafe IntPtr SdkCreateStreamWithHeadersProvider( + IntPtr sdk, + [MarshalAs(UnmanagedType.LPUTF8Str)] string tableName, + byte* descriptorProtoBytes, + nuint descriptorProtoLen, + HeadersProviderCallback headersCallback, + IntPtr userData, + ref CStreamConfigurationOptions options, + ref CResult result); + + [DllImport(LibName, CallingConvention = CallingConvention.Cdecl, EntryPoint = "zerobus_sdk_create_stream_with_headers_provider_async")] + [return: MarshalAs(UnmanagedType.U1)] + public static extern unsafe bool SdkCreateStreamWithHeadersProviderAsync( + IntPtr sdk, + [MarshalAs(UnmanagedType.LPUTF8Str)] string tableName, + byte* descriptorProtoBytes, + nuint descriptorProtoLen, + HeadersProviderCallback headersCallback, + IntPtr userData, + ref CStreamConfigurationOptions options, + CreateStreamAsyncCallback callback, + IntPtr callbackUserData, + ref CResult result); + + [DllImport(LibName, CallingConvention = CallingConvention.Cdecl, EntryPoint = "zerobus_sdk_create_stream_async")] + [return: MarshalAs(UnmanagedType.U1)] + public static extern unsafe bool SdkCreateStreamAsync( + IntPtr sdk, + [MarshalAs(UnmanagedType.LPUTF8Str)] string tableName, + byte* descriptorProtoBytes, + nuint descriptorProtoLen, + [MarshalAs(UnmanagedType.LPUTF8Str)] string clientId, + [MarshalAs(UnmanagedType.LPUTF8Str)] string clientSecret, + ref CStreamConfigurationOptions options, + CreateStreamAsyncCallback callback, + IntPtr userData, + ref CResult result); + + [DllImport(LibName, CallingConvention = CallingConvention.Cdecl, EntryPoint = "zerobus_sdk_recreate_stream")] + public static extern IntPtr SdkRecreateStream( + IntPtr sdk, + IntPtr stream, + ref CResult result); + + [DllImport(LibName, CallingConvention = CallingConvention.Cdecl, EntryPoint = "zerobus_sdk_recreate_stream_async")] + [return: MarshalAs(UnmanagedType.U1)] + public static extern bool SdkRecreateStreamAsync( + IntPtr sdk, + IntPtr stream, + CreateStreamAsyncCallback callback, + IntPtr userData, + ref CResult result); + + [DllImport(LibName, CallingConvention = CallingConvention.Cdecl, EntryPoint = "zerobus_stream_is_closed")] + public static extern bool StreamIsClosed(IntPtr stream); + + [DllImport(LibName, CallingConvention = CallingConvention.Cdecl, EntryPoint = "zerobus_stream_free")] + public static extern void StreamFree(IntPtr stream); + + // --- Record ingestion --- + + [DllImport(LibName, CallingConvention = CallingConvention.Cdecl, EntryPoint = "zerobus_stream_ingest_proto_record")] + public static extern unsafe long StreamIngestProtoRecord( + IntPtr stream, + byte* data, + nuint dataLen, + ref CResult result); + + [DllImport(LibName, CallingConvention = CallingConvention.Cdecl, EntryPoint = "zerobus_stream_ingest_proto_record_async")] + [return: MarshalAs(UnmanagedType.U1)] + public static extern unsafe bool StreamIngestProtoRecordAsync( + IntPtr stream, + byte* data, + nuint dataLen, + OffsetAsyncCallback callback, + IntPtr userData, + ref CResult result); + + [DllImport(LibName, CallingConvention = CallingConvention.Cdecl, EntryPoint = "zerobus_stream_ingest_json_record")] + public static extern long StreamIngestJsonRecord( + IntPtr stream, + [MarshalAs(UnmanagedType.LPUTF8Str)] string jsonData, + ref CResult result); + + [DllImport(LibName, CallingConvention = CallingConvention.Cdecl, EntryPoint = "zerobus_stream_ingest_json_record_async")] + [return: MarshalAs(UnmanagedType.U1)] + public static extern bool StreamIngestJsonRecordAsync( + IntPtr stream, + [MarshalAs(UnmanagedType.LPUTF8Str)] string jsonData, + OffsetAsyncCallback callback, + IntPtr userData, + ref CResult result); + + [DllImport(LibName, CallingConvention = CallingConvention.Cdecl, EntryPoint = "zerobus_stream_ingest_proto_records")] + public static extern unsafe long StreamIngestProtoRecords( + IntPtr stream, + byte** records, + nuint* recordLens, + nuint numRecords, + ref CResult result); + + [DllImport(LibName, CallingConvention = CallingConvention.Cdecl, EntryPoint = "zerobus_stream_ingest_proto_records_async")] + [return: MarshalAs(UnmanagedType.U1)] + public static extern unsafe bool StreamIngestProtoRecordsAsync( + IntPtr stream, + byte** records, + nuint* recordLens, + nuint numRecords, + OffsetAsyncCallback callback, + IntPtr userData, + ref CResult result); + + [DllImport(LibName, CallingConvention = CallingConvention.Cdecl, EntryPoint = "zerobus_stream_ingest_json_records")] + public static extern unsafe long StreamIngestJsonRecords( + IntPtr stream, + byte** jsonRecords, + nuint numRecords, + ref CResult result); + + [DllImport(LibName, CallingConvention = CallingConvention.Cdecl, EntryPoint = "zerobus_stream_ingest_json_records_async")] + [return: MarshalAs(UnmanagedType.U1)] + public static extern unsafe bool StreamIngestJsonRecordsAsync( + IntPtr stream, + byte** jsonRecords, + nuint numRecords, + OffsetAsyncCallback callback, + IntPtr userData, + ref CResult result); + + // --- Acknowledgment / flush --- + + [DllImport(LibName, CallingConvention = CallingConvention.Cdecl, EntryPoint = "zerobus_stream_wait_for_offset")] + [return: MarshalAs(UnmanagedType.U1)] + public static extern bool StreamWaitForOffset( + IntPtr stream, + long offset, + ref CResult result); + + [DllImport(LibName, CallingConvention = CallingConvention.Cdecl, EntryPoint = "zerobus_stream_wait_for_offset_async")] + [return: MarshalAs(UnmanagedType.U1)] + public static extern bool StreamWaitForOffsetAsync( + IntPtr stream, + long offset, + BoolAsyncCallback callback, + IntPtr userData, + ref CResult result); + + [DllImport(LibName, CallingConvention = CallingConvention.Cdecl, EntryPoint = "zerobus_stream_flush")] + [return: MarshalAs(UnmanagedType.U1)] + public static extern bool StreamFlush(IntPtr stream, ref CResult result); + + [DllImport(LibName, CallingConvention = CallingConvention.Cdecl, EntryPoint = "zerobus_stream_flush_async")] + [return: MarshalAs(UnmanagedType.U1)] + public static extern bool StreamFlushAsync( + IntPtr stream, + BoolAsyncCallback callback, + IntPtr userData, + ref CResult result); + + [DllImport(LibName, CallingConvention = CallingConvention.Cdecl, EntryPoint = "zerobus_stream_get_unacked_records")] + public static extern CRecordArray StreamGetUnackedRecords(IntPtr stream, ref CResult result); + + [DllImport(LibName, CallingConvention = CallingConvention.Cdecl, EntryPoint = "zerobus_stream_get_unacked_records_async")] + [return: MarshalAs(UnmanagedType.U1)] + public static extern bool StreamGetUnackedRecordsAsync( + IntPtr stream, + RecordArrayAsyncCallback callback, + IntPtr userData, + ref CResult result); + + [DllImport(LibName, CallingConvention = CallingConvention.Cdecl, EntryPoint = "zerobus_free_record_array")] + public static extern void FreeRecordArray(CRecordArray array); + + // --- Stream close --- + + [DllImport(LibName, CallingConvention = CallingConvention.Cdecl, EntryPoint = "zerobus_stream_close")] + [return: MarshalAs(UnmanagedType.U1)] + public static extern bool StreamClose(IntPtr stream, ref CResult result); + + [DllImport(LibName, CallingConvention = CallingConvention.Cdecl, EntryPoint = "zerobus_stream_close_async")] + [return: MarshalAs(UnmanagedType.U1)] + public static extern bool StreamCloseAsync( + IntPtr stream, + BoolAsyncCallback callback, + IntPtr userData, + ref CResult result); + + // --- Memory management --- + + [DllImport(LibName, CallingConvention = CallingConvention.Cdecl, EntryPoint = "zerobus_free_error_message")] + public static extern void FreeErrorMessage(IntPtr message); + + [DllImport(LibName, CallingConvention = CallingConvention.Cdecl, EntryPoint = "zerobus_alloc_header_array")] + public static extern IntPtr AllocHeaderArray(nuint count); + + [DllImport(LibName, CallingConvention = CallingConvention.Cdecl, EntryPoint = "zerobus_alloc_cstring")] + public static extern unsafe IntPtr AllocCString(byte* data, nuint len); + + [DllImport(LibName, CallingConvention = CallingConvention.Cdecl, EntryPoint = "zerobus_free_headers")] + public static extern void FreeHeaders(CHeaders headers); + + // --- Configuration --- + + [DllImport(LibName, CallingConvention = CallingConvention.Cdecl, EntryPoint = "zerobus_get_default_config")] + public static extern CStreamConfigurationOptions GetDefaultConfig(); + + // --- SDK builder --- + + [DllImport(LibName, CallingConvention = CallingConvention.Cdecl, EntryPoint = "zerobus_sdk_builder_new")] + public static extern IntPtr SdkBuilderNew(); + + [DllImport(LibName, CallingConvention = CallingConvention.Cdecl, EntryPoint = "zerobus_sdk_builder_endpoint")] + public static extern void SdkBuilderEndpoint( + IntPtr builder, + [MarshalAs(UnmanagedType.LPUTF8Str)] string value); + + [DllImport(LibName, CallingConvention = CallingConvention.Cdecl, EntryPoint = "zerobus_sdk_builder_unity_catalog_url")] + public static extern void SdkBuilderUnityCatalogUrl( + IntPtr builder, + [MarshalAs(UnmanagedType.LPUTF8Str)] string value); + + [DllImport(LibName, CallingConvention = CallingConvention.Cdecl, EntryPoint = "zerobus_sdk_builder_sdk_identifier")] + public static extern void SdkBuilderSdkIdentifier( + IntPtr builder, + [MarshalAs(UnmanagedType.LPUTF8Str)] string value); + + [DllImport(LibName, CallingConvention = CallingConvention.Cdecl, EntryPoint = "zerobus_sdk_builder_application_name")] + public static extern void SdkBuilderApplicationName( + IntPtr builder, + [MarshalAs(UnmanagedType.LPUTF8Str)] string value); + + [DllImport(LibName, CallingConvention = CallingConvention.Cdecl, EntryPoint = "zerobus_sdk_builder_disable_tls")] + public static extern void SdkBuilderDisableTls(IntPtr builder); + + [DllImport(LibName, CallingConvention = CallingConvention.Cdecl, EntryPoint = "zerobus_sdk_builder_build")] + public static extern IntPtr SdkBuilderBuild(IntPtr builder, ref CResult result); + + [DllImport(LibName, CallingConvention = CallingConvention.Cdecl, EntryPoint = "zerobus_sdk_builder_free")] + public static extern void SdkBuilderFree(IntPtr builder); +} diff --git a/dotnet/src/Zerobus/Native/NativeInterop.cs b/dotnet/src/Zerobus/Native/NativeInterop.cs new file mode 100644 index 00000000..67ef06e6 --- /dev/null +++ b/dotnet/src/Zerobus/Native/NativeInterop.cs @@ -0,0 +1,1056 @@ +// High-level safe wrappers around P/Invoke calls. +// Handles marshalling, error conversion, and memory management. +// This is the .NET equivalent of the unexported ffi* functions in ffi.go. + +using System.Runtime.InteropServices; +using System.Text; + +namespace Databricks.Zerobus.Native; + +/// +/// Provides safe, managed wrappers around the raw P/Invoke layer. +/// All methods convert errors into . +/// +internal static class NativeInterop +{ + private const int StackAllocThresholdBytes = 4096; + + private static unsafe void ApplyResult(TaskCompletionSource tcs, CResult* result) + { + if (result->Success) + tcs.TrySetResult(); + else + tcs.TrySetException(ToException(result)); + } + + private static unsafe void ApplyResult(TaskCompletionSource tcs, CResult* result, T successValue) + { + if (result->Success) + tcs.TrySetResult(successValue); + else + tcs.TrySetException(ToException(result)); + } + + private static ReadOnlyMemory[] ConvertRecordArrayAndFree(CRecordArray cArray) + { + if ((int)cArray.Len == 0 || cArray.Records == IntPtr.Zero) + return []; + + try + { + var records = new ReadOnlyMemory[(int)cArray.Len]; + var recordSize = Marshal.SizeOf(); + + for (var i = 0; i < (int)cArray.Len; i++) + { + var cRecord = Marshal.PtrToStructure(cArray.Records + i * recordSize); + var data = new byte[(int)cRecord.DataLen]; + Marshal.Copy(cRecord.Data, data, 0, data.Length); + records[i] = data; + } + + return records; + } + finally + { + NativeMethods.FreeRecordArray(cArray); + } + } + + /// + /// Converts a transient pointer (valid only for the duration of + /// a native callback) to a . + /// Unlike , this overload does not free + /// the error message — Rust owns and frees it after the callback returns. + /// + private static unsafe ZerobusException ToException(CResult* result) + { + var message = result->ErrorMessage != IntPtr.Zero + ? Marshal.PtrToStringUTF8(result->ErrorMessage) ?? "unknown error" + : "unknown error"; + return new ZerobusException(message, result->IsRetryable); + } + + /// + /// Converts a to a (or null on success). + /// Frees the native error message string. + /// + internal static ZerobusException? ToException(ref CResult result) + { + if (result.Success) + return null; + + string message; + if (result.ErrorMessage != IntPtr.Zero) + { + message = Marshal.PtrToStringUTF8(result.ErrorMessage) ?? "unknown error"; + NativeMethods.FreeErrorMessage(result.ErrorMessage); + result.ErrorMessage = IntPtr.Zero; + } + else + { + message = "unknown error"; + } + + return new ZerobusException(message, result.IsRetryable); + } + + /// + /// Throws if the indicates failure. + /// + internal static void ThrowIfFailed(ref CResult result) + { + var ex = ToException(ref result); + if (ex is not null) + throw ex; + } + + /// + /// Creates a stream with OAuth credentials. + /// + public static unsafe IntPtr SdkCreateStream( + IntPtr sdkPtr, + string tableName, + ReadOnlySpan descriptorProto, + string clientId, + string clientSecret, + ref CStreamConfigurationOptions options) + { + var result = new CResult(); + IntPtr ptr; + + fixed (byte* descPtr = descriptorProto) + { + ptr = NativeMethods.SdkCreateStream( + sdkPtr, + tableName, + descPtr, + (nuint)descriptorProto.Length, + clientId, + clientSecret, + ref options, + ref result); + } + + if (ptr == IntPtr.Zero) + { + var ex = ToException(ref result); + throw ex ?? new ZerobusException("Failed to create stream", isRetryable: false); + } + + return ptr; + } + /// + /// Creates a stream with OAuth credentials asynchronously. + /// Returns immediately; the returned completes on the Tokio + /// thread when stream creation succeeds or fails. + /// + /// + /// is copied by Rust before this method returns, so + /// the span does not need to remain valid after the call. + /// + public static unsafe Task SdkCreateStreamAsync( + IntPtr sdkPtr, + string tableName, + ReadOnlySpan descriptorProto, + string clientId, + string clientSecret, + ref CStreamConfigurationOptions options) + { + var tcs = new TaskCompletionSource(TaskCreationOptions.RunContinuationsAsynchronously); + + CreateStreamAsyncCallback callbackDelegate = (stream, result, _) => + { + ApplyResult(tcs, (CResult*)result, stream); + }; + + var handle = GCHandle.Alloc(callbackDelegate); + + fixed (byte* descPtr = descriptorProto) + { + var result = new CResult(); + if (!NativeMethods.SdkCreateStreamAsync( + sdkPtr, + tableName, + descPtr, + (nuint)descriptorProto.Length, + clientId, + clientSecret, + ref options, + callbackDelegate, + IntPtr.Zero, + ref result)) + { + var ex = ToException(ref result) + ?? new ZerobusException("Failed to schedule async stream creation", isRetryable: false); + tcs.TrySetException(ex); + } + } + + _ = tcs.Task.ContinueWith( + _ => + { + if (handle.IsAllocated) + handle.Free(); + }, + TaskContinuationOptions.ExecuteSynchronously); + + return tcs.Task; + } + + + /// + /// Creates a stream with a custom headers provider callback. + /// + public static unsafe IntPtr SdkCreateStreamWithHeadersProvider( + IntPtr sdkPtr, + string tableName, + ReadOnlySpan descriptorProto, + HeadersProviderCallback callback, + IntPtr userData, + ref CStreamConfigurationOptions options) + { + var result = new CResult(); + IntPtr ptr; + + fixed (byte* descPtr = descriptorProto) + { + ptr = NativeMethods.SdkCreateStreamWithHeadersProvider( + sdkPtr, + tableName, + descPtr, + (nuint)descriptorProto.Length, + callback, + userData, + ref options, + ref result); + } + + if (ptr == IntPtr.Zero) + { + var ex = ToException(ref result); + throw ex ?? new ZerobusException("Failed to create stream with headers provider", isRetryable: false); + } + + return ptr; + } + + /// + /// Creates a stream with a custom headers provider callback asynchronously. + /// + public static unsafe Task SdkCreateStreamWithHeadersProviderAsync( + IntPtr sdkPtr, + string tableName, + byte[] descriptorProto, + HeadersProviderCallback headersCallback, + IntPtr userData, + CStreamConfigurationOptions options) + { + var tcs = new TaskCompletionSource(TaskCreationOptions.RunContinuationsAsynchronously); + + CreateStreamAsyncCallback completionCallback = (stream, result, _) => + { + ApplyResult(tcs, (CResult*)result, stream); + }; + + var callbackHandle = GCHandle.Alloc(completionCallback); + + fixed (byte* descPtr = descriptorProto) + { + var result = new CResult(); + if (!NativeMethods.SdkCreateStreamWithHeadersProviderAsync( + sdkPtr, + tableName, + descPtr, + (nuint)descriptorProto.Length, + headersCallback, + userData, + ref options, + completionCallback, + IntPtr.Zero, + ref result)) + { + var ex = ToException(ref result) + ?? new ZerobusException("Failed to schedule async stream creation with headers provider", isRetryable: false); + tcs.TrySetException(ex); + } + } + + _ = tcs.Task.ContinueWith( + _ => + { + if (callbackHandle.IsAllocated) + callbackHandle.Free(); + }, + TaskContinuationOptions.ExecuteSynchronously); + + return tcs.Task; + } + + /// + /// Recreates a stream from an existing stream. + /// + public static IntPtr SdkRecreateStream(IntPtr sdkPtr, IntPtr streamPtr) + { + var result = new CResult(); + var ptr = NativeMethods.SdkRecreateStream(sdkPtr, streamPtr, ref result); + + if (ptr == IntPtr.Zero) + { + var ex = ToException(ref result); + throw ex ?? new ZerobusException("Failed to recreate stream", isRetryable: false); + } + + return ptr; + } + + /// + /// Recreates a stream from an existing stream asynchronously. + /// + public static Task SdkRecreateStreamAsync( + IntPtr sdkPtr, + IntPtr streamPtr) + { + var tcs = new TaskCompletionSource(TaskCreationOptions.RunContinuationsAsynchronously); + + CreateStreamAsyncCallback callbackDelegate = (stream, result, _) => + { + unsafe { ApplyResult(tcs, (CResult*)result, stream); } + }; + + var handle = GCHandle.Alloc(callbackDelegate); + var scheduleResult = new CResult(); + if (!NativeMethods.SdkRecreateStreamAsync( + sdkPtr, + streamPtr, + callbackDelegate, + IntPtr.Zero, + ref scheduleResult)) + { + var ex = ToException(ref scheduleResult) + ?? new ZerobusException("Failed to schedule async stream recreation", isRetryable: false); + tcs.TrySetException(ex); + } + + _ = tcs.Task.ContinueWith( + _ => + { + if (handle.IsAllocated) + handle.Free(); + }, + TaskContinuationOptions.ExecuteSynchronously); + + return tcs.Task; + } + + /// + /// Ingests a single protobuf record and returns the offset. + /// + public static unsafe long StreamIngestProtoRecord(IntPtr streamPtr, ReadOnlySpan data) + { + if (data.IsEmpty) + throw new ZerobusException("empty data", isRetryable: false); + + var result = new CResult(); + long offset; + + fixed (byte* dataPtr = data) + { + offset = NativeMethods.StreamIngestProtoRecord( + streamPtr, + dataPtr, + (nuint)data.Length, + ref result); + } + + if (offset < 0) + { + ThrowIfFailed(ref result); + throw new ZerobusException("Ingest failed with unknown error", isRetryable: false); + } + + return offset; + } + + /// + /// Ingests a single protobuf record asynchronously and returns the offset. + /// + public static Task StreamIngestProtoRecordAsync( + IntPtr streamPtr, + byte[] data) + { + if (data.Length == 0) + return Task.FromException(new ZerobusException("empty data", isRetryable: false)); + + var tcs = new TaskCompletionSource(TaskCreationOptions.RunContinuationsAsynchronously); + + OffsetAsyncCallback callback = (offset, result, _) => + { + unsafe { ApplyResult(tcs, (CResult*)result, offset); } + }; + + var handle = GCHandle.Alloc(callback); + + unsafe + { + fixed (byte* dataPtr = data) + { + var scheduleResult = new CResult(); + if (!NativeMethods.StreamIngestProtoRecordAsync( + streamPtr, + dataPtr, + (nuint)data.Length, + callback, + IntPtr.Zero, + ref scheduleResult)) + { + var ex = ToException(ref scheduleResult) + ?? new ZerobusException("Failed to schedule async ingest", isRetryable: false); + tcs.TrySetException(ex); + } + } + } + + _ = tcs.Task.ContinueWith( + _ => + { + if (handle.IsAllocated) + handle.Free(); + }, + TaskContinuationOptions.ExecuteSynchronously); + + return tcs.Task; + } + + /// + /// Ingests a single JSON record and returns the offset. + /// + public static long StreamIngestJsonRecord(IntPtr streamPtr, string jsonData) + { + var result = new CResult(); + var offset = NativeMethods.StreamIngestJsonRecord(streamPtr, jsonData, ref result); + + if (offset < 0) + { + ThrowIfFailed(ref result); + throw new ZerobusException("Ingest failed with unknown error", isRetryable: false); + } + + return offset; + } + + /// + /// Ingests a single JSON record asynchronously and returns the offset. + /// + public static Task StreamIngestJsonRecordAsync( + IntPtr streamPtr, + string jsonData) + { + var tcs = new TaskCompletionSource(TaskCreationOptions.RunContinuationsAsynchronously); + + OffsetAsyncCallback callback = (offset, result, _) => + { + unsafe { ApplyResult(tcs, (CResult*)result, offset); } + }; + + var handle = GCHandle.Alloc(callback); + + var scheduleResult = new CResult(); + if (!NativeMethods.StreamIngestJsonRecordAsync( + streamPtr, + jsonData, + callback, + IntPtr.Zero, + ref scheduleResult)) + { + var ex = ToException(ref scheduleResult) + ?? new ZerobusException("Failed to schedule async ingest", isRetryable: false); + tcs.TrySetException(ex); + } + + _ = tcs.Task.ContinueWith( + _ => + { + if (handle.IsAllocated) + handle.Free(); + }, + TaskContinuationOptions.ExecuteSynchronously); + + return tcs.Task; + } + + /// + /// Ingests a batch of protobuf records and returns the last offset. + /// + public static unsafe long StreamIngestProtoRecords(IntPtr streamPtr, byte[][] records) + { + if (records.Length == 0) + return -1; + + var result = new CResult(); + var numRecords = (nuint)records.Length; + + // Pin all record buffers and collect pointers. + // Use stackalloc for small batches, heap array for large ones. + var handles = new GCHandle[records.Length]; + var ptrs = records.Length * IntPtr.Size <= StackAllocThresholdBytes + ? stackalloc IntPtr[records.Length] + : new IntPtr[records.Length]; + var lens = records.Length * IntPtr.Size <= StackAllocThresholdBytes + ? stackalloc nuint[records.Length] + : new nuint[records.Length]; + + try + { + for (var i = 0; i < records.Length; i++) + { + handles[i] = GCHandle.Alloc(records[i], GCHandleType.Pinned); + ptrs[i] = handles[i].AddrOfPinnedObject(); + lens[i] = (nuint)records[i].Length; + } + + fixed (IntPtr* pointers = ptrs) + fixed (nuint* lengths = lens) + { + var offset = NativeMethods.StreamIngestProtoRecords( + streamPtr, + (byte**)pointers, + lengths, + numRecords, + ref result); + + if (offset == -2) return -1; // empty batch + if (offset < 0) + { + ThrowIfFailed(ref result); + throw new ZerobusException("Batch ingest failed with unknown error", isRetryable: false); + } + + return offset; + } + } + finally + { + for (var i = 0; i < handles.Length; i++) + { + if (handles[i].IsAllocated) + handles[i].Free(); + } + } + } + + /// + /// Ingests a batch of protobuf records asynchronously and returns the last offset. + /// + public static Task StreamIngestProtoRecordsAsync( + IntPtr streamPtr, + byte[][] records) + { + if (records.Length == 0) + return Task.FromResult(-1L); + + var tcs = new TaskCompletionSource(TaskCreationOptions.RunContinuationsAsynchronously); + var numRecords = (nuint)records.Length; + var handles = new GCHandle[records.Length]; + + var callback = new OffsetAsyncCallback((offset, result, _) => + { + unsafe { ApplyResult(tcs, (CResult*)result, offset == -2 ? -1 : offset); } + }); + var callbackHandle = GCHandle.Alloc(callback); + + var ptrs = records.Length * IntPtr.Size <= StackAllocThresholdBytes + ? stackalloc IntPtr[records.Length] + : new IntPtr[records.Length]; + var lens = records.Length * IntPtr.Size <= StackAllocThresholdBytes + ? stackalloc nuint[records.Length] + : new nuint[records.Length]; + + try + { + for (var i = 0; i < records.Length; i++) + { + handles[i] = GCHandle.Alloc(records[i], GCHandleType.Pinned); + ptrs[i] = handles[i].AddrOfPinnedObject(); + lens[i] = (nuint)records[i].Length; + } + + unsafe + { + fixed (IntPtr* pointers = ptrs) + fixed (nuint* lengths = lens) + { + var scheduleResult = new CResult(); + if (!NativeMethods.StreamIngestProtoRecordsAsync( + streamPtr, + (byte**)pointers, + lengths, + numRecords, + callback, + IntPtr.Zero, + ref scheduleResult)) + { + var ex = ToException(ref scheduleResult) + ?? new ZerobusException("Failed to schedule async batch ingest", isRetryable: false); + tcs.TrySetException(ex); + } + } + } + } + finally + { + for (var i = 0; i < handles.Length; i++) + { + if (handles[i].IsAllocated) + handles[i].Free(); + } + } + + _ = tcs.Task.ContinueWith( + _ => + { + if (callbackHandle.IsAllocated) + callbackHandle.Free(); + }, + TaskContinuationOptions.ExecuteSynchronously); + + return tcs.Task; + } + + /// + /// Ingests a batch of JSON records and returns the last offset. + /// + public static unsafe long StreamIngestJsonRecords(IntPtr streamPtr, string[] records) + { + if (records.Length == 0) + return -1; + + var result = new CResult(); + var numRecords = (nuint)records.Length; + + // Encode each string as null-terminated UTF-8 and pin. + // Use stackalloc for small batches, heap array for large ones. + var handles = new GCHandle[records.Length]; + var ptrs = records.Length * IntPtr.Size <= StackAllocThresholdBytes + ? stackalloc IntPtr[records.Length] + : new IntPtr[records.Length]; + + try + { + for (var i = 0; i < records.Length; i++) + { + // Encode with null terminator + var byteCount = Encoding.UTF8.GetByteCount(records[i]); + var utf8 = new byte[byteCount + 1]; + Encoding.UTF8.GetBytes(records[i], utf8); + handles[i] = GCHandle.Alloc(utf8, GCHandleType.Pinned); + ptrs[i] = handles[i].AddrOfPinnedObject(); + } + + fixed (IntPtr* pointers = ptrs) + { + var offset = NativeMethods.StreamIngestJsonRecords( + streamPtr, + (byte**)pointers, + numRecords, + ref result); + + if (offset == -2) return -1; // empty batch + if (offset < 0) + { + ThrowIfFailed(ref result); + throw new ZerobusException("Batch ingest failed with unknown error", isRetryable: false); + } + + return offset; + } + } + finally + { + for (var i = 0; i < handles.Length; i++) + { + if (handles[i].IsAllocated) + handles[i].Free(); + } + } + } + + /// + /// Ingests a batch of JSON records asynchronously and returns the last offset. + /// + public static Task StreamIngestJsonRecordsAsync( + IntPtr streamPtr, + string[] records) + { + if (records.Length == 0) + return Task.FromResult(-1L); + + var tcs = new TaskCompletionSource(TaskCreationOptions.RunContinuationsAsynchronously); + var numRecords = (nuint)records.Length; + var handles = new GCHandle[records.Length]; + + var callback = new OffsetAsyncCallback((offset, result, _) => + { + unsafe { ApplyResult(tcs, (CResult*)result, offset == -2 ? -1 : offset); } + }); + var callbackHandle = GCHandle.Alloc(callback); + + var ptrs = records.Length * IntPtr.Size <= StackAllocThresholdBytes + ? stackalloc IntPtr[records.Length] + : new IntPtr[records.Length]; + + try + { + for (var i = 0; i < records.Length; i++) + { + var byteCount = Encoding.UTF8.GetByteCount(records[i]); + var utf8 = new byte[byteCount + 1]; + Encoding.UTF8.GetBytes(records[i], utf8); + handles[i] = GCHandle.Alloc(utf8, GCHandleType.Pinned); + ptrs[i] = handles[i].AddrOfPinnedObject(); + } + + unsafe + { + fixed (IntPtr* pointers = ptrs) + { + var scheduleResult = new CResult(); + if (!NativeMethods.StreamIngestJsonRecordsAsync( + streamPtr, + (byte**)pointers, + numRecords, + callback, + IntPtr.Zero, + ref scheduleResult)) + { + var ex = ToException(ref scheduleResult) + ?? new ZerobusException("Failed to schedule async batch ingest", isRetryable: false); + tcs.TrySetException(ex); + } + } + } + } + finally + { + for (var i = 0; i < handles.Length; i++) + { + if (handles[i].IsAllocated) + handles[i].Free(); + } + } + + _ = tcs.Task.ContinueWith( + _ => + { + if (callbackHandle.IsAllocated) + callbackHandle.Free(); + }, + TaskContinuationOptions.ExecuteSynchronously); + + return tcs.Task; + } + + /// + /// Waits for a specific offset to be acknowledged. + /// + public static void StreamWaitForOffset(IntPtr streamPtr, long offset) + { + var result = new CResult(); + var success = NativeMethods.StreamWaitForOffset(streamPtr, offset, ref result); + + if (!success) + ThrowIfFailed(ref result); + } + + /// + /// Waits for a specific offset to be acknowledged asynchronously. + /// + public static Task StreamWaitForOffsetAsync( + IntPtr streamPtr, + long offset) + { + var tcs = new TaskCompletionSource(TaskCreationOptions.RunContinuationsAsynchronously); + + BoolAsyncCallback callback = (_, result, _) => + { + unsafe { ApplyResult(tcs, (CResult*)result); } + }; + + var handle = GCHandle.Alloc(callback); + + var scheduleResult = new CResult(); + if (!NativeMethods.StreamWaitForOffsetAsync( + streamPtr, + offset, + callback, + IntPtr.Zero, + ref scheduleResult)) + { + var ex = ToException(ref scheduleResult) + ?? new ZerobusException("Failed to schedule async wait_for_offset", isRetryable: false); + tcs.TrySetException(ex); + } + + _ = tcs.Task.ContinueWith( + _ => + { + if (handle.IsAllocated) + handle.Free(); + }, + TaskContinuationOptions.ExecuteSynchronously); + + return tcs.Task; + } + + /// + /// Flushes all pending records. + /// + public static void StreamFlush(IntPtr streamPtr) + { + var result = new CResult(); + var success = NativeMethods.StreamFlush(streamPtr, ref result); + + if (!success) + ThrowIfFailed(ref result); + } + + /// + /// Flushes all pending records asynchronously. + /// + public static Task StreamFlushAsync(IntPtr streamPtr) + { + var tcs = new TaskCompletionSource(TaskCreationOptions.RunContinuationsAsynchronously); + + BoolAsyncCallback callback = (_, result, _) => + { + unsafe { ApplyResult(tcs, (CResult*)result); } + }; + + var handle = GCHandle.Alloc(callback); + + var scheduleResult = new CResult(); + if (!NativeMethods.StreamFlushAsync( + streamPtr, + callback, + IntPtr.Zero, + ref scheduleResult)) + { + var ex = ToException(ref scheduleResult) + ?? new ZerobusException("Failed to schedule async flush", isRetryable: false); + tcs.TrySetException(ex); + } + + _ = tcs.Task.ContinueWith( + _ => + { + if (handle.IsAllocated) + handle.Free(); + }, + TaskContinuationOptions.ExecuteSynchronously); + + return tcs.Task; + } + + /// + /// Retrieves all unacknowledged records from a closed/failed stream. + /// + public static ReadOnlyMemory[] StreamGetUnackedRecords(IntPtr streamPtr) + { + var result = new CResult(); + var cArray = NativeMethods.StreamGetUnackedRecords(streamPtr, ref result); + + if (cArray.Records == IntPtr.Zero) + { + if ((int)cArray.Len == 0) + { + var ex = ToException(ref result); + if (ex is not null) throw ex; + return []; + } + + ThrowIfFailed(ref result); + return []; + } + + if ((int)cArray.Len == 0) + return []; + + try + { + var records = new ReadOnlyMemory[(int)cArray.Len]; + var recordSize = Marshal.SizeOf(); + + for (var i = 0; i < (int)cArray.Len; i++) + { + var cRecord = Marshal.PtrToStructure(cArray.Records + i * recordSize); + var data = new byte[(int)cRecord.DataLen]; + Marshal.Copy(cRecord.Data, data, 0, data.Length); + + records[i] = data; + } + + return records; + } + finally + { + NativeMethods.FreeRecordArray(cArray); + } + } + + /// + /// Retrieves all unacknowledged records from a closed/failed stream asynchronously. + /// + public static Task[]> StreamGetUnackedRecordsAsync( + IntPtr streamPtr) + { + var tcs = new TaskCompletionSource[]>(TaskCreationOptions.RunContinuationsAsynchronously); + + RecordArrayAsyncCallback callback = (records, result, _) => + { + unsafe + { + var cResult = (CResult*)result; + if (!cResult->Success) + { + tcs.TrySetException(ToException(cResult)); + return; + } + } + + tcs.TrySetResult(ConvertRecordArrayAndFree(records)); + }; + + var handle = GCHandle.Alloc(callback); + + var scheduleResult = new CResult(); + if (!NativeMethods.StreamGetUnackedRecordsAsync( + streamPtr, + callback, + IntPtr.Zero, + ref scheduleResult)) + { + var ex = ToException(ref scheduleResult) + ?? new ZerobusException("Failed to schedule async get_unacked_records", isRetryable: false); + tcs.TrySetException(ex); + } + + _ = tcs.Task.ContinueWith( + _ => + { + if (handle.IsAllocated) + handle.Free(); + }, + TaskContinuationOptions.ExecuteSynchronously); + + return tcs.Task; + } + + /// + /// Closes the stream gracefully. + /// + public static void StreamClose(IntPtr streamPtr) + { + var result = new CResult(); + var success = NativeMethods.StreamClose(streamPtr, ref result); + + if (!success) + ThrowIfFailed(ref result); + } + + /// + /// Closes the stream gracefully asynchronously. + /// + public static Task StreamCloseAsync(IntPtr streamPtr) + { + var tcs = new TaskCompletionSource(TaskCreationOptions.RunContinuationsAsynchronously); + + BoolAsyncCallback callback = (_, result, _) => + { + unsafe { ApplyResult(tcs, (CResult*)result); } + }; + + var handle = GCHandle.Alloc(callback); + + var scheduleResult = new CResult(); + if (!NativeMethods.StreamCloseAsync( + streamPtr, + callback, + IntPtr.Zero, + ref scheduleResult)) + { + var ex = ToException(ref scheduleResult) + ?? new ZerobusException("Failed to schedule async close", isRetryable: false); + tcs.TrySetException(ex); + } + + _ = tcs.Task.ContinueWith( + _ => + { + if (handle.IsAllocated) + handle.Free(); + }, + TaskContinuationOptions.ExecuteSynchronously); + + return tcs.Task; + } + + /// + /// Converts managed to the native struct, + /// applying defaults for unset values. + /// + public static CStreamConfigurationOptions ConvertConfig(StreamConfigurationOptions? options) + { + if (options is null) + return NativeMethods.GetDefaultConfig(); + + var config = NativeMethods.GetDefaultConfig(); + + config.MaxInflightRequests = (nuint)(options.MaxInflightRequests ?? config.MaxInflightRequests); + config.Recovery = options.Recovery; + config.RecoveryTimeoutMs = options.RecoveryTimeoutMs ?? config.RecoveryTimeoutMs; + config.RecoveryBackoffMs = options.RecoveryBackoffMs ?? config.RecoveryBackoffMs; + config.RecoveryRetries = options.RecoveryRetries ?? config.RecoveryRetries; + config.ServerLackOfAckTimeoutMs = options.ServerLackOfAckTimeoutMs ?? config.ServerLackOfAckTimeoutMs; + config.FlushTimeoutMs = options.FlushTimeoutMs ?? config.FlushTimeoutMs; + if (options.RecordType != RecordType.Unspecified) + config.RecordType = (int)options.RecordType; + + if (options.StreamPausedMaxWaitTimeMs.HasValue) + { + config.StreamPausedMaxWaitTimeMs = options.StreamPausedMaxWaitTimeMs.Value; + config.HasStreamPausedMaxWaitTimeMs = true; + } + + return config; + } + + /// + /// Allocates a new native SDK builder. Must be terminated with + /// or . + /// + public static IntPtr SdkBuilderNew() + { + var ptr = NativeMethods.SdkBuilderNew(); + if (ptr == IntPtr.Zero) + throw new ZerobusException("Failed to allocate SDK builder", isRetryable: false); + return ptr; + } + + /// + /// Consumes the builder and returns a native SDK pointer. + /// The builder pointer must not be used after this call. + /// + public static IntPtr SdkBuilderBuild(IntPtr builder) + { + var result = new CResult(); + var ptr = NativeMethods.SdkBuilderBuild(builder, ref result); + + if (ptr == IntPtr.Zero) + { + var ex = ToException(ref result); + throw ex ?? new ZerobusException("Failed to build SDK", isRetryable: false); + } + + return ptr; + } +} diff --git a/dotnet/src/Zerobus/Properties/AssemblyInfo.cs b/dotnet/src/Zerobus/Properties/AssemblyInfo.cs new file mode 100644 index 00000000..335af740 --- /dev/null +++ b/dotnet/src/Zerobus/Properties/AssemblyInfo.cs @@ -0,0 +1,3 @@ +using System.Runtime.CompilerServices; + +[assembly: InternalsVisibleTo("Databricks.Zerobus.Tests")] diff --git a/dotnet/src/Zerobus/ProtoZerobusStream.cs b/dotnet/src/Zerobus/ProtoZerobusStream.cs new file mode 100644 index 00000000..68626457 --- /dev/null +++ b/dotnet/src/Zerobus/ProtoZerobusStream.cs @@ -0,0 +1,76 @@ +namespace Databricks.Zerobus; + +/// +/// A stream that accepts protobuf payloads only. +/// +public sealed class ProtoZerobusStream : TypedZerobusStream +{ + internal ProtoZerobusStream(ZerobusStream innerStream) + : base(innerStream) + { + } + + /// + public long IngestRecord(byte[] payload) + { + return InnerStream.IngestRecord(payload); + } + + /// + public Task IngestRecordAsync(byte[] payload) + { + return InnerStream.IngestRecordAsync(payload); + } + + /// + public long IngestRecord(ReadOnlySpan payload) + { + return InnerStream.IngestRecord(payload); + } + + /// + public long IngestRecords(byte[][] records) + { + return InnerStream.IngestRecords(records); + } + + /// + public Task IngestRecordsAsync(byte[][] records) + { + return InnerStream.IngestRecordsAsync(records); + } + + /// + /// Retrieves all records that have not yet been acknowledged by the server. + /// + /// Important: This should only be called after the stream has + /// closed or failed. Calling it on an active stream will return an error. + /// + /// + /// + /// An array of raw protobuf record payloads as of . + /// + /// Thrown if retrieval fails. + /// Thrown if the stream has been disposed. + public ReadOnlyMemory[] GetUnackedRecords() + { + return InnerStream.GetUnackedRecords(); + } + + /// + /// Asynchronously retrieves all records that have not yet been acknowledged by the server. + /// + /// Important: This should only be called after the stream has + /// closed or failed. Calling it on an active stream will return an error. + /// + /// + /// + /// A task that resolves to an array of raw protobuf record payloads as of . + /// + /// Thrown if retrieval fails. + /// Thrown if the stream has been disposed. + public Task[]> GetUnackedRecordsAsync() + { + return InnerStream.GetUnackedRecordsAsync(); + } +} \ No newline at end of file diff --git a/dotnet/src/Zerobus/RecordType.cs b/dotnet/src/Zerobus/RecordType.cs new file mode 100644 index 00000000..99e64992 --- /dev/null +++ b/dotnet/src/Zerobus/RecordType.cs @@ -0,0 +1,22 @@ +namespace Databricks.Zerobus; + +/// +/// Represents the type of records to ingest. +/// +public enum RecordType +{ + /// + /// No specific record type. + /// + Unspecified = 0, + + /// + /// Protocol Buffer encoded records. + /// + Proto = 1, + + /// + /// JSON encoded records. + /// + Json = 2, +} diff --git a/dotnet/src/Zerobus/StreamConfigurationOptions.cs b/dotnet/src/Zerobus/StreamConfigurationOptions.cs new file mode 100644 index 00000000..92199479 --- /dev/null +++ b/dotnet/src/Zerobus/StreamConfigurationOptions.cs @@ -0,0 +1,100 @@ +namespace Databricks.Zerobus; + +/// +/// Configuration options for creating a Zerobus stream. +/// Start with and override only the fields you need. +/// +/// +/// +/// var options = StreamConfigurationOptions.Default with +/// { +/// MaxInflightRequests = 50_000, +/// RecoveryRetries = 8, +/// }; +/// +/// +public sealed record StreamConfigurationOptions +{ + /// + /// Maximum number of requests that can be in-flight (pending acknowledgment) at once. + /// null uses the SDK default. + /// Default: 1,000,000. + /// + public ulong? MaxInflightRequests { get; init; } + + /// + /// Enable automatic stream recovery on retryable failures. + /// Default: true. + /// + public bool Recovery { get; init; } = true; + + /// + /// Timeout for each recovery attempt in milliseconds. + /// null uses the SDK default. + /// Default: 15,000 (15 seconds). + /// + public ulong? RecoveryTimeoutMs { get; init; } + + /// + /// Backoff delay between recovery attempts in milliseconds. + /// null uses the SDK default. + /// Default: 2,000 (2 seconds). + /// + public ulong? RecoveryBackoffMs { get; init; } + + /// + /// Maximum number of recovery retry attempts. + /// null uses the SDK default. + /// Default: 4. + /// + public uint? RecoveryRetries { get; init; } + + /// + /// Server acknowledgment timeout in milliseconds. + /// null uses the SDK default. + /// Default: 60,000 (60 seconds). + /// + public ulong? ServerLackOfAckTimeoutMs { get; init; } + + /// + /// Flush operation timeout in milliseconds. + /// null uses the SDK default. + /// Default: 300,000 (5 minutes). + /// + public ulong? FlushTimeoutMs { get; init; } + + /// + /// Type of record to ingest (Proto, Json, or Unspecified). + /// Typed factories such as + /// and + /// set this automatically. + /// Default: . + /// + public RecordType RecordType { get; init; } = RecordType.Proto; + + /// + /// Maximum time in milliseconds to wait during graceful stream close + /// when the server sends a CloseStreamSignal. + /// + /// null — Wait for the full server-specified duration (most graceful, default). + /// 0 — Immediate recovery, close stream right away. + /// A positive value — Wait up to min(value, serverDuration) milliseconds. + /// + /// Default: null. + /// + public ulong? StreamPausedMaxWaitTimeMs { get; init; } + + /// + /// Returns the default configuration options. + /// This is the idiomatic starting point — use C# record with expressions to override. + /// + public static StreamConfigurationOptions Default => new() + { + MaxInflightRequests = 1_000_000, + RecoveryTimeoutMs = 15_000, + RecoveryBackoffMs = 2_000, + RecoveryRetries = 4, + ServerLackOfAckTimeoutMs = 60_000, + FlushTimeoutMs = 300_000, + }; +} diff --git a/dotnet/src/Zerobus/TableProperties.cs b/dotnet/src/Zerobus/TableProperties.cs new file mode 100644 index 00000000..ae611f1b --- /dev/null +++ b/dotnet/src/Zerobus/TableProperties.cs @@ -0,0 +1,15 @@ +namespace Databricks.Zerobus; + +/// +/// Contains information about the target Databricks Delta table. +/// +/// +/// Fully qualified table name in the form catalog.schema.table. +/// +/// +/// Protocol buffer descriptor (required for streams, null for JSON). +/// This should be a serialized DescriptorProto. +/// +public sealed record TableProperties( + string TableName, + byte[]? DescriptorProto = null); diff --git a/dotnet/src/Zerobus/TypedZerobusStream.cs b/dotnet/src/Zerobus/TypedZerobusStream.cs new file mode 100644 index 00000000..fa081f07 --- /dev/null +++ b/dotnet/src/Zerobus/TypedZerobusStream.cs @@ -0,0 +1,70 @@ +namespace Databricks.Zerobus; + +/// +/// Common lifecycle and acknowledgment APIs shared by type-safe stream wrappers. +/// +public abstract class TypedZerobusStream : IDisposable, IAsyncDisposable +{ + private readonly ZerobusStream _innerStream; + + internal TypedZerobusStream(ZerobusStream innerStream) + { + _innerStream = innerStream ?? throw new ArgumentNullException(nameof(innerStream)); + } + + internal ZerobusStream InnerStream => _innerStream; + + /// + public bool IsClosed() + { + return _innerStream.IsClosed(); + } + + /// + public void WaitForOffset(long offset) + { + _innerStream.WaitForOffset(offset); + } + + /// + public Task WaitForOffsetAsync(long offset) + { + return _innerStream.WaitForOffsetAsync(offset); + } + + /// + public void Flush() + { + _innerStream.Flush(); + } + + /// + public Task FlushAsync() + { + return _innerStream.FlushAsync(); + } + + /// + public void Close() + { + _innerStream.Close(); + } + + /// + public Task CloseAsync() + { + return _innerStream.CloseAsync(); + } + + /// + public void Dispose() + { + _innerStream.Dispose(); + } + + /// + public ValueTask DisposeAsync() + { + return _innerStream.DisposeAsync(); + } +} \ No newline at end of file diff --git a/dotnet/src/Zerobus/Zerobus.csproj b/dotnet/src/Zerobus/Zerobus.csproj new file mode 100644 index 00000000..9b9c8536 --- /dev/null +++ b/dotnet/src/Zerobus/Zerobus.csproj @@ -0,0 +1,113 @@ + + + + net8.0;net10.0 + Databricks.Zerobus + Databricks.Zerobus + Databricks.Zerobus + 0.1.0 + Databricks + Databricks + High-performance .NET SDK for streaming data ingestion into Databricks Delta tables + using the Zerobus service. + databricks;zerobus;streaming;ingestion;delta + LICENSE + https://github.com/databricks/zerobus-sdk + git + true + true + README.md + + + false + + + true + true + + embedded + true + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + \ No newline at end of file diff --git a/dotnet/src/Zerobus/ZerobusException.cs b/dotnet/src/Zerobus/ZerobusException.cs new file mode 100644 index 00000000..fa0498e0 --- /dev/null +++ b/dotnet/src/Zerobus/ZerobusException.cs @@ -0,0 +1,42 @@ +namespace Databricks.Zerobus; + +/// +/// Represents an error from the Zerobus SDK. +/// Errors are categorised as retryable (transient) or non-retryable (fatal). +/// +/// +/// The SDK automatically recovers from retryable errors when +/// is enabled. +/// Non-retryable errors require manual intervention (e.g. fixing permissions or schema). +/// +public sealed class ZerobusException : Exception +{ + /// + /// Initialises a new . + /// + /// The error message from the native layer. + /// Whether this error is transient and may succeed on retry. + public ZerobusException(string message, bool isRetryable) + : base(FormatMessage(message, isRetryable)) + { + IsRetryable = isRetryable; + RawMessage = message; + } + + /// + /// Whether this error is transient and the operation may succeed if retried. + /// When is enabled the SDK + /// handles retryable errors automatically. + /// + public bool IsRetryable { get; } + + /// + /// The raw error message from the native layer, without the retryable prefix. + /// + public string RawMessage { get; } + + private static string FormatMessage(string message, bool isRetryable) + => isRetryable + ? $"ZerobusException (retryable): {message}" + : $"ZerobusException: {message}"; +} diff --git a/dotnet/src/Zerobus/ZerobusSdk.cs b/dotnet/src/Zerobus/ZerobusSdk.cs new file mode 100644 index 00000000..712430a0 --- /dev/null +++ b/dotnet/src/Zerobus/ZerobusSdk.cs @@ -0,0 +1,641 @@ +using System.Runtime.InteropServices; +using Databricks.Zerobus.Native; + +namespace Databricks.Zerobus; + +/// +/// The main entry point for interacting with the Zerobus ingestion service. +/// Manages the connection to the Zerobus endpoint and Unity Catalog. +/// +/// +/// +/// This class wraps a native Rust SDK via P/Invoke and manages the lifecycle +/// of the underlying unmanaged resource. Always dispose when finished. +/// +/// +/// The SDK is thread-safe — you may create multiple streams from one instance. +/// +/// +/// +/// +/// using var sdk = ZerobusSdk.CreateBuilder() +/// .Endpoint("https://your-shard.zerobus.databricks.com") +/// .UnityCatalogUrl("https://your-workspace.databricks.com") +/// .Build(); +/// +/// using var stream = sdk.CreateJsonStream( +/// "catalog.schema.table", +/// clientId, +/// clientSecret); +/// +/// +public sealed class ZerobusSdk : IDisposable +{ + private IntPtr _ptr; + private int _disposed; + + /// + /// Internal constructor used by . + /// Takes ownership of the native pointer. + /// + internal ZerobusSdk(IntPtr ptr) + { + _ptr = ptr; + } + + /// + /// Returns a new builder for constructing a with + /// optional settings such as application name, TLS override, or a Unity + /// Catalog URL that can be omitted when using a custom headers provider. + /// + /// A ready to configure and build. + /// + /// + /// using var sdk = ZerobusSdk.CreateBuilder() + /// .Endpoint("https://zerobus.databricks.com") + /// .UnityCatalogUrl("https://workspace.databricks.com") + /// .ApplicationName("my-service") + /// .Build(); + /// + /// + public static ZerobusSdkBuilder CreateBuilder() => new(); + + /// + /// Creates a new bidirectional gRPC stream for ingesting records into a Databricks table. + /// Uses OAuth 2.0 client credentials flow for authentication. + /// + /// Table properties including name and optional protobuf descriptor. + /// OAuth 2.0 client ID. + /// OAuth 2.0 client secret. + /// + /// Stream configuration options. Pass null or omit to use defaults. + /// + /// A new ready for record ingestion. + /// + /// Thrown if the stream cannot be created (auth failure, invalid table, etc.). + /// + /// Thrown if the SDK has been disposed. + /// + /// + /// using var stream = sdk.CreateStream( + /// new TableProperties("catalog.schema.table"), + /// clientId, + /// clientSecret); + /// + /// + public ZerobusStream CreateStream( + TableProperties tableProperties, + string clientId, + string clientSecret, + StreamConfigurationOptions? options = null) + { + ObjectDisposedException.ThrowIf(_disposed != 0, this); + ArgumentNullException.ThrowIfNull(tableProperties); + ArgumentNullException.ThrowIfNull(clientId); + ArgumentNullException.ThrowIfNull(clientSecret); + + return CreateStreamCore(tableProperties, clientId, clientSecret, options); + } + + /// + /// Creates a new bidirectional gRPC stream asynchronously. + /// + public Task CreateStreamAsync( + TableProperties tableProperties, + string clientId, + string clientSecret, + StreamConfigurationOptions? options = null) + { + ObjectDisposedException.ThrowIf(_disposed != 0, this); + ArgumentNullException.ThrowIfNull(tableProperties); + ArgumentNullException.ThrowIfNull(clientId); + ArgumentNullException.ThrowIfNull(clientSecret); + + return CreateStreamCoreAsync(tableProperties, clientId, clientSecret, options); + } + + /// + /// Creates a JSON-only stream with OAuth 2.0 client credentials authentication. + /// This factory sets to + /// automatically and returns a stream wrapper that + /// exposes JSON ingest overloads only. + /// + /// Fully qualified table name in the form catalog.schema.table. + /// OAuth 2.0 client ID. + /// OAuth 2.0 client secret. + /// Optional stream configuration overrides. + /// A new ready for JSON ingestion. + public JsonZerobusStream CreateJsonStream( + string tableName, + string clientId, + string clientSecret, + StreamConfigurationOptions? options = null) + { + ArgumentNullException.ThrowIfNull(tableName); + ArgumentNullException.ThrowIfNull(clientId); + ArgumentNullException.ThrowIfNull(clientSecret); + + var stream = CreateStreamCore( + new TableProperties(tableName), + clientId, + clientSecret, + NormalizeStreamOptions(options, RecordType.Json)); + + return new JsonZerobusStream(stream); + } + + /// + /// Creates a JSON-only stream asynchronously. + /// + public async Task CreateJsonStreamAsync( + string tableName, + string clientId, + string clientSecret, + StreamConfigurationOptions? options = null) + { + ArgumentNullException.ThrowIfNull(tableName); + ArgumentNullException.ThrowIfNull(clientId); + ArgumentNullException.ThrowIfNull(clientSecret); + + var stream = await CreateStreamCoreAsync( + new TableProperties(tableName), + clientId, + clientSecret, + NormalizeStreamOptions(options, RecordType.Json)) + .ConfigureAwait(false); + + return new JsonZerobusStream(stream); + } + + /// + /// Creates a protobuf-only stream with OAuth 2.0 client credentials authentication. + /// This factory sets to + /// automatically and returns a stream wrapper that + /// exposes protobuf ingest overloads only. + /// + /// Fully qualified table name in the form catalog.schema.table. + /// Serialized protobuf descriptor for the target table schema. + /// OAuth 2.0 client ID. + /// OAuth 2.0 client secret. + /// Optional stream configuration overrides. + /// A new ready for protobuf ingestion. + public ProtoZerobusStream CreateProtoStream( + string tableName, + byte[] descriptorProto, + string clientId, + string clientSecret, + StreamConfigurationOptions? options = null) + { + ArgumentNullException.ThrowIfNull(tableName); + ArgumentNullException.ThrowIfNull(descriptorProto); + ArgumentNullException.ThrowIfNull(clientId); + ArgumentNullException.ThrowIfNull(clientSecret); + + var stream = CreateStreamCore( + new TableProperties(tableName, descriptorProto), + clientId, + clientSecret, + NormalizeStreamOptions(options, RecordType.Proto)); + + return new ProtoZerobusStream(stream); + } + + /// + /// Creates a protobuf-only stream asynchronously. + /// + public async Task CreateProtoStreamAsync( + string tableName, + byte[] descriptorProto, + string clientId, + string clientSecret, + StreamConfigurationOptions? options = null) + { + ArgumentNullException.ThrowIfNull(tableName); + ArgumentNullException.ThrowIfNull(descriptorProto); + ArgumentNullException.ThrowIfNull(clientId); + ArgumentNullException.ThrowIfNull(clientSecret); + + var stream = await CreateStreamCoreAsync( + new TableProperties(tableName, descriptorProto), + clientId, + clientSecret, + NormalizeStreamOptions(options, RecordType.Proto)) + .ConfigureAwait(false); + + return new ProtoZerobusStream(stream); + } + + private ZerobusStream CreateStreamCore( + TableProperties tableProperties, + string clientId, + string clientSecret, + StreamConfigurationOptions? options) + { + ValidateStreamConfiguration(tableProperties, options); + + var nativeOpts = NativeInterop.ConvertConfig(options); + + var streamPtr = NativeInterop.SdkCreateStream( + _ptr, + tableProperties.TableName, + tableProperties.DescriptorProto ?? [], + clientId, + clientSecret, + ref nativeOpts); + + return new ZerobusStream(streamPtr); + } + + private async Task CreateStreamCoreAsync( + TableProperties tableProperties, + string clientId, + string clientSecret, + StreamConfigurationOptions? options) + { + ValidateStreamConfiguration(tableProperties, options); + + var nativeOpts = NativeInterop.ConvertConfig(options); + + var streamPtr = await NativeInterop.SdkCreateStreamAsync( + _ptr, + tableProperties.TableName, + tableProperties.DescriptorProto ?? [], + clientId, + clientSecret, + ref nativeOpts) + .ConfigureAwait(false); + + return new ZerobusStream(streamPtr); + } + + /// + /// Creates a new bidirectional gRPC stream using a custom headers provider. + /// Use this for custom authentication logic (managed identity, vaults, etc.). + /// + /// Table properties including name and optional protobuf descriptor. + /// Custom implementation of . + /// + /// Stream configuration options. Pass null or omit to use defaults. + /// + /// A new ready for record ingestion. + /// + /// Thrown if the stream cannot be created (headers provider error, network issues, etc.). + /// + /// Thrown if the SDK has been disposed. + /// + /// + /// var provider = new CustomHeadersProvider(); + /// using var stream = sdk.CreateStreamWithHeadersProvider( + /// new TableProperties("catalog.schema.table"), + /// provider); + /// + /// + public ZerobusStream CreateStreamWithHeadersProvider( + TableProperties tableProperties, + IHeadersProvider headersProvider, + StreamConfigurationOptions? options = null) + { + ObjectDisposedException.ThrowIf(_disposed != 0, this); + ArgumentNullException.ThrowIfNull(tableProperties); + ArgumentNullException.ThrowIfNull(headersProvider); + + return CreateStreamWithHeadersProviderCore(tableProperties, headersProvider, options); + } + + /// + /// Creates a new bidirectional gRPC stream using a custom headers provider asynchronously. + /// + public Task CreateStreamWithHeadersProviderAsync( + TableProperties tableProperties, + IHeadersProvider headersProvider, + StreamConfigurationOptions? options = null) + { + ObjectDisposedException.ThrowIf(_disposed != 0, this); + ArgumentNullException.ThrowIfNull(tableProperties); + ArgumentNullException.ThrowIfNull(headersProvider); + + return CreateStreamWithHeadersProviderCoreAsync(tableProperties, headersProvider, options); + } + + /// + /// Creates a JSON-only stream using a custom headers provider. + /// + /// Fully qualified table name in the form catalog.schema.table. + /// Custom authentication headers provider. + /// Optional stream configuration overrides. + /// A new ready for JSON ingestion. + public JsonZerobusStream CreateJsonStreamWithHeadersProvider( + string tableName, + IHeadersProvider headersProvider, + StreamConfigurationOptions? options = null) + { + ArgumentNullException.ThrowIfNull(tableName); + ArgumentNullException.ThrowIfNull(headersProvider); + + var stream = CreateStreamWithHeadersProviderCore( + new TableProperties(tableName), + headersProvider, + NormalizeStreamOptions(options, RecordType.Json)); + + return new JsonZerobusStream(stream); + } + + /// + /// Creates a JSON-only stream using a custom headers provider asynchronously. + /// + public async Task CreateJsonStreamWithHeadersProviderAsync( + string tableName, + IHeadersProvider headersProvider, + StreamConfigurationOptions? options = null) + { + ArgumentNullException.ThrowIfNull(tableName); + ArgumentNullException.ThrowIfNull(headersProvider); + + var stream = await CreateStreamWithHeadersProviderAsync( + new TableProperties(tableName), + headersProvider, + NormalizeStreamOptions(options, RecordType.Json)) + .ConfigureAwait(false); + + return new JsonZerobusStream(stream); + } + + /// + /// Creates a protobuf-only stream using a custom headers provider. + /// + /// Fully qualified table name in the form catalog.schema.table. + /// Serialized protobuf descriptor for the target table schema. + /// Custom authentication headers provider. + /// Optional stream configuration overrides. + /// A new ready for protobuf ingestion. + public ProtoZerobusStream CreateProtoStreamWithHeadersProvider( + string tableName, + byte[] descriptorProto, + IHeadersProvider headersProvider, + StreamConfigurationOptions? options = null) + { + ArgumentNullException.ThrowIfNull(tableName); + ArgumentNullException.ThrowIfNull(descriptorProto); + ArgumentNullException.ThrowIfNull(headersProvider); + + var stream = CreateStreamWithHeadersProviderCore( + new TableProperties(tableName, descriptorProto), + headersProvider, + NormalizeStreamOptions(options, RecordType.Proto)); + + return new ProtoZerobusStream(stream); + } + + /// + /// Creates a protobuf-only stream using a custom headers provider asynchronously. + /// + public async Task CreateProtoStreamWithHeadersProviderAsync( + string tableName, + byte[] descriptorProto, + IHeadersProvider headersProvider, + StreamConfigurationOptions? options = null) + { + ArgumentNullException.ThrowIfNull(tableName); + ArgumentNullException.ThrowIfNull(descriptorProto); + ArgumentNullException.ThrowIfNull(headersProvider); + + var stream = await CreateStreamWithHeadersProviderAsync( + new TableProperties(tableName, descriptorProto), + headersProvider, + NormalizeStreamOptions(options, RecordType.Proto)) + .ConfigureAwait(false); + + return new ProtoZerobusStream(stream); + } + + private ZerobusStream CreateStreamWithHeadersProviderCore( + TableProperties tableProperties, + IHeadersProvider headersProvider, + StreamConfigurationOptions? options) + { + ValidateStreamConfiguration(tableProperties, options); + + var nativeOpts = NativeInterop.ConvertConfig(options); + + // Create the callback bridge that the native code will invoke. + var bridge = new HeadersProviderBridge(headersProvider); + var callback = new HeadersProviderCallback(bridge.NativeCallback); + + // GCHandle keeps the bridge + callback alive for the lifetime of the stream. + var handle = GCHandle.Alloc(bridge); + + IntPtr streamPtr; + try + { + streamPtr = NativeInterop.SdkCreateStreamWithHeadersProvider( + _ptr, + tableProperties.TableName, + tableProperties.DescriptorProto ?? [], + callback, + GCHandle.ToIntPtr(handle), + ref nativeOpts); + } + catch + { + handle.Free(); + throw; + } + + return new ZerobusStream(streamPtr, handle, callback); + } + + private async Task CreateStreamWithHeadersProviderCoreAsync( + TableProperties tableProperties, + IHeadersProvider headersProvider, + StreamConfigurationOptions? options) + { + ValidateStreamConfiguration(tableProperties, options); + + var nativeOpts = NativeInterop.ConvertConfig(options); + + var bridge = new HeadersProviderBridge(headersProvider); + var callback = new HeadersProviderCallback(bridge.NativeCallback); + + var handle = GCHandle.Alloc(bridge); + + IntPtr streamPtr; + try + { + streamPtr = await NativeInterop.SdkCreateStreamWithHeadersProviderAsync( + _ptr, + tableProperties.TableName, + tableProperties.DescriptorProto ?? [], + callback, + GCHandle.ToIntPtr(handle), + nativeOpts) + .ConfigureAwait(false); + } + catch + { + if (handle.IsAllocated) + handle.Free(); + throw; + } + + return new ZerobusStream(streamPtr, handle, callback); + } + + /// + /// Recreates a new bidirectional gRPC stream from an existing stream. + /// This is used for recovery scenarios where a stream needs to be re-established + /// using the existing stream's configuration and state. + /// + /// The existing stream to recreate from. + /// + /// + /// This method transfers ownership from to the returned stream. + /// The input stream is disposed as part of recreation and cannot be used afterward. + /// + /// + /// A later call on the original stream wrapper is a no-op, + /// so using declarations remain safe: + /// + /// using var stream = sdk.CreateStream(...); + /// stream.Close(); + /// using var recreated = sdk.RecreateStream(stream); + /// + /// + /// + /// A new ready for record ingestion. + /// Thrown if the stream cannot be recreated. + /// Thrown if the SDK has been disposed. + public ZerobusStream RecreateStream(ZerobusStream stream) + { + ObjectDisposedException.ThrowIf(_disposed != 0, this); + ArgumentNullException.ThrowIfNull(stream); + + var ptr = NativeInterop.SdkRecreateStream(_ptr, stream.NativePointer); + return stream.Recreate(ptr); + } + + /// + /// Recreates a new bidirectional gRPC stream asynchronously from an existing stream. + /// + public async Task RecreateStreamAsync( + ZerobusStream stream) + { + ObjectDisposedException.ThrowIf(_disposed != 0, this); + ArgumentNullException.ThrowIfNull(stream); + + var ptr = await NativeInterop.SdkRecreateStreamAsync(_ptr, stream.NativePointer) + .ConfigureAwait(false); + return stream.Recreate(ptr); + } + + /// + /// Recreates a JSON-only stream after the original stream has failed or closed. + /// + public JsonZerobusStream RecreateStream(JsonZerobusStream stream) + { + ObjectDisposedException.ThrowIf(_disposed != 0, this); + ArgumentNullException.ThrowIfNull(stream); + + return new JsonZerobusStream(RecreateStream(stream.InnerStream)); + } + + /// + /// Recreates a JSON-only stream asynchronously after the original stream has failed or closed. + /// + public async Task RecreateStreamAsync( + JsonZerobusStream stream) + { + ObjectDisposedException.ThrowIf(_disposed != 0, this); + ArgumentNullException.ThrowIfNull(stream); + + var recreated = await RecreateStreamAsync(stream.InnerStream).ConfigureAwait(false); + return new JsonZerobusStream(recreated); + } + + /// + /// Recreates a protobuf-only stream after the original stream has failed or closed. + /// + public ProtoZerobusStream RecreateStream(ProtoZerobusStream stream) + { + ObjectDisposedException.ThrowIf(_disposed != 0, this); + ArgumentNullException.ThrowIfNull(stream); + + return new ProtoZerobusStream(RecreateStream(stream.InnerStream)); + } + + /// + /// Recreates a protobuf-only stream asynchronously after the original stream has failed or closed. + /// + public async Task RecreateStreamAsync(ProtoZerobusStream stream) + { + ObjectDisposedException.ThrowIf(_disposed != 0, this); + ArgumentNullException.ThrowIfNull(stream); + + var recreated = await RecreateStreamAsync(stream.InnerStream).ConfigureAwait(false); + return new ProtoZerobusStream(recreated); + } + + private static StreamConfigurationOptions NormalizeStreamOptions( + StreamConfigurationOptions? options, + RecordType recordType) + { + return (options ?? StreamConfigurationOptions.Default) with { RecordType = recordType }; + } + + private static void ValidateStreamConfiguration( + TableProperties tableProperties, + StreamConfigurationOptions? options) + { + ArgumentNullException.ThrowIfNull(tableProperties); + + var descriptor = tableProperties.DescriptorProto; + var effectiveRecordType = options?.RecordType switch + { + null => StreamConfigurationOptions.Default.RecordType, + RecordType.Unspecified => StreamConfigurationOptions.Default.RecordType, + var recordType => recordType, + }; + + if (descriptor is { Length: 0 }) + { + throw new ArgumentException( + "DescriptorProto must be null for JSON streams or a non-empty serialized DescriptorProto for protobuf streams.", + nameof(tableProperties)); + } + + var hasDescriptor = descriptor is { Length: > 0 }; + + if (effectiveRecordType == RecordType.Json && hasDescriptor) + { + throw new ArgumentException( + "JSON streams cannot specify DescriptorProto. Use TableProperties(tableName) or CreateJsonStream(...).", + nameof(tableProperties)); + } + + if (effectiveRecordType == RecordType.Proto && !hasDescriptor) + { + throw new ArgumentException( + "Proto streams require a non-empty DescriptorProto. Use TableProperties(tableName, descriptorProto) or CreateProtoStream(...).", + nameof(tableProperties)); + } + } + + private void Free() + { + var ptr = Interlocked.Exchange(ref _ptr, IntPtr.Zero); + if (ptr != IntPtr.Zero) + { + NativeMethods.SdkFree(ptr); + } + } + + /// + public void Dispose() + { + if (Interlocked.CompareExchange(ref _disposed, 1, 0) != 0) return; + Free(); + GC.SuppressFinalize(this); + } + + /// Safety-net release of native memory for leaked instances. + ~ZerobusSdk() => Free(); +} diff --git a/dotnet/src/Zerobus/ZerobusSdkBuilder.cs b/dotnet/src/Zerobus/ZerobusSdkBuilder.cs new file mode 100644 index 00000000..2862734d --- /dev/null +++ b/dotnet/src/Zerobus/ZerobusSdkBuilder.cs @@ -0,0 +1,175 @@ +using Databricks.Zerobus.Native; + +namespace Databricks.Zerobus; + +/// +/// A builder for creating instances with optional configuration. +/// +/// +/// +/// Obtain an instance via and call the +/// fluent setter methods before calling . +/// +/// +/// The builder is single-use: calling or +/// consumes it. Any further method calls after either will throw +/// . +/// +/// +/// +/// OAuth credentials: +/// +/// using var sdk = ZerobusSdk.CreateBuilder() +/// .Endpoint("https://zerobus.databricks.com") +/// .UnityCatalogUrl("https://workspace.databricks.com") +/// .ApplicationName("my-service") +/// .Build(); +/// +/// +/// Custom headers provider (Unity Catalog URL not required): +/// +/// using var sdk = ZerobusSdk.CreateBuilder() +/// .Endpoint("https://zerobus.databricks.com") +/// .ApplicationName("my-service") +/// .Build(); +/// +/// +public sealed class ZerobusSdkBuilder : IDisposable +{ + private string? _endpoint; + private string? _unityCatalogUrl; + private string? _sdkIdentifier; + private string? _applicationName; + private bool _disableTls; + private int _consumed; // 0 = live, 1 = consumed/disposed + + internal ZerobusSdkBuilder() { } + + /// + /// Sets the Zerobus gRPC endpoint URL (required). + /// + /// The gRPC endpoint (e.g. https://zerobus.databricks.com). + /// This builder, for chaining. + /// Thrown if is null. + /// Thrown if the builder has been consumed or disposed. + public ZerobusSdkBuilder Endpoint(string endpoint) + { + ThrowIfConsumed(); + ArgumentNullException.ThrowIfNull(endpoint); + _endpoint = endpoint; + return this; + } + + /// + /// Sets the Unity Catalog URL used for OAuth token acquisition. + /// Required when using OAuth credentials; optional with a custom headers provider. + /// + /// The Unity Catalog URL (e.g. https://workspace.databricks.com). + /// This builder, for chaining. + /// Thrown if is null. + /// Thrown if the builder has been consumed or disposed. + public ZerobusSdkBuilder UnityCatalogUrl(string unityCatalogUrl) + { + ThrowIfConsumed(); + ArgumentNullException.ThrowIfNull(unityCatalogUrl); + _unityCatalogUrl = unityCatalogUrl; + return this; + } + + /// + /// Overrides the SDK identifier portion of the User-Agent header. + /// Wrapper SDKs use this to identify themselves; end-user code should prefer + /// instead. + /// + /// The SDK identifier string. + /// This builder, for chaining. + /// Thrown if is null. + /// Thrown if the builder has been consumed or disposed. + public ZerobusSdkBuilder SdkIdentifier(string sdkIdentifier) + { + ThrowIfConsumed(); + ArgumentNullException.ThrowIfNull(sdkIdentifier); + _sdkIdentifier = sdkIdentifier; + return this; + } + + /// + /// Appends an application name to the User-Agent header. + /// Useful for identifying which application is sending data in server-side logs. + /// + /// The application name. + /// This builder, for chaining. + /// Thrown if is null. + /// Thrown if the builder has been consumed or disposed. + public ZerobusSdkBuilder ApplicationName(string applicationName) + { + ThrowIfConsumed(); + ArgumentNullException.ThrowIfNull(applicationName); + _applicationName = applicationName; + return this; + } + + /// + /// Disables TLS for the gRPC connection. TLS is enabled by default. + /// Only use this for local development or testing. + /// + /// This builder, for chaining. + /// Thrown if the builder has been consumed or disposed. + public ZerobusSdkBuilder DisableTls() + { + ThrowIfConsumed(); + _disableTls = true; + return this; + } + + /// + /// Builds and returns a instance. + /// Consumes the builder — any further calls will throw . + /// + /// A fully initialised . + /// Thrown if the SDK cannot be initialised. + /// Thrown if the builder has already been consumed or disposed. + public ZerobusSdk Build() + { + ThrowIfConsumed(); + + if (Interlocked.CompareExchange(ref _consumed, 1, 0) != 0) + throw new ObjectDisposedException(nameof(ZerobusSdkBuilder)); + + var builderPtr = NativeInterop.SdkBuilderNew(); + + try + { + if (_endpoint is not null) + NativeMethods.SdkBuilderEndpoint(builderPtr, _endpoint); + if (_unityCatalogUrl is not null) + NativeMethods.SdkBuilderUnityCatalogUrl(builderPtr, _unityCatalogUrl); + if (_sdkIdentifier is not null) + NativeMethods.SdkBuilderSdkIdentifier(builderPtr, _sdkIdentifier); + if (_applicationName is not null) + NativeMethods.SdkBuilderApplicationName(builderPtr, _applicationName); + if (_disableTls) + NativeMethods.SdkBuilderDisableTls(builderPtr); + } + catch + { + NativeMethods.SdkBuilderFree(builderPtr); + throw; + } + + // SdkBuilderBuild consumes the native pointer on both success and failure. + var sdkPtr = NativeInterop.SdkBuilderBuild(builderPtr); + return new ZerobusSdk(sdkPtr); + } + + private void ThrowIfConsumed() + { + ObjectDisposedException.ThrowIf(_consumed != 0, this); + } + + /// + public void Dispose() + { + Interlocked.Exchange(ref _consumed, 1); + } +} diff --git a/dotnet/src/Zerobus/ZerobusStream.cs b/dotnet/src/Zerobus/ZerobusStream.cs new file mode 100644 index 00000000..db7c38f3 --- /dev/null +++ b/dotnet/src/Zerobus/ZerobusStream.cs @@ -0,0 +1,597 @@ +using System.Runtime.CompilerServices; +using System.Runtime.ExceptionServices; +using System.Runtime.InteropServices; +using Databricks.Zerobus.Native; + +namespace Databricks.Zerobus; + +/// +/// Represents an active bidirectional gRPC stream for ingesting records. +/// Records can be ingested concurrently and will be acknowledged asynchronously. +/// +/// +/// +/// The stream is thread-safe — you may call from +/// multiple threads concurrently, just like the Go SDK supports goroutines. +/// +/// +/// performs a graceful close (flush + close) but keeps the +/// native stream alive so callers can inspect or +/// recover with ZerobusSdk.RecreateStream(...). Call +/// when you are completely finished with the stream to free native resources. +/// +/// +/// ZerobusSdk.RecreateStream(stream) consumes stream and returns a +/// replacement stream wrapper. The original wrapper is disposed during recreation. +/// performs the same graceful close as +/// without blocking the calling thread, enabling await using for streams. +/// +/// +public sealed class ZerobusStream : IDisposable, IAsyncDisposable +{ + private IntPtr _ptr; + private int _disposed; + private readonly ReaderWriterLockSlim _lifetimeLock = new(); + private int _inflightAsyncOperations; + private readonly ManualResetEventSlim _asyncOperationsDrained = new(initialState: true); + + // Prevent the GCHandle / delegate from being collected while the native code holds a reference. + // Not readonly: GCHandle is not a readonly struct, so calling Free() on a readonly field + // creates a defensive copy — the field is never actually mutated, causing double-free. + private GCHandle _bridgeHandle; + private readonly HeadersProviderCallback? _callbackRef; + + internal ZerobusStream(IntPtr ptr) + { + _ptr = ptr; + } + + internal ZerobusStream(IntPtr ptr, GCHandle bridgeHandle, HeadersProviderCallback callbackRef) + { + _ptr = ptr; + _bridgeHandle = bridgeHandle; + _callbackRef = callbackRef; + } + + /// + /// Returns whether the stream has been closed + /// + /// + public bool IsClosed() + { + return WithReadLock(NativeMethods.StreamIsClosed); + } + + // ── Single-record ingestion ────────────────────────────────────────── + + /// + /// Ingests a single record and returns the offset. + /// This is the primary API for record ingestion. + /// + /// + /// The record payload. Pass a for JSON records + /// or a byte[] / of + /// for Protocol Buffer records. + /// + /// The offset of the ingested record. + /// Thrown if ingestion fails. + /// Thrown if the stream has been disposed. + /// + /// Thrown if the payload type is not string or byte[]. + /// + /// + /// + /// // JSON + /// long offset = stream.IngestRecord("{\"id\": 1, \"message\": \"Hello\"}"); + /// + /// // Protobuf + /// byte[] protoBytes = SerializeMyProto(myMessage); + /// long offset = stream.IngestRecord(protoBytes); + /// + /// + public long IngestRecord(string payload) + { + ArgumentNullException.ThrowIfNull(payload); + return WithReadLock(ptr => NativeInterop.StreamIngestJsonRecord(ptr, payload)); + } + + /// + public Task IngestRecordAsync(string payload) + { + ArgumentNullException.ThrowIfNull(payload); + return WithReadLockAsync(ptr => NativeInterop.StreamIngestJsonRecordAsync(ptr, payload)); + } + + /// + public long IngestRecord(byte[] payload) + { + ArgumentNullException.ThrowIfNull(payload); + return WithReadLock(ptr => NativeInterop.StreamIngestProtoRecord(ptr, payload)); + } + + /// + public Task IngestRecordAsync(byte[] payload) + { + ArgumentNullException.ThrowIfNull(payload); + return WithReadLockAsync(ptr => NativeInterop.StreamIngestProtoRecordAsync(ptr, payload)); + } + + /// + public long IngestRecord(ReadOnlySpan payload) + { + using var handle = WithReadLock(); + return NativeInterop.StreamIngestProtoRecord(GetNativePointerForCall(), payload); + } + + // ── Batch ingestion ────────────────────────────────────────────────── + + /// + /// Ingests a batch of JSON records and returns one offset for the entire batch. + /// All records in the batch must be JSON strings. + /// + /// The JSON record strings to ingest. + /// The offset representing the entire batch, or -1 if the batch is empty. + /// Thrown if ingestion fails. + /// Thrown if the stream has been disposed. + /// + /// + /// string[] records = + /// [ + /// "{\"device\": \"sensor-001\", \"temp\": 20}", + /// "{\"device\": \"sensor-002\", \"temp\": 21}", + /// ]; + /// long batchOffset = stream.IngestRecords(records); + /// + /// + public long IngestRecords(string[] records) + { + ArgumentNullException.ThrowIfNull(records); + return WithReadLock(ptr => NativeInterop.StreamIngestJsonRecords(ptr, records)); + } + + /// + public Task IngestRecordsAsync(string[] records) + { + ArgumentNullException.ThrowIfNull(records); + return WithReadLockAsync(ptr => NativeInterop.StreamIngestJsonRecordsAsync(ptr, records)); + } + + /// + /// Ingests a batch of protobuf records and returns one offset for the entire batch. + /// All records in the batch must be serialised protobuf byte arrays. + /// + /// The protobuf record byte spans to ingest. + /// The offset representing the entire batch, or -1 if the batch is empty. + /// Thrown if ingestion fails. + /// Thrown if the stream has been disposed. + public long IngestRecords(byte[][] records) + { + ArgumentNullException.ThrowIfNull(records); + return WithReadLock(ptr => NativeInterop.StreamIngestProtoRecords(ptr, records)); + } + + /// + public Task IngestRecordsAsync(byte[][] records) + { + ArgumentNullException.ThrowIfNull(records); + return WithReadLockAsync(ptr => NativeInterop.StreamIngestProtoRecordsAsync(ptr, records)); + } + + // ── Acknowledgment / flush ─────────────────────────────────────────── + + /// + /// Blocks until the server acknowledges the record at the specified offset. + /// Use this with offsets returned from to wait for + /// specific records to be durably written without waiting for all pending records. + /// + /// The offset to wait for. + /// Thrown if the wait fails. + /// Thrown if the stream has been disposed. + /// + /// + /// long offset = stream.IngestRecord(data); + /// // ... do other work ... + /// stream.WaitForOffset(offset); + /// + /// + public void WaitForOffset(long offset) + { + WithReadLock(ptr => NativeInterop.StreamWaitForOffset(ptr, offset)); + } + + /// + public Task WaitForOffsetAsync(long offset) + { + return WithReadLockAsync(ptr => NativeInterop.StreamWaitForOffsetAsync(ptr, offset)); + } + + /// + /// Blocks until all pending records have been acknowledged by the server. + /// This ensures durability guarantees before proceeding. + /// + /// + /// Thrown if the flush times out or a record fails with a non-retryable error. + /// + /// Thrown if the stream has been disposed. + /// + /// + /// stream.Flush(); + /// Console.WriteLine("All records durably stored."); + /// + /// + public void Flush() + { + WithReadLock(NativeInterop.StreamFlush); + } + + /// + public Task FlushAsync() + { + return WithReadLockAsync(NativeInterop.StreamFlushAsync); + } + + // ── Unacknowledged records ─────────────────────────────────────────── + + /// + /// Retrieves all records that have not yet been acknowledged by the server. + /// + /// Important: This should only be called after the stream has + /// closed or failed. Calling it on an active stream will return an error. + /// + /// + /// + /// An array of raw record payloads as of . + /// JSON payloads are UTF-8 encoded; callers can decode as needed. + /// + /// Thrown if retrieval fails. + /// Thrown if the stream has been disposed. + /// + /// + /// try + /// { + /// stream.Flush(); + /// } + /// catch (ZerobusException) + /// { + /// var unacked = stream.GetUnackedRecords(); + /// Console.WriteLine($"Failed to acknowledge {unacked.Length} records"); + /// foreach (var payload in unacked) + /// Console.WriteLine($"{payload.Length} bytes"); + /// } + /// + /// + public ReadOnlyMemory[] GetUnackedRecords() + { + return WithReadLock(NativeInterop.StreamGetUnackedRecords); + } + + /// + public Task[]> GetUnackedRecordsAsync() + { + return WithReadLockAsync(NativeInterop.StreamGetUnackedRecordsAsync); + } + + // ── Close / Dispose ────────────────────────────────────────────────── + + /// + /// Gracefully closes the stream after flushing all pending records. + /// After close, ingestion APIs are no longer usable, but the stream remains + /// readable for recovery paths (for example ). + /// + /// + /// This method does not free native resources. Call + /// after you have finished any recovery operations. + /// + /// + /// Thrown if flush or close fails. + /// + public void Close() + { + using var handle = WithWriteLock(); + ObjectDisposedException.ThrowIf(_disposed != 0, this); + var ptr = GetNativePointerForCall(); + if (NativeMethods.StreamIsClosed(ptr)) return; + + NativeInterop.StreamClose(ptr); + } + + /// + public Task CloseAsync() + { + return WithWriteLockAsync( + ptr => + { + if (NativeMethods.StreamIsClosed(ptr)) + return Task.CompletedTask; + return NativeInterop.StreamCloseAsync(ptr); + }); + } + + /// + public void Dispose() + { + Dispose(true); + GC.SuppressFinalize(this); + } + + /// + public async ValueTask DisposeAsync() + { + if (Interlocked.CompareExchange(ref _disposed, 1, 0) != 0) + return; + + IntPtr ptr = IntPtr.Zero; + var shouldClose = false; + Exception? closeError = null; + + using (WithWriteLock()) + { + ptr = Interlocked.Exchange(ref _ptr, IntPtr.Zero); + if (ptr == IntPtr.Zero) + { + FreeBridgeHandle(); + return; + } + + shouldClose = !NativeMethods.StreamIsClosed(ptr); + } + + if (shouldClose) + { + try + { + await NativeInterop.StreamCloseAsync(ptr).ConfigureAwait(false); + } + catch (Exception ex) + { + closeError = ex; + } + } + + NativeMethods.StreamFree(ptr); + FreeBridgeHandle(); + GC.SuppressFinalize(this); + + if (closeError is not null) + ExceptionDispatchInfo.Capture(closeError).Throw(); + } + + private void Dispose(bool disposing) + { + if (Interlocked.CompareExchange(ref _disposed, 1, 0) != 0) return; + + using var handle = WithWriteLock(); + var ptr = Interlocked.Exchange(ref _ptr, IntPtr.Zero); + if (ptr == IntPtr.Zero) + { + FreeBridgeHandle(); + return; + } + + Exception? closeError = null; + + if (disposing && !NativeMethods.StreamIsClosed(ptr)) + { + try + { + NativeInterop.StreamClose(ptr); + } + catch (Exception ex) + { + closeError = ex; + } + } + + NativeMethods.StreamFree(ptr); + FreeBridgeHandle(); + + if (closeError is not null) + ExceptionDispatchInfo.Capture(closeError).Throw(); + } + + /// + /// Safety-net release of native memory for leaked instances. + /// Only frees memory — does NOT call , + /// which performs blocking gRPC I/O and may trigger managed callbacks on + /// Rust-created threads that corrupt the .NET execution context + /// (infinite recursion in CultureInfo.CurrentUICulture). + /// + ~ZerobusStream() + { + Dispose(false); + } + + private void FreeBridgeHandle() + { + if (_bridgeHandle.IsAllocated) + _bridgeHandle.Free(); + } + + /// + /// Returns the raw native pointer to the underlying C stream. + /// Internal use only. + /// + internal IntPtr NativePointer + { + get + { + return WithReadLock(ptr => ptr); + } + } + + /// + /// Attempts to retrieve the bridge handle and callback reference for the stream. + /// Internal use only. + /// + internal ZerobusStream Recreate(IntPtr newPtr) + { + var disposed = Interlocked.CompareExchange(ref _disposed, 1, 0); + ObjectDisposedException.ThrowIf(disposed != 0, this); + + using var handle = WithWriteLock(); + + var oldPtr = Interlocked.Exchange(ref _ptr, IntPtr.Zero); + var newStream = _bridgeHandle.IsAllocated + ? new ZerobusStream(newPtr, _bridgeHandle, _callbackRef!) + : new ZerobusStream(newPtr); + + // Prevent accidental double-free of the shared handle from the old wrapper + _bridgeHandle = default; + + // Free the old native stream (don't Close — it's already failed/closed) + if (oldPtr != IntPtr.Zero) + NativeMethods.StreamFree(oldPtr); + + return newStream; + } + + private IntPtr GetNativePointerForCall() + { + ObjectDisposedException.ThrowIf(_disposed != 0, this); + var ptr = _ptr; + ObjectDisposedException.ThrowIf(ptr == IntPtr.Zero, this); + return ptr; + } + + [MethodImpl(MethodImplOptions.AggressiveInlining)] + private T WithReadLock(Func call) + { + using var handle = WithReadLock(); + return call(GetNativePointerForCall()); + } + + [MethodImpl(MethodImplOptions.AggressiveInlining)] + private void WithReadLock(Action call) + { + using var handle = WithReadLock(); + call(GetNativePointerForCall()); + } + + [MethodImpl(MethodImplOptions.AggressiveInlining)] + private Task WithReadLockAsync(Func> call) + { + return WithAsyncLock(writeLock: false, call); + } + + [MethodImpl(MethodImplOptions.AggressiveInlining)] + private Task WithReadLockAsync(Func call) + { + return WithAsyncLock(writeLock: false, call); + } + + [MethodImpl(MethodImplOptions.AggressiveInlining)] + private Task WithWriteLockAsync(Func call) + { + return WithAsyncLock(writeLock: true, call); + } + + [MethodImpl(MethodImplOptions.AggressiveInlining)] + private IDisposable WithReadLock() + { + _lifetimeLock.EnterReadLock(); + return new Disposable(() => _lifetimeLock.ExitReadLock()); + } + + [MethodImpl(MethodImplOptions.AggressiveInlining)] + private IDisposable WithWriteLock() + { + _lifetimeLock.EnterWriteLock(); + _asyncOperationsDrained.Wait(); + return new Disposable(() => _lifetimeLock.ExitWriteLock()); + } + + [MethodImpl(MethodImplOptions.AggressiveInlining)] + private Task WithAsyncLock(bool writeLock, Func> call) + { + var ptr = BeginAsyncOperation(writeLock); + try + { + var task = call(ptr); + return CompleteAsyncOperation(task); + } + catch + { + EndAsyncOperation(); + throw; + } + } + + [MethodImpl(MethodImplOptions.AggressiveInlining)] + private Task WithAsyncLock(bool writeLock, Func call) + { + var ptr = BeginAsyncOperation(writeLock); + try + { + var task = call(ptr); + return CompleteAsyncOperation(task); + } + catch + { + EndAsyncOperation(); + throw; + } + } + + [MethodImpl(MethodImplOptions.AggressiveInlining)] + private IntPtr BeginAsyncOperation(bool writeLock) + { + if (writeLock) + _lifetimeLock.EnterWriteLock(); + else + _lifetimeLock.EnterReadLock(); + + try + { + var ptr = GetNativePointerForCall(); + if (Interlocked.Increment(ref _inflightAsyncOperations) == 1) + _asyncOperationsDrained.Reset(); + return ptr; + } + finally + { + if (writeLock) + _lifetimeLock.ExitWriteLock(); + else + _lifetimeLock.ExitReadLock(); + } + } + + [MethodImpl(MethodImplOptions.AggressiveInlining)] + private void EndAsyncOperation() + { + if (Interlocked.Decrement(ref _inflightAsyncOperations) == 0) + _asyncOperationsDrained.Set(); + } + + [MethodImpl(MethodImplOptions.AggressiveInlining)] + private async Task CompleteAsyncOperation(Task task) + { + try + { + return await task.ConfigureAwait(false); + } + finally + { + EndAsyncOperation(); + } + } + + [MethodImpl(MethodImplOptions.AggressiveInlining)] + private async Task CompleteAsyncOperation(Task task) + { + try + { + await task.ConfigureAwait(false); + } + finally + { + EndAsyncOperation(); + } + } + + private class Disposable(Action action) : IDisposable + { + public void Dispose() => action(); + } + +} diff --git a/dotnet/tests/Zerobus.IntegrationTests/HeadersProviderIntegrationTests.cs b/dotnet/tests/Zerobus.IntegrationTests/HeadersProviderIntegrationTests.cs new file mode 100644 index 00000000..9d30e4a3 --- /dev/null +++ b/dotnet/tests/Zerobus.IntegrationTests/HeadersProviderIntegrationTests.cs @@ -0,0 +1,78 @@ +using NUnit.Framework; + +namespace Databricks.Zerobus.IntegrationTests; + +[TestFixture] +[Parallelizable(ParallelScope.Children)] +public class HeadersProviderIntegrationTests : IntegrationTestBase +{ + [Test] + public async Task CreateStreamWithHeadersProvider_InvokesCallbackAndSendsHeaders() + { + await using var fixture = await MockServerFixture.StartAsync(); + + fixture.MockServer.InjectResponses(TestTableName, + [ + MockResponses.CreateStreamResponse("test_stream_headers_provider"), + ]); + + var provider = new CountingHeadersProvider(new Dictionary + { + ["authorization"] = "Bearer callback_token", + ["x-databricks-zerobus-table-name"] = TestTableName, + ["x-test-callback-header"] = "headers-bridge-path", + }); + + using var sdk = CreateDefaultSdk(fixture); + + var tableProps = CreateTableProperties(); + + using var stream = sdk.CreateStreamWithHeadersProvider( + tableProps, + provider, + StreamConfigurationOptions.Default with { Recovery = false }); + + Assert.That(provider.CallCount, Is.GreaterThan(0)); + + var observedHeaders = fixture.MockServer.GetLastRequestHeaders(TestTableName); + Assert.That(observedHeaders.TryGetValue("x-test-callback-header", out var callbackHeader), Is.True); + Assert.That(callbackHeader, Is.EqualTo("headers-bridge-path")); + Assert.That(observedHeaders.TryGetValue("authorization", out var authHeader), Is.True); + Assert.That(authHeader, Is.EqualTo("Bearer callback_token")); + } + + [Test] + public async Task CreateStreamWithHeadersProviderAsync_InvokesCallbackAndSendsHeaders() + { + await using var fixture = await MockServerFixture.StartAsync(); + + fixture.MockServer.InjectResponses(TestTableName, + [ + MockResponses.CreateStreamResponse("test_stream_headers_provider_async"), + ]); + + var provider = new CountingHeadersProvider(new Dictionary + { + ["authorization"] = "Bearer callback_token", + ["x-databricks-zerobus-table-name"] = TestTableName, + ["x-test-callback-header"] = "headers-bridge-path", + }); + + using var sdk = CreateDefaultSdk(fixture); + + var tableProps = CreateTableProperties(); + + await using var stream = await sdk.CreateStreamWithHeadersProviderAsync( + tableProps, + provider, + StreamConfigurationOptions.Default with { Recovery = false }); + + Assert.That(provider.CallCount, Is.GreaterThan(0)); + + var observedHeaders = fixture.MockServer.GetLastRequestHeaders(TestTableName); + Assert.That(observedHeaders.TryGetValue("x-test-callback-header", out var callbackHeader), Is.True); + Assert.That(callbackHeader, Is.EqualTo("headers-bridge-path")); + Assert.That(observedHeaders.TryGetValue("authorization", out var authHeader), Is.True); + Assert.That(authHeader, Is.EqualTo("Bearer callback_token")); + } +} diff --git a/dotnet/tests/Zerobus.IntegrationTests/IngestionIntegrationTests.cs b/dotnet/tests/Zerobus.IntegrationTests/IngestionIntegrationTests.cs new file mode 100644 index 00000000..97739f82 --- /dev/null +++ b/dotnet/tests/Zerobus.IntegrationTests/IngestionIntegrationTests.cs @@ -0,0 +1,386 @@ +using NUnit.Framework; + +namespace Databricks.Zerobus.IntegrationTests; + +[TestFixture] +[Parallelizable(ParallelScope.Children)] +public class IngestionIntegrationTests : IntegrationTestBase +{ + [Test] + public async Task IngestSingleRecord() + { + await using var fixture = await MockServerFixture.StartAsync(); + + fixture.MockServer.InjectResponses(TestTableName, + [ + MockResponses.CreateStreamResponse("test_stream_1"), + MockResponses.RecordAckResponse(0), + ]); + + using var sdk = CreateDefaultSdk(fixture); + var tableProps = CreateTableProperties(); + var options = CreateDefaultOptions(); + + using var stream = sdk.CreateStreamWithHeadersProvider(tableProps, new TestHeadersProvider(), options); + + var testRecord = "test record data"u8.ToArray(); + var offsetId = stream.IngestRecord(testRecord); + + Assert.That(offsetId, Is.EqualTo(0)); + + await Task.Delay(100); + + var writeCount = fixture.MockServer.GetWriteCount(); + var maxOffset = fixture.MockServer.GetMaxOffsetSent(); + + Assert.That(writeCount, Is.EqualTo(1)); + Assert.That(maxOffset, Is.EqualTo(0)); + } + + [Test] + public async Task IngestSingleRecord_AsyncApi() + { + await using var fixture = await MockServerFixture.StartAsync(); + + fixture.MockServer.InjectResponses(TestTableName, + [ + MockResponses.CreateStreamResponse("test_stream_async_1"), + MockResponses.RecordAckResponse(0), + ]); + + using var sdk = CreateDefaultSdk(fixture); + var tableProps = CreateTableProperties(); + var options = CreateDefaultOptions(); + + await using var stream = await sdk.CreateStreamWithHeadersProviderAsync(tableProps, new TestHeadersProvider(), options); + + var testRecord = "test record data"u8.ToArray(); + var offsetId = await stream.IngestRecordAsync(testRecord); + + Assert.That(offsetId, Is.EqualTo(0)); + + await stream.WaitForOffsetAsync(offsetId); + + var writeCount = fixture.MockServer.GetWriteCount(); + var maxOffset = fixture.MockServer.GetMaxOffsetSent(); + + Assert.That(writeCount, Is.EqualTo(1)); + Assert.That(maxOffset, Is.EqualTo(0)); + } + + [Test] + public async Task IngestMultipleRecords() + { + await using var fixture = await MockServerFixture.StartAsync(); + + fixture.MockServer.InjectResponses(TestTableName, + [ + MockResponses.CreateStreamResponse("test_stream_1"), + MockResponses.RecordAckResponse(0), + MockResponses.RecordAckResponse(1), + MockResponses.RecordAckResponse(2), + ]); + + using var sdk = CreateDefaultSdk(fixture); + var tableProps = CreateTableProperties(); + var options = CreateDefaultOptions(); + + using var stream = sdk.CreateStreamWithHeadersProvider(tableProps, new TestHeadersProvider(), options); + + const int numRecords = 3; + for (var i = 0; i < numRecords; i++) + { + var testRecord = System.Text.Encoding.UTF8.GetBytes($"test record {i}"); + var offsetId = stream.IngestRecord(testRecord); + + Assert.That(offsetId, Is.EqualTo(i)); + } + + await Task.Delay(100); + + var writeCount = fixture.MockServer.GetWriteCount(); + var maxOffset = fixture.MockServer.GetMaxOffsetSent(); + + Assert.That(writeCount, Is.EqualTo((ulong)numRecords)); + Assert.That(maxOffset, Is.EqualTo(numRecords - 1)); + } + + [Test] + public async Task IngestMultipleRecords_AsyncApi() + { + await using var fixture = await MockServerFixture.StartAsync(); + + fixture.MockServer.InjectResponses(TestTableName, + [ + MockResponses.CreateStreamResponse("test_stream_async_2"), + MockResponses.RecordAckResponse(0), + MockResponses.RecordAckResponse(1), + MockResponses.RecordAckResponse(2), + ]); + + using var sdk = CreateDefaultSdk(fixture); + var tableProps = CreateTableProperties(); + var options = CreateDefaultOptions(); + + await using var stream = await sdk.CreateStreamWithHeadersProviderAsync(tableProps, new TestHeadersProvider(), options); + + const int numRecords = 3; + for (var i = 0; i < numRecords; i++) + { + var testRecord = System.Text.Encoding.UTF8.GetBytes($"test record {i}"); + var offsetId = await stream.IngestRecordAsync(testRecord); + + Assert.That(offsetId, Is.EqualTo(i)); + } + + await stream.WaitForOffsetAsync(numRecords - 1); + + var writeCount = fixture.MockServer.GetWriteCount(); + var maxOffset = fixture.MockServer.GetMaxOffsetSent(); + + Assert.That(writeCount, Is.EqualTo((ulong)numRecords)); + Assert.That(maxOffset, Is.EqualTo(numRecords - 1)); + } + + [Test] + public async Task IngestSingleStreamConcurrently() + { + await using var fixture = await MockServerFixture.StartAsync(); + + const int numRecords = 32; + + var responses = new List + { + MockResponses.CreateStreamResponse("test_stream_concurrent_1"), + }; + + for (var i = 0; i < numRecords; i++) + { + responses.Add(MockResponses.RecordAckResponse(i)); + } + + fixture.MockServer.InjectResponses(TestTableName, responses); + + using var sdk = CreateDefaultSdk(fixture); + var tableProps = CreateTableProperties(); + var options = CreateDefaultOptions(); + + using var stream = sdk.CreateStreamWithHeadersProvider(tableProps, new TestHeadersProvider(), options); + + var offsets = new long[numRecords]; + var tasks = new Task[numRecords]; + + for (var i = 0; i < numRecords; i++) + { + var index = i; + tasks[index] = Task.Run(() => + { + var payload = System.Text.Encoding.UTF8.GetBytes($"concurrent test record {index}"); + offsets[index] = stream.IngestRecord(payload); + }); + } + + await Task.WhenAll(tasks); + + var sortedOffsets = (long[])offsets.Clone(); + Array.Sort(sortedOffsets); + + for (var i = 0; i < numRecords; i++) + { + Assert.That(sortedOffsets[i], Is.EqualTo(i)); + } + + stream.WaitForOffset(numRecords - 1); + + var writeCount = fixture.MockServer.GetWriteCount(); + var maxOffset = fixture.MockServer.GetMaxOffsetSent(); + + Assert.That(writeCount, Is.EqualTo((ulong)numRecords)); + Assert.That(maxOffset, Is.EqualTo(numRecords - 1)); + } + + [Test] + public async Task IngestSingleStreamConcurrently_AsyncApi() + { + await using var fixture = await MockServerFixture.StartAsync(); + + const int numRecords = 32; + + var responses = new List + { + MockResponses.CreateStreamResponse("test_stream_concurrent_async_1"), + }; + + for (var i = 0; i < numRecords; i++) + { + responses.Add(MockResponses.RecordAckResponse(i)); + } + + fixture.MockServer.InjectResponses(TestTableName, responses); + + using var sdk = CreateDefaultSdk(fixture); + var tableProps = CreateTableProperties(); + var options = CreateDefaultOptions(); + + await using var stream = await sdk.CreateStreamWithHeadersProviderAsync(tableProps, new TestHeadersProvider(), options); + + var offsets = new long[numRecords]; + var tasks = new Task[numRecords]; + + for (var i = 0; i < numRecords; i++) + { + var index = i; + tasks[index] = Task.Run(async () => + { + var payload = System.Text.Encoding.UTF8.GetBytes($"concurrent test record {index}"); + offsets[index] = await stream.IngestRecordAsync(payload); + }); + } + + await Task.WhenAll(tasks); + + var sortedOffsets = (long[])offsets.Clone(); + Array.Sort(sortedOffsets); + + for (var i = 0; i < numRecords; i++) + { + Assert.That(sortedOffsets[i], Is.EqualTo(i)); + } + + await stream.WaitForOffsetAsync(numRecords - 1); + + var writeCount = fixture.MockServer.GetWriteCount(); + var maxOffset = fixture.MockServer.GetMaxOffsetSent(); + + Assert.That(writeCount, Is.EqualTo((ulong)numRecords)); + Assert.That(maxOffset, Is.EqualTo(numRecords - 1)); + } + + [Test] + public async Task IngestBatchRecords() + { + await using var fixture = await MockServerFixture.StartAsync(); + + fixture.MockServer.InjectResponses(TestTableName, + [ + MockResponses.CreateStreamResponse("test_stream_1"), + MockResponses.RecordAckResponse(0), + ]); + + using var sdk = CreateDefaultSdk(fixture); + var tableProps = CreateTableProperties(); + var options = CreateDefaultOptions(); + + using var stream = sdk.CreateStreamWithHeadersProvider(tableProps, new TestHeadersProvider(), options); + + byte[][] batch = + [ + "record 1"u8.ToArray(), + "record 2"u8.ToArray(), + "record 3"u8.ToArray(), + "record 4"u8.ToArray(), + "record 5"u8.ToArray(), + ]; + + var offsetId = stream.IngestRecords(batch); + + Assert.That(offsetId, Is.EqualTo(0)); + + await Task.Delay(100); + + var writeCount = fixture.MockServer.GetWriteCount(); + var maxOffset = fixture.MockServer.GetMaxOffsetSent(); + + Assert.That(writeCount, Is.EqualTo((ulong)batch.Length)); + Assert.That(maxOffset, Is.EqualTo(0)); + } + + [Test] + public async Task IngestBatchRecords_AsyncApi() + { + await using var fixture = await MockServerFixture.StartAsync(); + + fixture.MockServer.InjectResponses(TestTableName, + [ + MockResponses.CreateStreamResponse("test_stream_async_batch"), + MockResponses.RecordAckResponse(0), + ]); + + using var sdk = CreateDefaultSdk(fixture); + var tableProps = CreateTableProperties(); + var options = CreateDefaultOptions(); + + await using var stream = await sdk.CreateStreamWithHeadersProviderAsync(tableProps, new TestHeadersProvider(), options); + + byte[][] batch = + [ + "record 1"u8.ToArray(), + "record 2"u8.ToArray(), + "record 3"u8.ToArray(), + "record 4"u8.ToArray(), + "record 5"u8.ToArray(), + ]; + + var offsetId = await stream.IngestRecordsAsync(batch); + + Assert.That(offsetId, Is.EqualTo(0)); + + await stream.WaitForOffsetAsync(offsetId); + + var writeCount = fixture.MockServer.GetWriteCount(); + var maxOffset = fixture.MockServer.GetMaxOffsetSent(); + + Assert.That(writeCount, Is.EqualTo((ulong)batch.Length)); + Assert.That(maxOffset, Is.EqualTo(0)); + } + + [Test] + public async Task IngestRecordsAfterClose() + { + await using var fixture = await MockServerFixture.StartAsync(); + + fixture.MockServer.InjectResponses(TestTableName, + [ + MockResponses.CreateStreamResponse("test_stream_batch_after_close"), + ]); + + using var sdk = CreateDefaultSdk(fixture); + var tableProps = CreateTableProperties(); + var options = CreateDefaultOptions(); + + var stream = sdk.CreateStreamWithHeadersProvider(tableProps, new TestHeadersProvider(), options); + + stream.Close(); + + byte[][] batch = ["record 1"u8.ToArray(), "record 2"u8.ToArray()]; + + Assert.That(() => stream.IngestRecords(batch), + Throws.InstanceOf()); + } + + [Test] + public async Task IngestRecordsAfterClose_AsyncApi() + { + await using var fixture = await MockServerFixture.StartAsync(); + + fixture.MockServer.InjectResponses(TestTableName, + [ + MockResponses.CreateStreamResponse("test_stream_batch_after_close_async"), + ]); + + using var sdk = CreateDefaultSdk(fixture); + var tableProps = CreateTableProperties(); + var options = CreateDefaultOptions(); + + var stream = await sdk.CreateStreamWithHeadersProviderAsync(tableProps, new TestHeadersProvider(), options); + + await stream.CloseAsync(); + + byte[][] batch = ["record 1"u8.ToArray(), "record 2"u8.ToArray()]; + + Assert.ThrowsAsync(async () => + { + await stream.IngestRecordsAsync(batch); + }); + } +} diff --git a/dotnet/tests/Zerobus.IntegrationTests/IntegrationTestBase.cs b/dotnet/tests/Zerobus.IntegrationTests/IntegrationTestBase.cs new file mode 100644 index 00000000..5a613db6 --- /dev/null +++ b/dotnet/tests/Zerobus.IntegrationTests/IntegrationTestBase.cs @@ -0,0 +1,31 @@ +using NUnit.Framework; + +namespace Databricks.Zerobus.IntegrationTests; + +public abstract class IntegrationTestBase +{ + protected const string TestTableName = "test_catalog.test_schema.test_table"; + + protected static ZerobusSdk CreateDefaultSdk(MockServerFixture fixture) + { + return ZerobusSdk.CreateBuilder() + .Endpoint(fixture.ServerUrl) + .UnityCatalogUrl("https://mock-uc.com") + .DisableTls() + .Build(); + } + + protected static TableProperties CreateTableProperties() + { + return new TableProperties(TestTableName, TestDescriptor.CreateTestDescriptorProto()); + } + + protected static StreamConfigurationOptions CreateDefaultOptions() + { + return StreamConfigurationOptions.Default with + { + MaxInflightRequests = 100, + Recovery = false, + }; + } +} diff --git a/dotnet/tests/Zerobus.IntegrationTests/LifecycleRecoveryIntegrationTests.cs b/dotnet/tests/Zerobus.IntegrationTests/LifecycleRecoveryIntegrationTests.cs new file mode 100644 index 00000000..e7eaee84 --- /dev/null +++ b/dotnet/tests/Zerobus.IntegrationTests/LifecycleRecoveryIntegrationTests.cs @@ -0,0 +1,712 @@ +using System.Collections.Concurrent; +using Grpc.Core; +using NUnit.Framework; + +namespace Databricks.Zerobus.IntegrationTests; + +[TestFixture] +[Parallelizable(ParallelScope.Children)] +public class LifecycleRecoveryIntegrationTests : IntegrationTestBase +{ + [Test] + public async Task GracefulClose() + { + await using var fixture = await MockServerFixture.StartAsync(); + + fixture.MockServer.InjectResponses(TestTableName, + [ + MockResponses.CreateStreamResponse("test_stream_1"), + MockResponses.RecordAckResponse(0, delayMs: 100), + ]); + + using var sdk = CreateDefaultSdk(fixture); + var tableProps = CreateTableProperties(); + var options = CreateDefaultOptions(); + + var stream = sdk.CreateStreamWithHeadersProvider(tableProps, new TestHeadersProvider(), options); + + var testRecord = "test record data"u8.ToArray(); + var offsetId = stream.IngestRecord(testRecord); + + Assert.That(offsetId, Is.EqualTo(0)); + + stream.Close(); + + var writeCount = fixture.MockServer.GetWriteCount(); + var maxOffset = fixture.MockServer.GetMaxOffsetSent(); + + Assert.That(writeCount, Is.EqualTo(1)); + Assert.That(maxOffset, Is.EqualTo(0)); + } + + [Test] + public async Task GracefulClose_AsyncApi() + { + await using var fixture = await MockServerFixture.StartAsync(); + + fixture.MockServer.InjectResponses(TestTableName, + [ + MockResponses.CreateStreamResponse("test_stream_async_close"), + MockResponses.RecordAckResponse(0, delayMs: 100), + ]); + + using var sdk = CreateDefaultSdk(fixture); + var tableProps = CreateTableProperties(); + var options = CreateDefaultOptions(); + + var stream = await sdk.CreateStreamWithHeadersProviderAsync(tableProps, new TestHeadersProvider(), options); + + var testRecord = "test record data"u8.ToArray(); + var offsetId = await stream.IngestRecordAsync(testRecord); + + Assert.That(offsetId, Is.EqualTo(0)); + + await stream.CloseAsync(); + + var writeCount = fixture.MockServer.GetWriteCount(); + var maxOffset = fixture.MockServer.GetMaxOffsetSent(); + + Assert.That(writeCount, Is.EqualTo(1)); + Assert.That(maxOffset, Is.EqualTo(0)); + } + + [Test] + public async Task AsyncDispose_GracefullyClosesStream() + { + await using var fixture = await MockServerFixture.StartAsync(); + + fixture.MockServer.InjectResponses(TestTableName, + [ + MockResponses.CreateStreamResponse("test_stream_async_dispose"), + MockResponses.RecordAckResponse(0, delayMs: 100), + ]); + + using var sdk = CreateDefaultSdk(fixture); + var tableProps = CreateTableProperties(); + var options = CreateDefaultOptions(); + + await using (var stream = await sdk.CreateStreamWithHeadersProviderAsync(tableProps, new TestHeadersProvider(), options)) + { + var testRecord = "test record data"u8.ToArray(); + var offsetId = await stream.IngestRecordAsync(testRecord); + + Assert.That(offsetId, Is.EqualTo(0)); + } + + var writeCount = fixture.MockServer.GetWriteCount(); + var maxOffset = fixture.MockServer.GetMaxOffsetSent(); + + Assert.That(writeCount, Is.EqualTo(1)); + Assert.That(maxOffset, Is.EqualTo(0)); + } + + [Test] + public async Task IdempotentClose() + { + await using var fixture = await MockServerFixture.StartAsync(); + + fixture.MockServer.InjectResponses(TestTableName, + [ + MockResponses.CreateStreamResponse("test_stream_1"), + ]); + + using var sdk = CreateDefaultSdk(fixture); + var tableProps = CreateTableProperties(); + var options = CreateDefaultOptions(); + + var stream = sdk.CreateStreamWithHeadersProvider(tableProps, new TestHeadersProvider(), options); + + stream.Close(); + stream.Close(); + } + + [Test] + public async Task IdempotentClose_AsyncApi() + { + await using var fixture = await MockServerFixture.StartAsync(); + + fixture.MockServer.InjectResponses(TestTableName, + [ + MockResponses.CreateStreamResponse("test_stream_async_idempotent_close"), + ]); + + using var sdk = CreateDefaultSdk(fixture); + var tableProps = CreateTableProperties(); + var options = CreateDefaultOptions(); + + var stream = await sdk.CreateStreamWithHeadersProviderAsync(tableProps, new TestHeadersProvider(), options); + + await stream.CloseAsync(); + await stream.CloseAsync(); + } + + [Test] + public async Task IngestAfterClose() + { + await using var fixture = await MockServerFixture.StartAsync(); + + fixture.MockServer.InjectResponses(TestTableName, + [ + MockResponses.CreateStreamResponse("test_stream_1"), + ]); + + using var sdk = CreateDefaultSdk(fixture); + var tableProps = CreateTableProperties(); + var options = CreateDefaultOptions(); + + var stream = sdk.CreateStreamWithHeadersProvider(tableProps, new TestHeadersProvider(), options); + stream.Close(); + + Assert.That(() => stream.IngestRecord("test record data"u8.ToArray()), + Throws.InstanceOf()); + } + + [Test] + public async Task IngestAfterClose_AsyncApi() + { + await using var fixture = await MockServerFixture.StartAsync(); + + fixture.MockServer.InjectResponses(TestTableName, + [ + MockResponses.CreateStreamResponse("test_stream_async_after_close"), + ]); + + using var sdk = CreateDefaultSdk(fixture); + var tableProps = CreateTableProperties(); + var options = CreateDefaultOptions(); + + var stream = await sdk.CreateStreamWithHeadersProviderAsync(tableProps, new TestHeadersProvider(), options); + await stream.CloseAsync(); + + Assert.ThrowsAsync(async () => + { + await stream.IngestRecordAsync("test record data"u8.ToArray()); + }); + } + + [Test] + public async Task ConcurrentIngestAndClose_DoesNotProduceUnexpectedExceptions() + { + await using var fixture = await MockServerFixture.StartAsync(); + + const int workerCount = 12; + const int recordsPerWorker = 60; + const int maxRecords = workerCount * recordsPerWorker; + + var responses = new List + { + MockResponses.CreateStreamResponse("test_stream_concurrent_close"), + }; + + for (var i = 0; i < maxRecords; i++) + { + responses.Add(MockResponses.RecordAckResponse(i)); + } + + fixture.MockServer.InjectResponses(TestTableName, responses); + + using var sdk = CreateDefaultSdk(fixture); + var tableProps = CreateTableProperties(); + var options = CreateDefaultOptions(); + + using var stream = sdk.CreateStreamWithHeadersProvider(tableProps, new TestHeadersProvider(), options); + + var startGate = new ManualResetEventSlim(false); + var unexpected = new ConcurrentBag(); + + var workers = Enumerable.Range(0, workerCount) + .Select(workerId => Task.Run(() => + { + startGate.Wait(); + + for (var i = 0; i < recordsPerWorker; i++) + { + try + { + var payload = System.Text.Encoding.UTF8.GetBytes($"close-race-{workerId}-{i}"); + _ = stream.IngestRecord(payload); + } + catch (ObjectDisposedException) + { + break; + } + catch (ZerobusException) + { + break; + } + catch (Exception ex) + { + unexpected.Add(ex); + break; + } + } + })) + .ToArray(); + + var closer = Task.Run(async () => + { + startGate.Wait(); + await Task.Delay(20); + + try + { + stream.Close(); + } + catch (ZerobusException) + { + // Close races with active ingesters by design in this stress test. + } + }); + + startGate.Set(); + await Task.WhenAll(workers.Concat([closer])); + + Assert.That(unexpected, Is.Empty); + } + + [Test] + public async Task ConcurrentIngestAndDispose_DoesNotProduceUnexpectedExceptions() + { + await using var fixture = await MockServerFixture.StartAsync(); + + const int workerCount = 12; + const int recordsPerWorker = 60; + const int maxRecords = workerCount * recordsPerWorker; + + var responses = new List + { + MockResponses.CreateStreamResponse("test_stream_concurrent_dispose"), + }; + + for (var i = 0; i < maxRecords; i++) + { + responses.Add(MockResponses.RecordAckResponse(i)); + } + + fixture.MockServer.InjectResponses(TestTableName, responses); + + using var sdk = CreateDefaultSdk(fixture); + var tableProps = CreateTableProperties(); + var options = CreateDefaultOptions(); + + var stream = sdk.CreateStreamWithHeadersProvider(tableProps, new TestHeadersProvider(), options); + + var startGate = new ManualResetEventSlim(false); + var unexpected = new ConcurrentBag(); + + var workers = Enumerable.Range(0, workerCount) + .Select(workerId => Task.Run(() => + { + startGate.Wait(); + + for (var i = 0; i < recordsPerWorker; i++) + { + try + { + var payload = System.Text.Encoding.UTF8.GetBytes($"dispose-race-{workerId}-{i}"); + _ = stream.IngestRecord(payload); + } + catch (ObjectDisposedException) + { + break; + } + catch (ZerobusException) + { + break; + } + catch (Exception ex) + { + unexpected.Add(ex); + break; + } + } + })) + .ToArray(); + + var disposer = Task.Run(async () => + { + startGate.Wait(); + await Task.Delay(20); + + try + { + stream.Dispose(); + } + catch (ZerobusException) + { + // Dispose may surface close failure when racing with active ingesters. + } + }); + + startGate.Set(); + await Task.WhenAll(workers.Concat([disposer])); + + Assert.That(unexpected, Is.Empty); + } + + [Test] + public async Task ConcurrentIngestThenCloseAndRecreate_DoesNotProduceUnexpectedExceptions() + { + await using var fixture = await MockServerFixture.StartAsync(); + + var responses = new List + { + MockResponses.CreateStreamResponse("test_stream_concurrent_recreate_1"), + }; + + for (var i = 0; i < 48; i++) + { + responses.Add(MockResponses.RecordAckResponse(i)); + } + + responses.Add(MockResponses.CreateStreamResponse("test_stream_concurrent_recreate_2")); + responses.Add(MockResponses.RecordAckResponse(0)); + + fixture.MockServer.InjectResponses(TestTableName, responses); + + using var sdk = CreateDefaultSdk(fixture); + + var tableProps = CreateTableProperties(); + var options = CreateDefaultOptions(); + + var currentStream = sdk.CreateStreamWithHeadersProvider(tableProps, new TestHeadersProvider(), options); + var sync = new object(); + var startGate = new ManualResetEventSlim(false); + var unexpected = new ConcurrentBag(); + + const int workerCount = 4; + const int iterationsPerWorker = 4; + + var workers = Enumerable.Range(0, workerCount) + .Select(workerId => Task.Run(() => + { + startGate.Wait(); + + for (var i = 0; i < iterationsPerWorker; i++) + { + ZerobusStream target; + lock (sync) + { + target = currentStream; + } + + try + { + var payload = System.Text.Encoding.UTF8.GetBytes($"recreate-race-{workerId}-{i}"); + _ = target.IngestRecord(payload); + } + catch (ObjectDisposedException) + { + break; + } + catch (ZerobusException) + { + break; + } + catch (Exception ex) + { + unexpected.Add(ex); + break; + } + } + })) + .ToArray(); + + var recreater = Task.Run(async () => + { + startGate.Wait(); + await Task.Delay(5); + + ZerobusStream oldStream; + lock (sync) + { + oldStream = currentStream; + } + + try + { + oldStream.Close(); + } + catch (ZerobusException) + { + // Close can race with active ingest; stream recreation should still be valid. + } + + using var recreated = sdk.RecreateStream(oldStream); + var recreatedOffset = recreated.IngestRecord("post-recreate"u8.ToArray()); + recreated.WaitForOffset(recreatedOffset); + }); + + startGate.Set(); + await Task.WhenAll(workers.Concat([recreater])); + + Assert.That(unexpected, Is.Empty); + } + + [Test] + public async Task RecreateStream_ClosedStream_RecreatedStreamUsable() + { + await using var fixture = await MockServerFixture.StartAsync(); + + fixture.MockServer.InjectResponses(TestTableName, + [ + MockResponses.CreateStreamResponse("test_stream_1"), + MockResponses.CreateStreamResponse("test_stream_2"), + MockResponses.RecordAckResponse(0), + ]); + + using var sdk = CreateDefaultSdk(fixture); + var tableProps = CreateTableProperties(); + var options = CreateDefaultOptions(); + + var stream = sdk.CreateStreamWithHeadersProvider(tableProps, new TestHeadersProvider(), options); + stream.Close(); + + var unacked = stream.GetUnackedRecords(); + Assert.That(unacked, Is.Empty); + + using var recreatedStream = sdk.RecreateStream(stream); + + Assert.That(recreatedStream, Is.Not.Null); + Assert.That(recreatedStream, Is.Not.SameAs(stream)); + + Assert.That(() => stream.IngestRecord("old stream"u8.ToArray()), + Throws.InstanceOf()); + + var offsetId = recreatedStream.IngestRecord("recreated stream"u8.ToArray()); + Assert.That(offsetId, Is.EqualTo(0)); + + recreatedStream.WaitForOffset(offsetId); + + var writeCount = fixture.MockServer.GetWriteCount(); + var maxOffset = fixture.MockServer.GetMaxOffsetSent(); + + Assert.That(writeCount, Is.EqualTo(1)); + Assert.That(maxOffset, Is.EqualTo(0)); + } + + [Test] + public async Task RecreateStream_ClosedStream_RecreatedStreamUsable_AsyncApi() + { + await using var fixture = await MockServerFixture.StartAsync(); + + fixture.MockServer.InjectResponses(TestTableName, + [ + MockResponses.CreateStreamResponse("test_stream_async_recreate_1"), + MockResponses.CreateStreamResponse("test_stream_async_recreate_2"), + MockResponses.RecordAckResponse(0), + ]); + + using var sdk = CreateDefaultSdk(fixture); + var tableProps = CreateTableProperties(); + var options = CreateDefaultOptions(); + + var stream = await sdk.CreateStreamWithHeadersProviderAsync(tableProps, new TestHeadersProvider(), options); + await stream.CloseAsync(); + + var unacked = await stream.GetUnackedRecordsAsync(); + Assert.That(unacked, Is.Empty); + + await using var recreatedStream = await sdk.RecreateStreamAsync(stream); + + Assert.That(recreatedStream, Is.Not.Null); + Assert.That(recreatedStream, Is.Not.SameAs(stream)); + + Assert.ThrowsAsync(async () => + { + await stream.IngestRecordAsync("old stream"u8.ToArray()); + }); + + var offsetId = await recreatedStream.IngestRecordAsync("recreated stream"u8.ToArray()); + Assert.That(offsetId, Is.EqualTo(0)); + + await recreatedStream.WaitForOffsetAsync(offsetId); + + var writeCount = fixture.MockServer.GetWriteCount(); + var maxOffset = fixture.MockServer.GetMaxOffsetSent(); + + Assert.That(writeCount, Is.EqualTo(1)); + Assert.That(maxOffset, Is.EqualTo(0)); + } + + [Test] + public async Task RecreateStream_JsonTypedStream_RecreatedStreamUsable() + { + await using var fixture = await MockServerFixture.StartAsync(); + + fixture.MockServer.InjectResponses(TestTableName, + [ + MockResponses.CreateStreamResponse("test_stream_json_1"), + MockResponses.CreateStreamResponse("test_stream_json_2"), + MockResponses.RecordAckResponse(0), + ]); + + using var sdk = CreateDefaultSdk(fixture); + var options = CreateDefaultOptions(); + + var stream = sdk.CreateJsonStreamWithHeadersProvider( + TestTableName, + new TestHeadersProvider(), + options); + stream.Close(); + + using var recreatedStream = sdk.RecreateStream(stream); + + var offsetId = recreatedStream.IngestRecord("{\"message\":\"recreated\"}"); + recreatedStream.WaitForOffset(offsetId); + + Assert.That(offsetId, Is.EqualTo(0)); + } + + [Test] + public async Task RecreateStream_JsonTypedStream_RecreatedStreamUsable_AsyncApi() + { + await using var fixture = await MockServerFixture.StartAsync(); + + fixture.MockServer.InjectResponses(TestTableName, + [ + MockResponses.CreateStreamResponse("test_stream_json_async_1"), + MockResponses.CreateStreamResponse("test_stream_json_async_2"), + MockResponses.RecordAckResponse(0), + ]); + + using var sdk = CreateDefaultSdk(fixture); + var options = CreateDefaultOptions(); + + var stream = await sdk.CreateJsonStreamWithHeadersProviderAsync( + TestTableName, + new TestHeadersProvider(), + options); + await stream.CloseAsync(); + + await using var recreatedStream = await sdk.RecreateStreamAsync(stream); + + var offsetId = await recreatedStream.IngestRecordAsync("{\"message\":\"recreated\"}"); + await recreatedStream.WaitForOffsetAsync(offsetId); + + Assert.That(offsetId, Is.EqualTo(0)); + } + + [Test] + public async Task FlushFailure_GetUnackedRecords_ThenRecreateStream_Works() + { + await using var fixture = await MockServerFixture.StartAsync(); + + fixture.MockServer.InjectResponses(TestTableName, + [ + MockResponses.CreateStreamResponse("test_stream_flush_failure_1"), + MockResponses.ErrorResponse(StatusCode.Unavailable, "transient ack failure"), + MockResponses.CreateStreamResponse("test_stream_flush_failure_2"), + MockResponses.RecordAckResponse(1), + ]); + + using var sdk = CreateDefaultSdk(fixture); + + var tableProps = CreateTableProperties(); + var options = CreateDefaultOptions(); + + var stream = sdk.CreateStreamWithHeadersProvider(tableProps, new TestHeadersProvider(), options); + + var failedPayload = "failed-before-flush"u8.ToArray(); + _ = stream.IngestRecord(failedPayload); + + Assert.That(() => stream.Flush(), Throws.InstanceOf()); + + var unacked = stream.GetUnackedRecords(); + Assert.That(unacked, Has.Length.EqualTo(1)); + Assert.That(unacked[0].ToArray(), Is.EqualTo(failedPayload)); + + using var recreated = sdk.RecreateStream(stream); + + var recoveredOffset = recreated.IngestRecord("recovered-record"u8.ToArray()); + Assert.That(recoveredOffset, Is.GreaterThanOrEqualTo(0)); + } + + [Test] + public async Task RecreateStream_RetryableFailure_ThrowsRetryableException() + { + await using var fixture = await MockServerFixture.StartAsync(); + + fixture.MockServer.InjectResponses(TestTableName, + [ + MockResponses.CreateStreamResponse("test_stream_recreate_retryable_1"), + MockResponses.ErrorResponse(StatusCode.Unavailable, "recreate unavailable"), + ]); + + using var sdk = CreateDefaultSdk(fixture); + + var tableProps = CreateTableProperties(); + var options = CreateDefaultOptions(); + + using var stream = sdk.CreateStreamWithHeadersProvider(tableProps, new TestHeadersProvider(), options); + stream.Close(); + + var ex = Assert.Throws(() => sdk.RecreateStream(stream)); + Assert.That(ex, Is.Not.Null); + Assert.That(ex!.IsRetryable, Is.True); + } + + [Test] + public async Task RecreateStream_RetryableFailure_ThrowsRetryableException_AsyncApi() + { + await using var fixture = await MockServerFixture.StartAsync(); + + fixture.MockServer.InjectResponses(TestTableName, + [ + MockResponses.CreateStreamResponse("test_stream_recreate_retryable_async_1"), + MockResponses.ErrorResponse(StatusCode.Unavailable, "recreate unavailable"), + ]); + + using var sdk = CreateDefaultSdk(fixture); + + var tableProps = CreateTableProperties(); + var options = CreateDefaultOptions(); + + await using var stream = await sdk.CreateStreamWithHeadersProviderAsync(tableProps, new TestHeadersProvider(), options); + await stream.CloseAsync(); + + var ex = Assert.ThrowsAsync(async () => await sdk.RecreateStreamAsync(stream)); + Assert.That(ex, Is.Not.Null); + Assert.That(ex!.IsRetryable, Is.True); + } + + [Test] + public async Task RecreateStream_ActiveStream_ThrowsZerobusException() + { + await using var fixture = await MockServerFixture.StartAsync(); + + fixture.MockServer.InjectResponses(TestTableName, + [ + MockResponses.CreateStreamResponse("test_stream_1"), + ]); + + using var sdk = CreateDefaultSdk(fixture); + var tableProps = CreateTableProperties(); + var options = CreateDefaultOptions(); + + using var stream = sdk.CreateStreamWithHeadersProvider(tableProps, new TestHeadersProvider(), options); + + var ex = Assert.Throws(() => sdk.RecreateStream(stream)); + Assert.That(ex!.Message, Does.Contain("active stream")); + } + + [Test] + public async Task RecreateStream_ActiveStream_ThrowsZerobusException_AsyncApi() + { + await using var fixture = await MockServerFixture.StartAsync(); + + fixture.MockServer.InjectResponses(TestTableName, + [ + MockResponses.CreateStreamResponse("test_stream_active_async_1"), + ]); + + using var sdk = CreateDefaultSdk(fixture); + var tableProps = CreateTableProperties(); + var options = CreateDefaultOptions(); + + await using var stream = await sdk.CreateStreamWithHeadersProviderAsync(tableProps, new TestHeadersProvider(), options); + + var ex = Assert.ThrowsAsync(async () => await sdk.RecreateStreamAsync(stream)); + Assert.That(ex!.Message, Does.Contain("active stream")); + } +} diff --git a/dotnet/tests/Zerobus.IntegrationTests/MockZerobusServer.cs b/dotnet/tests/Zerobus.IntegrationTests/MockZerobusServer.cs new file mode 100644 index 00000000..0e767aaf --- /dev/null +++ b/dotnet/tests/Zerobus.IntegrationTests/MockZerobusServer.cs @@ -0,0 +1,420 @@ +using Grpc.Core; +using Google.Protobuf.WellKnownTypes; +using Databricks.Zerobus.IntegrationTests.Protos; + +namespace Databricks.Zerobus.IntegrationTests; + +/// +/// The type of a mock response that can be injected into the mock server. +/// +public enum MockResponseType +{ + CreateStream, + RecordAck, + CloseStreamSignal, + Error, +} + +/// +/// Represents a response that can be injected into the mock server. +/// +public sealed class MockResponse +{ + public MockResponseType Type { get; init; } + public string? StreamId { get; init; } + public long DelayMs { get; init; } + public long AckUpToOffset { get; init; } + public long DurationSeconds { get; init; } + public StatusCode? GrpcStatusCode { get; init; } + public string? GrpcMessage { get; init; } +} + +/// +/// Mock gRPC server that implements the Zerobus EphemeralStream RPC for integration testing. +/// Faithfully ports the Go mock_server.go implementation. +/// +public sealed class MockZerobusServer : Protos.Zerobus.ZerobusBase +{ + private readonly object _responsesLock = new(); + private readonly object _counterLock = new(); + private readonly object _offsetLock = new(); + private readonly object _writeCountLock = new(); + private readonly object _responseIndicesLock = new(); + private readonly object _requestHeadersLock = new(); + + private Dictionary> _responses = new(); + private Dictionary _responseIndices = new(); + private Dictionary> _lastRequestHeadersByTable = new(); + private int _streamCounter; + private long _maxOffsetSent = -1; + private ulong _writeCount; + + /// + /// Configures responses for a specific table. + /// + public void InjectResponses(string tableName, List responses) + { + lock (_responsesLock) + { + _responses[tableName] = new List(responses); + } + + lock (_responseIndicesLock) + { + _responseIndices[tableName] = 0; + } + } + + /// + /// Returns the maximum offset sent by clients. + /// + public long GetMaxOffsetSent() + { + lock (_offsetLock) + { + return _maxOffsetSent; + } + } + + /// + /// Returns the number of writes received. + /// + public ulong GetWriteCount() + { + lock (_writeCountLock) + { + return _writeCount; + } + } + + /// + /// Returns the request headers seen for the latest create stream call for a table. + /// + public IDictionary GetLastRequestHeaders(string tableName) + { + lock (_requestHeadersLock) + { + return _lastRequestHeadersByTable.TryGetValue(tableName, out var headers) + ? new Dictionary(headers) + : new Dictionary(); + } + } + + /// + /// Resets all server state. + /// + public void Reset() + { + lock (_responsesLock) { _responses = new Dictionary>(); } + lock (_responseIndicesLock) { _responseIndices = new Dictionary(); } + lock (_requestHeadersLock) { _lastRequestHeadersByTable = new Dictionary>(); } + lock (_offsetLock) { _maxOffsetSent = -1; } + lock (_writeCountLock) { _writeCount = 0; } + lock (_counterLock) { _streamCounter = 0; } + } + + /// + /// Implements the bidirectional streaming EphemeralStream RPC. + /// + public override async Task EphemeralStream( + IAsyncStreamReader requestStream, + IServerStreamWriter responseStream, + ServerCallContext context) + { + // Send initial headers/metadata to signal readiness. + // This is CRITICAL: Without this, the Rust SDK will wait indefinitely. + await context.WriteResponseHeadersAsync(new Metadata()); + + // Read the first message (CreateStream request). + if (!await requestStream.MoveNext(context.CancellationToken)) + { + return; + } + + var firstRequest = requestStream.Current; + var createReq = firstRequest.CreateStream; + if (createReq is null) + { + throw new RpcException(new Status(StatusCode.InvalidArgument, "first message must be CreateStream")); + } + + var tableName = createReq.TableName ?? ""; + + // Capture incoming metadata so tests can validate custom header callback behavior. + var requestHeaders = new Dictionary(StringComparer.OrdinalIgnoreCase); + foreach (var entry in context.RequestHeaders) + { + requestHeaders[entry.Key] = entry.IsBinary + ? Convert.ToBase64String(entry.ValueBytes ?? []) + : entry.Value ?? string.Empty; + } + + lock (_requestHeadersLock) + { + _lastRequestHeadersByTable[tableName] = requestHeaders; + } + + // Get stream ID. + string streamId; + lock (_counterLock) + { + _streamCounter++; + streamId = $"test_stream_{_streamCounter}"; + } + + // Get configured responses for this table. + List streamResponses; + lock (_responsesLock) + { + streamResponses = _responses.TryGetValue(tableName, out var responses) + ? new List(responses) + : []; + } + + // Get current response index. + int responseIndex; + lock (_responseIndicesLock) + { + _responseIndices.TryGetValue(tableName, out responseIndex); + } + + // Search for the next CreateStream response starting from responseIndex. + var createStreamFound = false; + for (var idx = responseIndex; idx < streamResponses.Count; idx++) + { + var mockResp = streamResponses[idx]; + if (mockResp.Type == MockResponseType.CreateStream) + { + responseIndex = idx; + createStreamFound = true; + + if (mockResp.DelayMs > 0) + { + await Task.Delay(TimeSpan.FromMilliseconds(mockResp.DelayMs)); + } + + var customId = mockResp.StreamId ?? streamId; + + await responseStream.WriteAsync(new EphemeralStreamResponse + { + CreateStreamResponse = new CreateIngestStreamResponse + { + StreamId = customId, + }, + }); + + responseIndex++; + lock (_responseIndicesLock) + { + _responseIndices[tableName] = responseIndex; + } + + break; + } + else if (mockResp.Type == MockResponseType.Error) + { + if (mockResp.DelayMs > 0) + { + await Task.Delay(TimeSpan.FromMilliseconds(mockResp.DelayMs)); + } + + lock (_responseIndicesLock) + { + _responseIndices[tableName] = idx + 1; + } + + throw new RpcException(new Status( + mockResp.GrpcStatusCode ?? StatusCode.Internal, + mockResp.GrpcMessage ?? "Unknown error")); + } + } + + // If no CreateStream response was configured, send default. + if (!createStreamFound) + { + await responseStream.WriteAsync(new EphemeralStreamResponse + { + CreateStreamResponse = new CreateIngestStreamResponse + { + StreamId = streamId, + }, + }); + } + + // Process subsequent requests. + while (await requestStream.MoveNext(context.CancellationToken)) + { + var request = requestStream.Current; + + if (request.IngestRecord is { } ingestRecord) + { + responseIndex = await HandleIngestRecord(responseStream, ingestRecord, streamResponses, responseIndex, tableName); + } + else if (request.IngestRecordBatch is { } ingestBatch) + { + responseIndex = await HandleIngestRecordBatch(responseStream, ingestBatch, streamResponses, responseIndex, tableName); + } + } + } + + private async Task HandleIngestRecord( + IServerStreamWriter responseStream, + IngestRecordRequest req, + List streamResponses, + int responseIndex, + string tableName) + { + // Update max offset. + if (req.HasOffsetId) + { + lock (_offsetLock) + { + if (req.OffsetId > _maxOffsetSent) + { + _maxOffsetSent = req.OffsetId; + } + } + } + + // Increment write count. + lock (_writeCountLock) + { + _writeCount++; + } + + // Process mock response. + if (responseIndex < streamResponses.Count) + { + var (shouldContinue, indexIncrement) = await HandleMockResponse( + responseStream, streamResponses[responseIndex], req.HasOffsetId ? req.OffsetId : null, tableName); + responseIndex += indexIncrement; + + lock (_responseIndicesLock) + { + _responseIndices[tableName] = responseIndex; + } + + if (!shouldContinue) + { + throw new RpcException(new Status(StatusCode.Internal, "mock response indicated stop")); + } + } + + return responseIndex; + } + + private async Task HandleIngestRecordBatch( + IServerStreamWriter responseStream, + IngestRecordBatchRequest req, + List streamResponses, + int responseIndex, + string tableName) + { + // Count records in the batch. + var recordCount = 0; + if (req.ProtoEncodedBatch is { } protoBatch) + { + recordCount = protoBatch.Records.Count; + } + else if (req.JsonBatch is { } jsonBatch) + { + recordCount = jsonBatch.Records.Count; + } + + // Update max offset. + if (req.HasOffsetId) + { + lock (_offsetLock) + { + if (req.OffsetId > _maxOffsetSent) + { + _maxOffsetSent = req.OffsetId; + } + } + } + + // Increment write count by number of records. + lock (_writeCountLock) + { + _writeCount += (ulong)recordCount; + } + + // Process mock response. + if (responseIndex < streamResponses.Count) + { + var (shouldContinue, indexIncrement) = await HandleMockResponse( + responseStream, streamResponses[responseIndex], req.HasOffsetId ? req.OffsetId : null, tableName); + responseIndex += indexIncrement; + + lock (_responseIndicesLock) + { + _responseIndices[tableName] = responseIndex; + } + + if (!shouldContinue) + { + throw new RpcException(new Status(StatusCode.Internal, "mock response indicated stop")); + } + } + + return responseIndex; + } + + private async Task<(bool ShouldContinue, int IndexIncrement)> HandleMockResponse( + IServerStreamWriter responseStream, + MockResponse mockResp, + long? offset, + string tableName) + { + switch (mockResp.Type) + { + case MockResponseType.RecordAck: + if (offset.HasValue && offset.Value == mockResp.AckUpToOffset) + { + if (mockResp.DelayMs > 0) + { + await Task.Delay(TimeSpan.FromMilliseconds(mockResp.DelayMs)); + } + + await responseStream.WriteAsync(new EphemeralStreamResponse + { + IngestRecordResponse = new IngestRecordResponse + { + DurabilityAckUpToOffset = mockResp.AckUpToOffset, + }, + }); + + return (true, 1); + } + return (true, 0); + + case MockResponseType.CloseStreamSignal: + if (mockResp.DelayMs > 0) + { + await Task.Delay(TimeSpan.FromMilliseconds(mockResp.DelayMs)); + } + + await responseStream.WriteAsync(new EphemeralStreamResponse + { + CloseStreamSignal = new CloseStreamSignal + { + Duration = Duration.FromTimeSpan(TimeSpan.FromSeconds(mockResp.DurationSeconds)), + }, + }); + return (true, 1); + + case MockResponseType.Error: + if (mockResp.DelayMs > 0) + { + await Task.Delay(TimeSpan.FromMilliseconds(mockResp.DelayMs)); + } + return (false, 1); + + case MockResponseType.CreateStream: + return (true, 1); + + default: + return (true, 0); + } + } +} diff --git a/dotnet/tests/Zerobus.IntegrationTests/Protos/zerobus_service.proto b/dotnet/tests/Zerobus.IntegrationTests/Protos/zerobus_service.proto new file mode 100644 index 00000000..144dd217 --- /dev/null +++ b/dotnet/tests/Zerobus.IntegrationTests/Protos/zerobus_service.proto @@ -0,0 +1,79 @@ +syntax = "proto2"; + +import "google/protobuf/duration.proto"; + +package databricks.zerobus; + +option csharp_namespace = "Databricks.Zerobus.IntegrationTests.Protos"; + +service Zerobus { + rpc EphemeralStream(stream EphemeralStreamRequest) returns (stream EphemeralStreamResponse); +} + +enum RecordType { + RECORD_TYPE_UNSPECIFIED = 0; + PROTO = 1; + JSON = 2; +} + +message JsonRecordBatch { + repeated string records = 1; +} + +message ProtoEncodedRecordBatch { + repeated bytes records = 1; +} + +message CreateIngestStreamRequest { + optional string table_name = 1; + reserved "stream_id"; + reserved 2; + optional bytes descriptor_proto = 3; + optional RecordType record_type = 4; +} + +message CreateIngestStreamResponse { + optional string stream_id = 1; + reserved "last_offset_id"; + reserved 2; +} + +message IngestRecordRequest { + optional int64 offset_id = 1; + oneof record { + bytes proto_encoded_record = 2; + string json_record = 3; + } +} + +message IngestRecordBatchRequest { + optional int64 offset_id = 1; + oneof batch { + ProtoEncodedRecordBatch proto_encoded_batch = 2; + JsonRecordBatch json_batch = 3; + } +} + +message EphemeralStreamRequest { + oneof payload { + CreateIngestStreamRequest create_stream = 1; + IngestRecordRequest ingest_record = 2; + IngestRecordBatchRequest ingest_record_batch = 3; + } +} + +message IngestRecordResponse { + optional int64 durability_ack_up_to_offset = 1; +} + +message CloseStreamSignal { + optional google.protobuf.Duration duration = 1; +} + +message EphemeralStreamResponse { + oneof payload { + CreateIngestStreamResponse create_stream_response = 1; + IngestRecordResponse ingest_record_response = 2; + CloseStreamSignal close_stream_signal = 3; + } +} diff --git a/dotnet/tests/Zerobus.IntegrationTests/StreamCreationIntegrationTests.cs b/dotnet/tests/Zerobus.IntegrationTests/StreamCreationIntegrationTests.cs new file mode 100644 index 00000000..3b75b7aa --- /dev/null +++ b/dotnet/tests/Zerobus.IntegrationTests/StreamCreationIntegrationTests.cs @@ -0,0 +1,461 @@ +using System.Diagnostics; +using Grpc.Core; +using NUnit.Framework; + +namespace Databricks.Zerobus.IntegrationTests; + +[TestFixture] +[Parallelizable(ParallelScope.Children)] +public class StreamCreationIntegrationTests : IntegrationTestBase +{ + [Test] + public async Task CreateStream_NullTableProperties_ThrowsArgumentNullException() + { + await using var fixture = await MockServerFixture.StartAsync(); + + using var sdk = CreateDefaultSdk(fixture); + + Assert.Throws(() => + { + sdk.CreateStream(null!, "id", "secret"); + }); + } + + [Test] + public async Task CreateStreamAsync_NullTableProperties_ThrowsArgumentNullException() + { + await using var fixture = await MockServerFixture.StartAsync(); + + using var sdk = CreateDefaultSdk(fixture); + + Assert.ThrowsAsync(async () => + { + await sdk.CreateStreamAsync(null!, "id", "secret"); + }); + } + + [Test] + public async Task CreateStream_NullClientId_ThrowsArgumentNullException() + { + await using var fixture = await MockServerFixture.StartAsync(); + + using var sdk = CreateDefaultSdk(fixture); + + Assert.Throws(() => + { + sdk.CreateStream(new TableProperties("test_table"), null!, "secret"); + }); + } + + [Test] + public async Task CreateStreamAsync_NullClientId_ThrowsArgumentNullException() + { + await using var fixture = await MockServerFixture.StartAsync(); + + using var sdk = CreateDefaultSdk(fixture); + + Assert.ThrowsAsync(async () => + { + await sdk.CreateStreamAsync(new TableProperties("test_table"), null!, "secret"); + }); + } + + [Test] + public async Task CreateStream_NullClientSecret_ThrowsArgumentNullException() + { + await using var fixture = await MockServerFixture.StartAsync(); + + using var sdk = CreateDefaultSdk(fixture); + + Assert.Throws(() => + { + sdk.CreateStream(new TableProperties("test_table"), "id", null!); + }); + } + + [Test] + public async Task CreateStreamAsync_NullClientSecret_ThrowsArgumentNullException() + { + await using var fixture = await MockServerFixture.StartAsync(); + + using var sdk = CreateDefaultSdk(fixture); + + Assert.ThrowsAsync(async () => + { + await sdk.CreateStreamAsync(new TableProperties("test_table"), "id", null!); + }); + } + + [Test] + public async Task SuccessfulStreamCreation() + { + await using var fixture = await MockServerFixture.StartAsync(); + + fixture.MockServer.InjectResponses(TestTableName, + [ + MockResponses.CreateStreamResponse("test_stream_1"), + ]); + + using var sdk = CreateDefaultSdk(fixture); + var tableProps = CreateTableProperties(); + var options = CreateDefaultOptions(); + + using var stream = sdk.CreateStreamWithHeadersProvider(tableProps, new TestHeadersProvider(), options); + + Assert.That(stream, Is.Not.Null); + } + + [Test] + public async Task SuccessfulStreamCreation_AsyncApi() + { + await using var fixture = await MockServerFixture.StartAsync(); + + fixture.MockServer.InjectResponses(TestTableName, + [ + MockResponses.CreateStreamResponse("test_stream_async_success"), + ]); + + using var sdk = CreateDefaultSdk(fixture); + var tableProps = CreateTableProperties(); + var options = CreateDefaultOptions(); + + await using var stream = await sdk.CreateStreamWithHeadersProviderAsync(tableProps, new TestHeadersProvider(), options); + + Assert.That(stream, Is.Not.Null); + } + + [Test] + public async Task CreateJsonStreamWithHeadersProvider_ReturnsTypedJsonStream() + { + await using var fixture = await MockServerFixture.StartAsync(); + + fixture.MockServer.InjectResponses(TestTableName, + [ + MockResponses.CreateStreamResponse("test_stream_json"), + MockResponses.RecordAckResponse(0), + ]); + + using var sdk = CreateDefaultSdk(fixture); + var options = CreateDefaultOptions(); + + using var stream = sdk.CreateJsonStreamWithHeadersProvider( + TestTableName, + new TestHeadersProvider(), + options); + + var offset = stream.IngestRecord("{\"message\":\"json\"}"); + stream.WaitForOffset(offset); + + Assert.That(offset, Is.EqualTo(0)); + } + + [Test] + public async Task CreateJsonStreamWithHeadersProviderAsync_ReturnsTypedJsonStream() + { + await using var fixture = await MockServerFixture.StartAsync(); + + fixture.MockServer.InjectResponses(TestTableName, + [ + MockResponses.CreateStreamResponse("test_stream_json_async"), + MockResponses.RecordAckResponse(0), + ]); + + using var sdk = CreateDefaultSdk(fixture); + var options = CreateDefaultOptions(); + + await using var stream = await sdk.CreateJsonStreamWithHeadersProviderAsync( + TestTableName, + new TestHeadersProvider(), + options); + + var offset = await stream.IngestRecordAsync("{\"message\":\"json\"}"); + await stream.WaitForOffsetAsync(offset); + + Assert.That(offset, Is.EqualTo(0)); + } + + [Test] + public async Task CreateProtoStreamWithHeadersProvider_ReturnsTypedProtoStream() + { + await using var fixture = await MockServerFixture.StartAsync(); + + fixture.MockServer.InjectResponses(TestTableName, + [ + MockResponses.CreateStreamResponse("test_stream_proto"), + MockResponses.RecordAckResponse(0), + ]); + + using var sdk = CreateDefaultSdk(fixture); + var options = CreateDefaultOptions(); + + using var stream = sdk.CreateProtoStreamWithHeadersProvider( + TestTableName, + TestDescriptor.CreateTestDescriptorProto(), + new TestHeadersProvider(), + options); + + var offset = stream.IngestRecord("proto-payload"u8.ToArray()); + stream.WaitForOffset(offset); + + Assert.That(offset, Is.EqualTo(0)); + } + + [Test] + public async Task CreateProtoStreamWithHeadersProviderAsync_ReturnsTypedProtoStream() + { + await using var fixture = await MockServerFixture.StartAsync(); + + fixture.MockServer.InjectResponses(TestTableName, + [ + MockResponses.CreateStreamResponse("test_stream_proto_async"), + MockResponses.RecordAckResponse(0), + ]); + + using var sdk = CreateDefaultSdk(fixture); + var options = CreateDefaultOptions(); + + await using var stream = await sdk.CreateProtoStreamWithHeadersProviderAsync( + TestTableName, + TestDescriptor.CreateTestDescriptorProto(), + new TestHeadersProvider(), + options); + + var offset = await stream.IngestRecordAsync("proto-payload"u8.ToArray()); + await stream.WaitForOffsetAsync(offset); + + Assert.That(offset, Is.EqualTo(0)); + } + + [Test] + public async Task CreateStream_JsonRecordTypeWithDescriptor_ThrowsArgumentException() + { + await using var fixture = await MockServerFixture.StartAsync(); + + using var sdk = CreateDefaultSdk(fixture); + var options = CreateDefaultOptions() with { RecordType = RecordType.Json }; + var tableProperties = CreateTableProperties(); + + var ex = Assert.Throws(() => + { + sdk.CreateStreamWithHeadersProvider(tableProperties, new TestHeadersProvider(), options); + }); + + Assert.That(ex!.Message, Does.Contain("JSON streams cannot specify DescriptorProto")); + } + + [Test] + public async Task CreateStream_ProtoRecordTypeWithoutDescriptor_ThrowsArgumentException() + { + await using var fixture = await MockServerFixture.StartAsync(); + + using var sdk = CreateDefaultSdk(fixture); + var options = CreateDefaultOptions() with { RecordType = RecordType.Proto }; + + var ex = Assert.Throws(() => + { + sdk.CreateStreamWithHeadersProvider( + new TableProperties(TestTableName), + new TestHeadersProvider(), + options); + }); + + Assert.That(ex!.Message, Does.Contain("Proto streams require a non-empty DescriptorProto")); + } + + [Test] + public async Task TimeoutedStreamCreation() + { + await using var fixture = await MockServerFixture.StartAsync(); + + fixture.MockServer.InjectResponses(TestTableName, + [ + MockResponses.CreateStreamResponse("test_stream_1", delayMs: 300), + ]); + + using var sdk = CreateDefaultSdk(fixture); + var tableProps = CreateTableProperties(); + + var options = StreamConfigurationOptions.Default with + { + MaxInflightRequests = 100, + RecoveryTimeoutMs = 100, + Recovery = false, + }; + + Assert.Throws(() => + { + sdk.CreateStreamWithHeadersProvider(tableProps, new TestHeadersProvider(), options); + }); + + await Task.Delay(100); + } + + [Test] + public async Task TimeoutedStreamCreation_AsyncApi() + { + await using var fixture = await MockServerFixture.StartAsync(); + + fixture.MockServer.InjectResponses(TestTableName, + [ + MockResponses.CreateStreamResponse("test_stream_timeout_async", delayMs: 300), + ]); + + using var sdk = CreateDefaultSdk(fixture); + var tableProps = CreateTableProperties(); + + var options = StreamConfigurationOptions.Default with + { + MaxInflightRequests = 100, + RecoveryTimeoutMs = 100, + Recovery = false, + }; + + Assert.ThrowsAsync(async () => + { + await sdk.CreateStreamWithHeadersProviderAsync(tableProps, new TestHeadersProvider(), options); + }); + } + + [Test] + public async Task NonRetriableErrorDuringStreamCreation() + { + await using var fixture = await MockServerFixture.StartAsync(); + + fixture.MockServer.InjectResponses(TestTableName, + [ + MockResponses.ErrorResponse(StatusCode.Unauthenticated, "Non-retriable error"), + ]); + + using var sdk = CreateDefaultSdk(fixture); + var tableProps = CreateTableProperties(); + + var options = StreamConfigurationOptions.Default with + { + MaxInflightRequests = 100, + Recovery = true, + }; + + Assert.Throws(() => + { + sdk.CreateStreamWithHeadersProvider(tableProps, new TestHeadersProvider(), options); + }); + } + + [Test] + public async Task NonRetriableErrorDuringStreamCreation_AsyncApi() + { + await using var fixture = await MockServerFixture.StartAsync(); + + fixture.MockServer.InjectResponses(TestTableName, + [ + MockResponses.ErrorResponse(StatusCode.Unauthenticated, "Non-retriable error"), + ]); + + using var sdk = CreateDefaultSdk(fixture); + var tableProps = CreateTableProperties(); + + var options = StreamConfigurationOptions.Default with + { + MaxInflightRequests = 100, + Recovery = true, + }; + + Assert.ThrowsAsync(async () => + { + await sdk.CreateStreamWithHeadersProviderAsync(tableProps, new TestHeadersProvider(), options); + }); + } + + [Test] + public async Task RetriableErrorWithoutRecoveryDuringStreamCreation() + { + await using var fixture = await MockServerFixture.StartAsync(); + + fixture.MockServer.InjectResponses(TestTableName, + [ + MockResponses.ErrorResponse(StatusCode.Unavailable, "Retriable error"), + ]); + + using var sdk = CreateDefaultSdk(fixture); + var tableProps = CreateTableProperties(); + + var options = StreamConfigurationOptions.Default with + { + MaxInflightRequests = 100, + Recovery = false, + RecoveryTimeoutMs = 100, + RecoveryBackoffMs = 100, + }; + + var sw = Stopwatch.StartNew(); + + var ex = Assert.Throws(() => + { + sdk.CreateStreamWithHeadersProvider(tableProps, new TestHeadersProvider(), options); + }); + + Assert.That(ex, Is.Not.Null); + Assert.That(ex!.IsRetryable, Is.True); + + sw.Stop(); + + Assert.That(sw.ElapsedMilliseconds, Is.LessThan(1000), + $"Expected reasonable failure time, but took {sw.ElapsedMilliseconds}ms"); + } + + [Test] + public async Task RetriableErrorWithoutRecoveryDuringStreamCreation_AsyncApi() + { + await using var fixture = await MockServerFixture.StartAsync(); + + fixture.MockServer.InjectResponses(TestTableName, + [ + MockResponses.ErrorResponse(StatusCode.Unavailable, "Retriable error"), + ]); + + using var sdk = CreateDefaultSdk(fixture); + var tableProps = CreateTableProperties(); + + var options = StreamConfigurationOptions.Default with + { + MaxInflightRequests = 100, + Recovery = false, + RecoveryTimeoutMs = 100, + RecoveryBackoffMs = 100, + }; + + var ex = Assert.ThrowsAsync(async () => + { + await sdk.CreateStreamWithHeadersProviderAsync(tableProps, new TestHeadersProvider(), options); + }); + + Assert.That(ex, Is.Not.Null); + Assert.That(ex!.IsRetryable, Is.True); + } + + [Test] + public async Task Builder_CreatesWorkingSdk() + { + await using var fixture = await MockServerFixture.StartAsync(); + + fixture.MockServer.InjectResponses(TestTableName, + [ + MockResponses.CreateStreamResponse("test_stream_builder"), + ]); + + using var sdk = ZerobusSdk.CreateBuilder() + .Endpoint(fixture.ServerUrl) + .ApplicationName("integration-test") + .DisableTls() + .Build(); + + var tableProps = CreateTableProperties(); + + using var stream = sdk.CreateStreamWithHeadersProvider( + tableProps, + new TestHeadersProvider(), + StreamConfigurationOptions.Default with { Recovery = false }); + + Assert.That(stream, Is.Not.Null); + } +} diff --git a/dotnet/tests/Zerobus.IntegrationTests/TestHelpers.cs b/dotnet/tests/Zerobus.IntegrationTests/TestHelpers.cs new file mode 100644 index 00000000..f1ecefb3 --- /dev/null +++ b/dotnet/tests/Zerobus.IntegrationTests/TestHelpers.cs @@ -0,0 +1,205 @@ +using System.Net; +using Grpc.Core; +using Grpc.Core.Interceptors; +using Microsoft.AspNetCore.Builder; +using Microsoft.AspNetCore.Hosting; +using Microsoft.AspNetCore.Server.Kestrel.Core; +using Microsoft.Extensions.DependencyInjection; +using Microsoft.Extensions.Logging; + +namespace Databricks.Zerobus.IntegrationTests; + +/// +/// Simple headers provider for testing. +/// Returns a fixed Bearer token and table name header. +/// +public sealed class TestHeadersProvider : IHeadersProvider +{ + public IDictionary GetHeaders() + { + return new Dictionary + { + ["authorization"] = "Bearer test_token", + ["x-databricks-zerobus-table-name"] = "test_table", + }; + } +} + +/// +/// Headers provider that tracks callback invocation count. +/// Useful for asserting that the native callback bridge path was exercised. +/// +public sealed class CountingHeadersProvider : IHeadersProvider +{ + private readonly Dictionary _headers; + private int _callCount; + + public CountingHeadersProvider(IDictionary headers) + { + _headers = new Dictionary(headers, StringComparer.OrdinalIgnoreCase); + } + + public int CallCount => Volatile.Read(ref _callCount); + + public IDictionary GetHeaders() + { + Interlocked.Increment(ref _callCount); + return new Dictionary(_headers, StringComparer.OrdinalIgnoreCase); + } +} + +/// +/// Helper methods for creating mock responses and test data. +/// +public static class MockResponses +{ + public static MockResponse CreateStreamResponse(string streamId, long delayMs = 0) => new() + { + Type = MockResponseType.CreateStream, + StreamId = streamId, + DelayMs = delayMs, + }; + + public static MockResponse RecordAckResponse(long offset, long delayMs = 0) => new() + { + Type = MockResponseType.RecordAck, + AckUpToOffset = offset, + DelayMs = delayMs, + }; + + public static MockResponse ErrorResponse(StatusCode code, string message, long delayMs = 0) => new() + { + Type = MockResponseType.Error, + GrpcStatusCode = code, + GrpcMessage = message, + DelayMs = delayMs, + }; + + public static MockResponse CloseStreamSignalResponse(long durationSeconds, long delayMs = 0) => new() + { + Type = MockResponseType.CloseStreamSignal, + DurationSeconds = durationSeconds, + DelayMs = delayMs, + }; +} + +/// +/// Creates a simple test protobuf descriptor. +/// This is a pre-serialized FileDescriptorSet for: +/// message TestMessage { int64 id = 1; string message = 2; } +/// Matches the Go CreateTestDescriptorProto(). +/// +public static class TestDescriptor +{ + public static byte[] CreateTestDescriptorProto() => + [ + 0x0a, 0x45, 0x0a, 0x0b, 0x74, 0x65, 0x73, 0x74, 0x2e, 0x70, 0x72, 0x6f, + 0x74, 0x6f, 0x22, 0x36, 0x0a, 0x0b, 0x54, 0x65, 0x73, 0x74, 0x4d, 0x65, + 0x73, 0x73, 0x61, 0x67, 0x65, 0x12, 0x0e, 0x0a, 0x02, 0x69, 0x64, 0x18, + 0x01, 0x20, 0x01, 0x28, 0x03, 0x52, 0x02, 0x69, 0x64, 0x12, 0x17, 0x0a, + 0x07, 0x6d, 0x65, 0x73, 0x73, 0x61, 0x67, 0x65, 0x18, 0x02, 0x20, 0x01, + 0x28, 0x09, 0x52, 0x07, 0x6d, 0x65, 0x73, 0x73, 0x61, 0x67, 0x65, 0x62, + 0x06, 0x70, 0x72, 0x6f, 0x74, 0x6f, 0x33, + ]; +} + +/// +/// Manages the lifecycle of the mock gRPC server for integration tests. +/// Starts an ASP.NET Core Kestrel server with the gRPC mock service. +/// +public sealed class MockServerFixture : IAsyncDisposable +{ + public MockZerobusServer MockServer { get; } + public string ServerUrl { get; } + + private readonly WebApplication _app; + + private MockServerFixture(MockZerobusServer mockServer, string serverUrl, WebApplication app) + { + MockServer = mockServer; + ServerUrl = serverUrl; + _app = app; + } + + /// + /// Starts a new mock gRPC server on a random available port. + /// + public static async Task StartAsync() + { + var mockServer = new MockZerobusServer(); + + var builder = WebApplication.CreateBuilder(); + builder.Logging.ClearProviders(); // Suppress noisy ASP.NET Core logs in test output. + builder.Services.AddGrpc(options => + { + options.Interceptors.Add(); + }); + builder.Services.AddSingleton(mockServer); + builder.WebHost.ConfigureKestrel(options => + { + // Listen on 127.0.0.1 with a random port and HTTP/2 (required for gRPC). + options.Listen(IPAddress.Loopback, 0, listenOptions => + { + listenOptions.Protocols = HttpProtocols.Http2; + }); + }); + + var app = builder.Build(); + app.MapGrpcService(); + + await app.StartAsync(); + + // Give the server a moment to be fully ready for connections. + await Task.Delay(100); + + // Get the actual bound address from the server features. + var addresses = app.Urls; + var serverUrl = addresses.FirstOrDefault() + ?? throw new InvalidOperationException("No server URL was bound."); + + return new MockServerFixture(mockServer, serverUrl, app); + } + + public async ValueTask DisposeAsync() + { + await _app.StopAsync(); + await _app.DisposeAsync(); + } +} + +/// +/// gRPC server interceptor that catches unhandled exceptions and converts them +/// to proper responses. Without this, ASP.NET Core +/// wraps unhandled exceptions as StatusCode.Unknown "Exception was thrown by handler." +/// which confuses the Rust SDK's error classification. +/// +internal sealed class ExceptionInterceptor : Interceptor +{ + public override async Task DuplexStreamingServerHandler( + IAsyncStreamReader requestStream, + IServerStreamWriter responseStream, + ServerCallContext context, + DuplexStreamingServerMethod continuation) + { + try + { + await continuation(requestStream, responseStream, context); + } + catch (RpcException) + { + throw; + } + catch (OperationCanceledException) + { + throw new RpcException(new Status(StatusCode.Cancelled, "Client disconnected")); + } + catch (IOException) + { + throw new RpcException(new Status(StatusCode.Cancelled, "Client disconnected")); + } + catch (InvalidOperationException) + { + throw new RpcException(new Status(StatusCode.Cancelled, "Client disconnected")); + } + } +} diff --git a/dotnet/tests/Zerobus.IntegrationTests/Zerobus.IntegrationTests.csproj b/dotnet/tests/Zerobus.IntegrationTests/Zerobus.IntegrationTests.csproj new file mode 100644 index 00000000..52136f3a --- /dev/null +++ b/dotnet/tests/Zerobus.IntegrationTests/Zerobus.IntegrationTests.csproj @@ -0,0 +1,32 @@ + + + + net8.0;net10.0 + Databricks.Zerobus.IntegrationTests + Databricks.Zerobus.IntegrationTests + false + true + + + + + + + + + + + + + + + + + + + + + + + + diff --git a/dotnet/tests/Zerobus.Tests/NativeInteropTests.cs b/dotnet/tests/Zerobus.Tests/NativeInteropTests.cs new file mode 100644 index 00000000..b2eccb19 --- /dev/null +++ b/dotnet/tests/Zerobus.Tests/NativeInteropTests.cs @@ -0,0 +1,147 @@ +using Databricks.Zerobus; +using Databricks.Zerobus.Native; +using NUnit.Framework; + +namespace Databricks.Zerobus.Tests; + +[TestFixture] +public class NativeInteropTests +{ + [Test] + public void ConvertConfig_CustomOptions_AppliesValues() + { + var options = StreamConfigurationOptions.Default with + { + MaxInflightRequests = 42, + Recovery = false, + RecoveryTimeoutMs = 100, + RecoveryBackoffMs = 200, + RecoveryRetries = 3, + ServerLackOfAckTimeoutMs = 1000, + FlushTimeoutMs = 2000, + RecordType = RecordType.Json, + StreamPausedMaxWaitTimeMs = 500, + }; + + var native = NativeInterop.ConvertConfig(options); + + Assert.That(native.MaxInflightRequests, Is.EqualTo((nuint)42)); + Assert.That(native.Recovery, Is.False); + Assert.That(native.RecoveryTimeoutMs, Is.EqualTo(100UL)); + Assert.That(native.RecoveryBackoffMs, Is.EqualTo(200UL)); + Assert.That(native.RecoveryRetries, Is.EqualTo(3U)); + Assert.That(native.ServerLackOfAckTimeoutMs, Is.EqualTo(1000UL)); + Assert.That(native.FlushTimeoutMs, Is.EqualTo(2000UL)); + Assert.That(native.RecordType, Is.EqualTo((int)RecordType.Json)); + Assert.That(native.StreamPausedMaxWaitTimeMs, Is.EqualTo(500UL)); + Assert.That(native.HasStreamPausedMaxWaitTimeMs, Is.True); + } + + [Test] + public void ConvertConfig_NullStreamPausedWait_SetsFlagToFalse() + { + var options = StreamConfigurationOptions.Default; + + var native = NativeInterop.ConvertConfig(options); + + Assert.That(native.HasStreamPausedMaxWaitTimeMs, Is.False); + } + + [Test] + public void ConvertConfig_ExplicitZeroValues_ArePreserved() + { + var options = StreamConfigurationOptions.Default with + { + MaxInflightRequests = 0, + RecoveryTimeoutMs = 0, + RecoveryBackoffMs = 0, + RecoveryRetries = 0, + ServerLackOfAckTimeoutMs = 0, + FlushTimeoutMs = 0, + }; + + var native = NativeInterop.ConvertConfig(options); + + Assert.That(native.MaxInflightRequests, Is.EqualTo((nuint)0)); + Assert.That(native.RecoveryTimeoutMs, Is.EqualTo(0UL)); + Assert.That(native.RecoveryBackoffMs, Is.EqualTo(0UL)); + Assert.That(native.RecoveryRetries, Is.EqualTo(0U)); + Assert.That(native.ServerLackOfAckTimeoutMs, Is.EqualTo(0UL)); + Assert.That(native.FlushTimeoutMs, Is.EqualTo(0UL)); + } + + [Test] + public void ConvertConfig_NullNumericValues_FallBackToDefaults() + { + var options = StreamConfigurationOptions.Default with + { + MaxInflightRequests = null, + RecoveryTimeoutMs = null, + RecoveryBackoffMs = null, + RecoveryRetries = null, + ServerLackOfAckTimeoutMs = null, + FlushTimeoutMs = null, + }; + + var nativeDefaults = NativeMethods.GetDefaultConfig(); + var native = NativeInterop.ConvertConfig(options); + + Assert.That(native.MaxInflightRequests, Is.EqualTo(nativeDefaults.MaxInflightRequests)); + Assert.That(native.RecoveryTimeoutMs, Is.EqualTo(nativeDefaults.RecoveryTimeoutMs)); + Assert.That(native.RecoveryBackoffMs, Is.EqualTo(nativeDefaults.RecoveryBackoffMs)); + Assert.That(native.RecoveryRetries, Is.EqualTo(nativeDefaults.RecoveryRetries)); + Assert.That(native.ServerLackOfAckTimeoutMs, Is.EqualTo(nativeDefaults.ServerLackOfAckTimeoutMs)); + Assert.That(native.FlushTimeoutMs, Is.EqualTo(nativeDefaults.FlushTimeoutMs)); + } + + [Test] + public void ConvertConfig_DefaultManagedOptions_MatchNativeDefaults() + { + var nativeDefaults = NativeMethods.GetDefaultConfig(); + var converted = NativeInterop.ConvertConfig(StreamConfigurationOptions.Default); + + Assert.That(converted.MaxInflightRequests, Is.EqualTo(nativeDefaults.MaxInflightRequests)); + Assert.That(converted.Recovery, Is.EqualTo(nativeDefaults.Recovery)); + Assert.That(converted.RecoveryTimeoutMs, Is.EqualTo(nativeDefaults.RecoveryTimeoutMs)); + Assert.That(converted.RecoveryBackoffMs, Is.EqualTo(nativeDefaults.RecoveryBackoffMs)); + Assert.That(converted.RecoveryRetries, Is.EqualTo(nativeDefaults.RecoveryRetries)); + Assert.That(converted.ServerLackOfAckTimeoutMs, Is.EqualTo(nativeDefaults.ServerLackOfAckTimeoutMs)); + Assert.That(converted.FlushTimeoutMs, Is.EqualTo(nativeDefaults.FlushTimeoutMs)); + Assert.That(converted.RecordType, Is.EqualTo(nativeDefaults.RecordType)); + Assert.That(converted.HasStreamPausedMaxWaitTimeMs, Is.EqualTo(nativeDefaults.HasStreamPausedMaxWaitTimeMs)); + Assert.That(converted.StreamPausedMaxWaitTimeMs, Is.EqualTo(nativeDefaults.StreamPausedMaxWaitTimeMs)); + } + + [Test] + public void ToException_SuccessResult_ReturnsNull() + { + var result = new CResult { Success = true, ErrorMessage = IntPtr.Zero, IsRetryable = false }; + + var ex = NativeInterop.ToException(ref result); + + Assert.That(ex, Is.Null); + } + + [Test] + public void ToException_FailureWithNullMessage_ReturnsUnknownError() + { + var result = new CResult { Success = false, ErrorMessage = IntPtr.Zero, IsRetryable = false }; + + var ex = NativeInterop.ToException(ref result); + + Assert.That(ex, Is.Not.Null); + Assert.That(ex!.RawMessage, Is.EqualTo("unknown error")); + Assert.That(ex.IsRetryable, Is.False); + } + + [Test] + public void ToException_RetryableFailure_SetsIsRetryable() + { + var result = new CResult { Success = false, ErrorMessage = IntPtr.Zero, IsRetryable = true }; + + var ex = NativeInterop.ToException(ref result); + + Assert.That(ex, Is.Not.Null); + Assert.That(ex!.IsRetryable, Is.True); + } +} diff --git a/dotnet/tests/Zerobus.Tests/RecordTypeTests.cs b/dotnet/tests/Zerobus.Tests/RecordTypeTests.cs new file mode 100644 index 00000000..39f2e838 --- /dev/null +++ b/dotnet/tests/Zerobus.Tests/RecordTypeTests.cs @@ -0,0 +1,17 @@ +using Databricks.Zerobus; +using NUnit.Framework; + +namespace Databricks.Zerobus.Tests; + +[TestFixture] +public class RecordTypeTests +{ + [Test] + public void EnumValues_MatchCBindings() + { + // These values must match the C enum used by the Rust FFI layer. + Assert.That((int)RecordType.Unspecified, Is.EqualTo(0)); + Assert.That((int)RecordType.Proto, Is.EqualTo(1)); + Assert.That((int)RecordType.Json, Is.EqualTo(2)); + } +} diff --git a/dotnet/tests/Zerobus.Tests/StreamConfigurationOptionsTests.cs b/dotnet/tests/Zerobus.Tests/StreamConfigurationOptionsTests.cs new file mode 100644 index 00000000..1d0f19d8 --- /dev/null +++ b/dotnet/tests/Zerobus.Tests/StreamConfigurationOptionsTests.cs @@ -0,0 +1,103 @@ +using Databricks.Zerobus; +using NUnit.Framework; + +namespace Databricks.Zerobus.Tests; + +[TestFixture] +public class StreamConfigurationOptionsTests +{ + [Test] + public void Default_ReturnsExpectedValues() + { + var options = StreamConfigurationOptions.Default; + + Assert.That(options.MaxInflightRequests, Is.EqualTo(1_000_000UL)); + Assert.That(options.Recovery, Is.True); + Assert.That(options.RecoveryTimeoutMs, Is.EqualTo(15_000UL)); + Assert.That(options.RecoveryBackoffMs, Is.EqualTo(2_000UL)); + Assert.That(options.RecoveryRetries, Is.EqualTo(4U)); + Assert.That(options.ServerLackOfAckTimeoutMs, Is.EqualTo(60_000UL)); + Assert.That(options.FlushTimeoutMs, Is.EqualTo(300_000UL)); + Assert.That(options.RecordType, Is.EqualTo(RecordType.Proto)); + Assert.That(options.StreamPausedMaxWaitTimeMs, Is.Null); + } + + [Test] + public void WithExpression_OverridesSpecificFields() + { + var options = StreamConfigurationOptions.Default with + { + MaxInflightRequests = 50_000, + RecordType = RecordType.Json, + }; + + Assert.That(options.MaxInflightRequests, Is.EqualTo(50_000UL)); + Assert.That(options.RecordType, Is.EqualTo(RecordType.Json)); + + // Other fields remain default. + Assert.That(options.Recovery, Is.True); + Assert.That(options.RecoveryTimeoutMs, Is.EqualTo(15_000UL)); + } + + [Test] + public void WithExpression_CanSetStreamPausedMaxWaitTime() + { + var options = StreamConfigurationOptions.Default with + { + StreamPausedMaxWaitTimeMs = 5_000, + }; + + Assert.That(options.StreamPausedMaxWaitTimeMs, Is.EqualTo(5_000UL)); + } + + [Test] + public void WithExpression_CanSetStreamPausedMaxWaitTimeToZero() + { + var options = StreamConfigurationOptions.Default with + { + StreamPausedMaxWaitTimeMs = 0, + }; + + Assert.That(options.StreamPausedMaxWaitTimeMs, Is.Not.Null); + Assert.That(options.StreamPausedMaxWaitTimeMs, Is.EqualTo(0UL)); + } + + [Test] + public void WithExpression_CanSetNumericFieldsToNull() + { + var options = StreamConfigurationOptions.Default with + { + MaxInflightRequests = null, + RecoveryTimeoutMs = null, + RecoveryBackoffMs = null, + RecoveryRetries = null, + ServerLackOfAckTimeoutMs = null, + FlushTimeoutMs = null, + }; + + Assert.That(options.MaxInflightRequests, Is.Null); + Assert.That(options.RecoveryTimeoutMs, Is.Null); + Assert.That(options.RecoveryBackoffMs, Is.Null); + Assert.That(options.RecoveryRetries, Is.Null); + Assert.That(options.ServerLackOfAckTimeoutMs, Is.Null); + Assert.That(options.FlushTimeoutMs, Is.Null); + } + + [Test] + public void Record_SupportsEquality() + { + var a = StreamConfigurationOptions.Default; + var b = StreamConfigurationOptions.Default; + + Assert.That(a, Is.EqualTo(b)); + } + + [Test] + public void Record_DifferentValues_AreNotEqual() + { + var a = StreamConfigurationOptions.Default; + var b = StreamConfigurationOptions.Default with { RecordType = RecordType.Json }; + + Assert.That(a, Is.Not.EqualTo(b)); + } +} diff --git a/dotnet/tests/Zerobus.Tests/TablePropertiesTests.cs b/dotnet/tests/Zerobus.Tests/TablePropertiesTests.cs new file mode 100644 index 00000000..9c05e999 --- /dev/null +++ b/dotnet/tests/Zerobus.Tests/TablePropertiesTests.cs @@ -0,0 +1,46 @@ +using Databricks.Zerobus; +using NUnit.Framework; + +namespace Databricks.Zerobus.Tests; + +[TestFixture] +public class TablePropertiesTests +{ + [Test] + public void Constructor_JsonTable_DescriptorProtoIsNull() + { + var props = new TableProperties("catalog.schema.table"); + + Assert.That(props.TableName, Is.EqualTo("catalog.schema.table")); + Assert.That(props.DescriptorProto, Is.Null); + } + + [Test] + public void Constructor_ProtoTable_HasDescriptorProto() + { + byte[] descriptor = [0x0A, 0x03, 0x66, 0x6F, 0x6F]; + var props = new TableProperties("catalog.schema.table", descriptor); + + Assert.That(props.TableName, Is.EqualTo("catalog.schema.table")); + Assert.That(props.DescriptorProto, Is.Not.Null); + Assert.That(props.DescriptorProto, Is.EqualTo(descriptor)); + } + + [Test] + public void Record_SupportsEquality() + { + var a = new TableProperties("catalog.schema.table"); + var b = new TableProperties("catalog.schema.table"); + + Assert.That(a, Is.EqualTo(b)); + } + + [Test] + public void Record_DifferentTableNames_AreNotEqual() + { + var a = new TableProperties("catalog.schema.table1"); + var b = new TableProperties("catalog.schema.table2"); + + Assert.That(a, Is.Not.EqualTo(b)); + } +} diff --git a/dotnet/tests/Zerobus.Tests/TypedZerobusStreamTests.cs b/dotnet/tests/Zerobus.Tests/TypedZerobusStreamTests.cs new file mode 100644 index 00000000..ff3d39fb --- /dev/null +++ b/dotnet/tests/Zerobus.Tests/TypedZerobusStreamTests.cs @@ -0,0 +1,52 @@ +using System.Reflection; +using Databricks.Zerobus; +using NUnit.Framework; + +namespace Databricks.Zerobus.Tests; + +[TestFixture] +public class TypedZerobusStreamTests +{ + [Test] + public void JsonStream_ExposesOnlyJsonIngestOverloads() + { + var singleRecord = typeof(JsonZerobusStream) + .GetMethods(BindingFlags.Instance | BindingFlags.Public) + .Where(method => method.Name == nameof(JsonZerobusStream.IngestRecord)) + .ToArray(); + + var batch = typeof(JsonZerobusStream) + .GetMethods(BindingFlags.Instance | BindingFlags.Public) + .Where(method => method.Name == nameof(JsonZerobusStream.IngestRecords)) + .ToArray(); + + Assert.That(singleRecord, Has.Length.EqualTo(1)); + Assert.That(singleRecord[0].GetParameters().Select(parameter => parameter.ParameterType), + Is.EqualTo(new[] { typeof(string) })); + + Assert.That(batch, Has.Length.EqualTo(1)); + Assert.That(batch[0].GetParameters().Select(parameter => parameter.ParameterType), + Is.EqualTo(new[] { typeof(string[]) })); + } + + [Test] + public void ProtoStream_ExposesOnlyProtoIngestOverloads() + { + var singleRecord = typeof(ProtoZerobusStream) + .GetMethods(BindingFlags.Instance | BindingFlags.Public) + .Where(method => method.Name == nameof(ProtoZerobusStream.IngestRecord)) + .ToArray(); + + var batch = typeof(ProtoZerobusStream) + .GetMethods(BindingFlags.Instance | BindingFlags.Public) + .Where(method => method.Name == nameof(ProtoZerobusStream.IngestRecords)) + .ToArray(); + + Assert.That(singleRecord.Select(method => method.GetParameters().Single().ParameterType), + Is.EquivalentTo(new[] { typeof(byte[]), typeof(ReadOnlySpan) })); + + Assert.That(batch, Has.Length.EqualTo(1)); + Assert.That(batch[0].GetParameters().Select(parameter => parameter.ParameterType), + Is.EqualTo(new[] { typeof(byte[][]) })); + } +} \ No newline at end of file diff --git a/dotnet/tests/Zerobus.Tests/Zerobus.Tests.csproj b/dotnet/tests/Zerobus.Tests/Zerobus.Tests.csproj new file mode 100644 index 00000000..0c005ea9 --- /dev/null +++ b/dotnet/tests/Zerobus.Tests/Zerobus.Tests.csproj @@ -0,0 +1,22 @@ + + + + net8.0;net10.0 + Databricks.Zerobus.Tests + Databricks.Zerobus.Tests + false + true + + + + + + + + + + + + + + diff --git a/dotnet/tests/Zerobus.Tests/ZerobusExceptionTests.cs b/dotnet/tests/Zerobus.Tests/ZerobusExceptionTests.cs new file mode 100644 index 00000000..0e249e04 --- /dev/null +++ b/dotnet/tests/Zerobus.Tests/ZerobusExceptionTests.cs @@ -0,0 +1,60 @@ +using Databricks.Zerobus; +using NUnit.Framework; + +namespace Databricks.Zerobus.Tests; + +[TestFixture] +public class ZerobusExceptionTests +{ + [Test] + public void Constructor_RetryableError_FormatsMessageCorrectly() + { + var ex = new ZerobusException("connection lost", isRetryable: true); + + Assert.That(ex.IsRetryable, Is.True); + Assert.That(ex.RawMessage, Is.EqualTo("connection lost")); + Assert.That(ex.Message, Does.Contain("retryable")); + Assert.That(ex.Message, Does.Contain("connection lost")); + } + + [Test] + public void Constructor_NonRetryableError_FormatsMessageCorrectly() + { + var ex = new ZerobusException("invalid table name", isRetryable: false); + + Assert.That(ex.IsRetryable, Is.False); + Assert.That(ex.RawMessage, Is.EqualTo("invalid table name")); + Assert.That(ex.Message, Does.Not.Contain("retryable")); + Assert.That(ex.Message, Does.Contain("invalid table name")); + } + + [Test] + public void Exception_IsStandardException() + { + var ex = new ZerobusException("test", isRetryable: false); + + Assert.That(ex, Is.InstanceOf()); + } + + [Test] + public void RetryableException_CanBeCaughtWithPattern() + { + var ex = new ZerobusException("timeout", isRetryable: true); + + try + { + throw ex; + } + catch (ZerobusException caught) when (caught.IsRetryable) + { + Assert.That(caught.IsRetryable, Is.True); + return; + } + catch + { + Assert.Fail("Should have been caught by the retryable catch clause"); + } + + Assert.Fail("Exception was not thrown"); + } +} diff --git a/dotnet/tests/Zerobus.Tests/ZerobusSdkBuilderTests.cs b/dotnet/tests/Zerobus.Tests/ZerobusSdkBuilderTests.cs new file mode 100644 index 00000000..b86f16b8 --- /dev/null +++ b/dotnet/tests/Zerobus.Tests/ZerobusSdkBuilderTests.cs @@ -0,0 +1,213 @@ +using Databricks.Zerobus; +using NUnit.Framework; + +namespace Databricks.Zerobus.Tests; + +[TestFixture] +public class ZerobusSdkBuilderTests +{ + // ------------------------------------------------------------------------- + // CreateBuilder / factory + // ------------------------------------------------------------------------- + + [Test] + public void CreateBuilder_ReturnsNonNullBuilder() + { + using var builder = ZerobusSdk.CreateBuilder(); + + Assert.That(builder, Is.Not.Null); + } + + // ------------------------------------------------------------------------- + // Fluent return-value contract (no native call needed) + // ------------------------------------------------------------------------- + + [Test] + public void Endpoint_ReturnsSameBuilderInstance() + { + using var builder = ZerobusSdk.CreateBuilder(); + + var returned = builder.Endpoint("https://zerobus.databricks.com"); + + Assert.That(returned, Is.SameAs(builder)); + } + + [Test] + public void UnityCatalogUrl_ReturnsSameBuilderInstance() + { + using var builder = ZerobusSdk.CreateBuilder(); + + var returned = builder.UnityCatalogUrl("https://workspace.databricks.com"); + + Assert.That(returned, Is.SameAs(builder)); + } + + [Test] + public void ApplicationName_ReturnsSameBuilderInstance() + { + using var builder = ZerobusSdk.CreateBuilder(); + + var returned = builder.ApplicationName("my-app"); + + Assert.That(returned, Is.SameAs(builder)); + } + + [Test] + public void SdkIdentifier_ReturnsSameBuilderInstance() + { + using var builder = ZerobusSdk.CreateBuilder(); + + var returned = builder.SdkIdentifier("zerobus-sdk-dotnet/1.0.0"); + + Assert.That(returned, Is.SameAs(builder)); + } + + [Test] + public void DisableTls_ReturnsSameBuilderInstance() + { + using var builder = ZerobusSdk.CreateBuilder(); + + var returned = builder.DisableTls(); + + Assert.That(returned, Is.SameAs(builder)); + } + + // ------------------------------------------------------------------------- + // Null argument guards + // ------------------------------------------------------------------------- + + [Test] + public void Endpoint_NullValue_ThrowsArgumentNullException() + { + using var builder = ZerobusSdk.CreateBuilder(); + + Assert.Throws(() => builder.Endpoint(null!)); + } + + [Test] + public void UnityCatalogUrl_NullValue_ThrowsArgumentNullException() + { + using var builder = ZerobusSdk.CreateBuilder(); + + Assert.Throws(() => builder.UnityCatalogUrl(null!)); + } + + [Test] + public void ApplicationName_NullValue_ThrowsArgumentNullException() + { + using var builder = ZerobusSdk.CreateBuilder(); + + Assert.Throws(() => builder.ApplicationName(null!)); + } + + [Test] + public void SdkIdentifier_NullValue_ThrowsArgumentNullException() + { + using var builder = ZerobusSdk.CreateBuilder(); + + Assert.Throws(() => builder.SdkIdentifier(null!)); + } + + // ------------------------------------------------------------------------- + // ObjectDisposedException after Dispose() + // ------------------------------------------------------------------------- + + [Test] + public void Endpoint_AfterDispose_ThrowsObjectDisposedException() + { + var builder = ZerobusSdk.CreateBuilder(); + builder.Dispose(); + + Assert.Throws(() => builder.Endpoint("https://zerobus.databricks.com")); + } + + [Test] + public void UnityCatalogUrl_AfterDispose_ThrowsObjectDisposedException() + { + var builder = ZerobusSdk.CreateBuilder(); + builder.Dispose(); + + Assert.Throws(() => builder.UnityCatalogUrl("https://workspace.databricks.com")); + } + + [Test] + public void ApplicationName_AfterDispose_ThrowsObjectDisposedException() + { + var builder = ZerobusSdk.CreateBuilder(); + builder.Dispose(); + + Assert.Throws(() => builder.ApplicationName("my-app")); + } + + [Test] + public void SdkIdentifier_AfterDispose_ThrowsObjectDisposedException() + { + var builder = ZerobusSdk.CreateBuilder(); + builder.Dispose(); + + Assert.Throws(() => builder.SdkIdentifier("zerobus-sdk-dotnet/1.0.0")); + } + + [Test] + public void DisableTls_AfterDispose_ThrowsObjectDisposedException() + { + var builder = ZerobusSdk.CreateBuilder(); + builder.Dispose(); + + Assert.Throws(() => builder.DisableTls()); + } + + [Test] + public void Build_AfterDispose_ThrowsObjectDisposedException() + { + var builder = ZerobusSdk.CreateBuilder(); + builder.Dispose(); + + Assert.Throws(() => builder.Build()); + } + + // ------------------------------------------------------------------------- + // Single-use contract: Build() consumes the builder + // ------------------------------------------------------------------------- + + [Test] + public void Build_CalledTwice_ThrowsObjectDisposedExceptionOnSecondCall() + { + var builder = ZerobusSdk.CreateBuilder() + .Endpoint("https://zerobus.databricks.com") + .UnityCatalogUrl("https://workspace.databricks.com"); + + using var sdk = builder.Build(); + + Assert.Throws(() => builder.Build()); + } + + [Test] + public void Endpoint_AfterBuild_ThrowsObjectDisposedException() + { + var builder = ZerobusSdk.CreateBuilder() + .Endpoint("https://zerobus.databricks.com") + .UnityCatalogUrl("https://workspace.databricks.com"); + + using var sdk = builder.Build(); + + Assert.Throws(() => builder.Endpoint("https://other.databricks.com")); + } + + // ------------------------------------------------------------------------- + // Dispose is idempotent + // ------------------------------------------------------------------------- + + [Test] + public void Dispose_CalledMultipleTimes_DoesNotThrow() + { + var builder = ZerobusSdk.CreateBuilder(); + + Assert.DoesNotThrow(() => + { + builder.Dispose(); + builder.Dispose(); + builder.Dispose(); + }); + } +} diff --git a/rust/ffi/NEXT_CHANGELOG.md b/rust/ffi/NEXT_CHANGELOG.md index 473e4c49..e05b9cc7 100644 --- a/rust/ffi/NEXT_CHANGELOG.md +++ b/rust/ffi/NEXT_CHANGELOG.md @@ -6,6 +6,8 @@ ### New Features and Improvements +- Add callback-based async overloads for all previously blocking stream operations: stream creation (`zerobus_sdk_create_stream_async`, `zerobus_sdk_create_stream_with_headers_provider_async`), stream recreation (`zerobus_sdk_recreate_stream_async`), offset-returning ingest calls (`zerobus_stream_ingest_proto_record_async`, `zerobus_stream_ingest_json_record_async`, `zerobus_stream_ingest_proto_records_async`, `zerobus_stream_ingest_json_records_async`), completion methods (`zerobus_stream_wait_for_offset_async`, `zerobus_stream_flush_async`, `zerobus_stream_close_async`), and unacked-record retrieval (`zerobus_stream_get_unacked_records_async`). These APIs return immediately after validation/scheduling and complete via callbacks; caller-owned string/descriptor/config inputs are copied before return, and SDK/stream handles must remain valid until callback completion. + ### Bug Fixes ### Documentation @@ -19,3 +21,5 @@ ### Deprecations ### API Changes + +- Add `CreateStreamAsyncCallback`, `OffsetAsyncCallback`, `BoolAsyncCallback`, and `RecordArrayAsyncCallback` plus the full async stream API set (`*_async` overloads for create/recreate/ingest/wait/flush/get_unacked_records/close) to `zerobus.h`. Callback `const CResult *` values are valid only for the duration of each callback; any error text must be copied during the call. diff --git a/rust/ffi/README.md b/rust/ffi/README.md index de176f03..328c9bc3 100644 --- a/rust/ffi/README.md +++ b/rust/ffi/README.md @@ -68,6 +68,20 @@ private static extern IntPtr zerobus_sdk_new(string endpoint, string ucUrl, ref // Link with -lzerobus_ffi ``` +### Async stream creation callback + +For callers that do not want to block a thread in the synchronous stream +creation functions, use `zerobus_sdk_create_stream_async` or +`zerobus_sdk_create_stream_with_headers_provider_async`. Each returns once the +request has been validated and queued on the Rust runtime, then invokes a +one-shot `CreateStreamAsyncCallback` with either the created +`CZerobusStream*` or a null stream plus a failure `CResult`. + +The callback receives a `const CResult*` that is valid only for the duration of +the call, so copy `error_message` if you need to keep it. All string, +descriptor, and config inputs are copied before the async create call returns, +but the `CZerobusSdk*` itself must remain valid until the callback fires. + ### Dynamic protobuf from a Unity Catalog schema (pure C) Build a protobuf descriptor and encode records straight from Unity Catalog diff --git a/rust/ffi/src/common.rs b/rust/ffi/src/common.rs index e2d30bd8..8b645419 100644 --- a/rust/ffi/src/common.rs +++ b/rust/ffi/src/common.rs @@ -167,6 +167,10 @@ pub struct CStreamConfigurationOptions { pub ack_user_data: *mut std::ffi::c_void, } +// Safety: this POD struct only carries scalars, function pointers, and an +// opaque user_data pointer whose synchronization remains the caller's contract. +unsafe impl Send for CStreamConfigurationOptions {} + // Helper to convert C string to Rust String pub(crate) unsafe fn c_str_to_string(c_str: *const c_char) -> Result { if c_str.is_null() { @@ -206,6 +210,52 @@ pub struct CHeaders { /// The caller is responsible for freeing the returned CHeaders using zerobus_free_headers pub type HeadersProviderCallback = extern "C" fn(user_data: *mut std::ffi::c_void) -> CHeaders; +/// Function pointer type for async stream creation completion. +/// +/// `stream` is non-null on success and null on failure. `result` points to a +/// `CResult` valid only for the duration of the call; copy any error text during +/// the callback if you need to retain it. +/// +/// Invoked from a background task, so it must be thread-safe and must not +/// unwind across the FFI boundary. +pub type CreateStreamAsyncCallback = extern "C" fn( + stream: *mut CZerobusStream, + result: *const CResult, + user_data: *mut std::ffi::c_void, +); + +/// Function pointer type for async offset-returning operations. +/// +/// `result` points to a `CResult` valid only for the duration of the call; copy +/// any error text during the callback if you need to retain it. +pub type OffsetAsyncCallback = extern "C" fn( + offset: i64, + result: *const CResult, + user_data: *mut std::ffi::c_void, +); + +/// Function pointer type for async bool-returning operations. +/// +/// `result` points to a `CResult` valid only for the duration of the call; copy +/// any error text during the callback if you need to retain it. +pub type BoolAsyncCallback = extern "C" fn( + value: bool, + result: *const CResult, + user_data: *mut std::ffi::c_void, +); + +/// Function pointer type for async `CRecordArray`-returning operations. +/// +/// On success, ownership of `records` transfers to the callback recipient, who +/// must free it with `zerobus_free_record_array`. `result` points to a `CResult` +/// valid only for the duration of the call; copy any error text during the +/// callback if you need to retain it. +pub type RecordArrayAsyncCallback = extern "C" fn( + records: CRecordArray, + result: *const CResult, + user_data: *mut std::ffi::c_void, +); + /// Rust struct that wraps a Go callback and implements HeadersProvider pub(crate) struct CallbackHeadersProvider { callback: HeadersProviderCallback, diff --git a/rust/ffi/src/stream.rs b/rust/ffi/src/stream.rs index 96a516d2..0f68ecd4 100644 --- a/rust/ffi/src/stream.rs +++ b/rust/ffi/src/stream.rs @@ -2,8 +2,11 @@ use crate::common::*; use databricks_zerobus_ingest_sdk::databricks::zerobus::RecordType; -use databricks_zerobus_ingest_sdk::{EncodedRecord, HeadersProvider, StreamBuilder, ZerobusStream}; +use databricks_zerobus_ingest_sdk::{ + EncodedRecord, HeadersProvider, StreamBuilder, ZerobusError, ZerobusStream, +}; use prost::Message; +use std::mem::ManuallyDrop; use std::ffi::CString; use std::os::raw::c_char; use std::ptr; @@ -46,6 +49,227 @@ fn apply_c_stream_options<'a>( } } +#[derive(Clone, Copy)] +struct SendPtr(*mut T); + +impl SendPtr { + const fn new(ptr: *mut T) -> Self { + Self(ptr) + } + + const fn get(self) -> *mut T { + self.0 + } +} + +// Safety: this wrapper carries an opaque FFI pointer across task boundaries. +// The pointee's lifetime and synchronization remain the caller's contract. +unsafe impl Send for SendPtr {} + +enum StreamCreateAuth { + OAuth { + client_id: String, + client_secret: String, + }, + HeadersProvider { + headers_callback: HeadersProviderCallback, + user_data: SendPtr, + }, +} + +async fn build_stream_from_parts( + sdk_ref: &databricks_zerobus_ingest_sdk::ZerobusSdk, + table_name: String, + descriptor_proto: Option, + auth: StreamCreateAuth, + options: Option, +) -> Result<*mut CZerobusStream, ZerobusError> { + let record_type = options + .as_ref() + .map(|c| c_record_type(c.record_type)) + .unwrap_or(RecordType::Proto); + + let base = match auth { + StreamCreateAuth::OAuth { + client_id, + client_secret, + } => sdk_ref + .stream_builder() + .table(table_name) + .oauth(client_id, client_secret), + StreamCreateAuth::HeadersProvider { + headers_callback, + user_data, + } => { + let headers_provider: Arc = + Arc::new(CallbackHeadersProvider::new(headers_callback, user_data.get())); + sdk_ref + .stream_builder() + .table(table_name) + .headers_provider(headers_provider) + } + }; + + let mut builder = match record_type { + RecordType::Proto => { + let desc = descriptor_proto.ok_or_else(|| { + ZerobusError::InvalidArgument( + "Proto descriptor is required for Proto record type".to_string(), + ) + })?; + base.compiled_proto(desc) + } + RecordType::Json => base.json(), + RecordType::Unspecified => { + return Err(ZerobusError::InvalidArgument( + "Record type is not specified".to_string(), + )) + } + }; + + if let Some(c) = options.as_ref() { + builder = apply_c_stream_options(builder, c); + } + + let stream = builder.build().await?; + Ok(Arc::into_raw(Arc::new(stream)) as *mut CZerobusStream) +} + +fn invoke_create_stream_async_callback( + callback: CreateStreamAsyncCallback, + stream: *mut CZerobusStream, + result: CResult, + user_data: *mut std::ffi::c_void, +) { + let callback_result = result; + let callback_result_ptr = &callback_result as *const CResult; + + if std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| { + callback(stream, callback_result_ptr, user_data) + })) + .is_err() + { + tracing::error!("async create_stream callback panicked; contained at FFI boundary"); + } + + if !callback_result.error_message.is_null() { + unsafe { + let _ = CString::from_raw(callback_result.error_message); + } + } +} + +fn invoke_offset_async_callback( + callback: OffsetAsyncCallback, + offset: i64, + result: CResult, + user_data: *mut std::ffi::c_void, +) { + let callback_result = result; + let callback_result_ptr = &callback_result as *const CResult; + + if std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| { + callback(offset, callback_result_ptr, user_data) + })) + .is_err() + { + tracing::error!(offset, "async offset callback panicked; contained at FFI boundary"); + } + + if !callback_result.error_message.is_null() { + unsafe { + let _ = CString::from_raw(callback_result.error_message); + } + } +} + +fn invoke_bool_async_callback( + callback: BoolAsyncCallback, + value: bool, + result: CResult, + user_data: *mut std::ffi::c_void, +) { + let callback_result = result; + let callback_result_ptr = &callback_result as *const CResult; + + if std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| { + callback(value, callback_result_ptr, user_data) + })) + .is_err() + { + tracing::error!(value, "async bool callback panicked; contained at FFI boundary"); + } + + if !callback_result.error_message.is_null() { + unsafe { + let _ = CString::from_raw(callback_result.error_message); + } + } +} + +fn invoke_record_array_async_callback( + callback: RecordArrayAsyncCallback, + records: CRecordArray, + result: CResult, + user_data: *mut std::ffi::c_void, +) { + let callback_result = result; + let callback_result_ptr = &callback_result as *const CResult; + let callback_records = ManuallyDrop::new(records); + + let panicked = std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| unsafe { + callback(ptr::read(&*callback_records), callback_result_ptr, user_data) + })) + .is_err(); + if panicked { + tracing::error!("async record-array callback panicked; contained at FFI boundary"); + zerobus_free_record_array(ManuallyDrop::into_inner(callback_records)); + } + + if !callback_result.error_message.is_null() { + unsafe { + let _ = CString::from_raw(callback_result.error_message); + } + } +} + +fn encoded_records_to_c_array(records_vec: Vec) -> CRecordArray { + let len = records_vec.len(); + + let mut c_records: Vec = records_vec + .into_iter() + .map(|record| match record { + EncodedRecord::Proto(data) => { + let data_len = data.len(); + let data_ptr = Box::into_raw(data.into_boxed_slice()) as *mut u8; + CRecord { + is_json: false, + data: data_ptr, + data_len, + } + } + EncodedRecord::Json(json_str) => { + let bytes = json_str.into_bytes(); + let data_len = bytes.len(); + let data_ptr = Box::into_raw(bytes.into_boxed_slice()) as *mut u8; + CRecord { + is_json: true, + data: data_ptr, + data_len, + } + } + }) + .collect(); + + let records_ptr = c_records.as_mut_ptr(); + std::mem::forget(c_records); + + CRecordArray { + records: records_ptr, + len, + } +} + /// Create a stream with OAuth authentication /// descriptor_proto_bytes: protobuf-encoded DescriptorProto (can be NULL for JSON streams) #[no_mangle] @@ -69,52 +293,49 @@ pub extern "C" fn zerobus_sdk_create_stream( }; let res = RUNTIME.block_on(async { - let table_name_str = unsafe { c_str_to_string(table_name).map_err(|e| e.to_string())? }; - let client_id_str = unsafe { c_str_to_string(client_id).map_err(|e| e.to_string())? }; - let client_secret_str = - unsafe { c_str_to_string(client_secret).map_err(|e| e.to_string())? }; + let table_name_str = unsafe { + c_str_to_string(table_name) + .map_err(|e| ZerobusError::InvalidArgument(e.to_string()))? + }; + let client_id_str = unsafe { + c_str_to_string(client_id) + .map_err(|e| ZerobusError::InvalidArgument(e.to_string()))? + }; + let client_secret_str = unsafe { + c_str_to_string(client_secret) + .map_err(|e| ZerobusError::InvalidArgument(e.to_string()))? + }; let descriptor_proto = if !descriptor_proto_bytes.is_null() && descriptor_proto_len > 0 { let bytes = unsafe { std::slice::from_raw_parts(descriptor_proto_bytes, descriptor_proto_len) }; - Some(prost_types::DescriptorProto::decode(bytes).map_err(|e| e.to_string())?) + Some( + prost_types::DescriptorProto::decode(bytes) + .map_err(|e| ZerobusError::InvalidArgument(e.to_string()))?, + ) } else { None }; let c_opts = if !options.is_null() { - Some(unsafe { &*options }) + Some(unsafe { *options }) } else { None }; - let record_type = c_opts - .map(|c| c_record_type(c.record_type)) - .unwrap_or(RecordType::Proto); - - let base = sdk_ref - .stream_builder() - .table(table_name_str) - .oauth(client_id_str, client_secret_str); - let mut builder = match record_type { - RecordType::Proto => { - let desc = descriptor_proto.ok_or_else(|| { - "Proto descriptor is required for Proto record type".to_string() - })?; - base.compiled_proto(desc) - } - RecordType::Json => base.json(), - RecordType::Unspecified => return Err("Record type is not specified".to_string()), - }; - if let Some(c) = c_opts { - builder = apply_c_stream_options(builder, c); - } - - let stream = builder.build().await.map_err(|e| e.to_string())?; - let arc = Arc::new(stream); - Ok::<*mut CZerobusStream, String>(Arc::into_raw(arc) as *mut CZerobusStream) + build_stream_from_parts( + sdk_ref, + table_name_str, + descriptor_proto, + StreamCreateAuth::OAuth { + client_id: client_id_str, + client_secret: client_secret_str, + }, + c_opts, + ) + .await }); match res { @@ -123,13 +344,134 @@ pub extern "C" fn zerobus_sdk_create_stream( stream_ptr } Err(err) => { - write_error_result(result, &err, false); + if !result.is_null() { + unsafe { + *result = CResult::error(err); + } + } ptr::null_mut() } } }) } +/// Create a stream with OAuth authentication on a background task. +/// +/// Returns `true` once the request has been validated and scheduled. The +/// callback is invoked exactly once with either a non-null stream pointer and a +/// success result, or a null stream pointer and a failure result. The SDK +/// handle must remain valid until the callback runs. +#[no_mangle] +pub extern "C" fn zerobus_sdk_create_stream_async( + sdk: *mut CZerobusSdk, + table_name: *const c_char, + descriptor_proto_bytes: *const u8, + descriptor_proto_len: usize, + client_id: *const c_char, + client_secret: *const c_char, + options: *const CStreamConfigurationOptions, + callback: CreateStreamAsyncCallback, + user_data: *mut std::ffi::c_void, + result: *mut CResult, +) -> bool { + ffi_guard(result, false, move || { + if let Err(msg) = validate_sdk_ptr(sdk) { + write_error_result(result, msg, false); + return false; + } + + let table_name_str = match unsafe { c_str_to_string(table_name) } { + Ok(s) => s, + Err(e) => { + write_error_result(result, e, false); + return false; + } + }; + let client_id_str = match unsafe { c_str_to_string(client_id) } { + Ok(s) => s, + Err(e) => { + write_error_result(result, e, false); + return false; + } + }; + let client_secret_str = match unsafe { c_str_to_string(client_secret) } { + Ok(s) => s, + Err(e) => { + write_error_result(result, e, false); + return false; + } + }; + + let descriptor_proto = if !descriptor_proto_bytes.is_null() && descriptor_proto_len > 0 { + let bytes = unsafe { + std::slice::from_raw_parts(descriptor_proto_bytes, descriptor_proto_len) + }; + match prost_types::DescriptorProto::decode(bytes) { + Ok(desc) => Some(desc), + Err(e) => { + write_error_result(result, &e.to_string(), false); + return false; + } + } + } else { + None + }; + + let c_opts = if !options.is_null() { + Some(unsafe { *options }) + } else { + None + }; + + let sdk_ptr = SendPtr::new(sdk); + let callback_user_data = SendPtr::new(user_data); + RUNTIME.spawn(async move { + let callback_result = match validate_sdk_ptr(sdk_ptr.get()) { + Ok(sdk_ref) => match build_stream_from_parts( + sdk_ref, + table_name_str, + descriptor_proto, + StreamCreateAuth::OAuth { + client_id: client_id_str, + client_secret: client_secret_str, + }, + c_opts, + ) + .await + { + Ok(stream_ptr) => { + invoke_create_stream_async_callback( + callback, + stream_ptr, + CResult::success(), + callback_user_data.get(), + ); + return; + } + Err(err) => CResult::error(err), + }, + Err(msg) => CResult { + success: false, + error_message: CString::new(msg) + .unwrap_or_else(|_| CString::new("SDK pointer is invalid").unwrap()) + .into_raw(), + is_retryable: false, + }, + }; + + invoke_create_stream_async_callback( + callback, + ptr::null_mut(), + callback_result, + callback_user_data.get(), + ); + }); + + write_success_result(result); + true + }) +} + /// Create a stream with a custom headers provider callback /// This allows you to provide custom authentication headers via a Go callback function #[no_mangle] @@ -153,52 +495,194 @@ pub extern "C" fn zerobus_sdk_create_stream_with_headers_provider( }; let res = RUNTIME.block_on(async { - let table_name_str = unsafe { c_str_to_string(table_name).map_err(|e| e.to_string())? }; + let table_name_str = unsafe { + c_str_to_string(table_name) + .map_err(|e| ZerobusError::InvalidArgument(e.to_string()))? + }; let descriptor_proto = if !descriptor_proto_bytes.is_null() && descriptor_proto_len > 0 { let bytes = unsafe { std::slice::from_raw_parts(descriptor_proto_bytes, descriptor_proto_len) }; - Some(prost_types::DescriptorProto::decode(bytes).map_err(|e| e.to_string())?) + Some( + prost_types::DescriptorProto::decode(bytes) + .map_err(|e| ZerobusError::InvalidArgument(e.to_string()))?, + ) } else { None }; let c_opts = if !options.is_null() { - Some(unsafe { &*options }) + Some(unsafe { *options }) } else { None }; - let record_type = c_opts - .map(|c| c_record_type(c.record_type)) - .unwrap_or(RecordType::Proto); - let headers_provider: Arc = - Arc::new(CallbackHeadersProvider::new(headers_callback, user_data)); + build_stream_from_parts( + sdk_ref, + table_name_str, + descriptor_proto, + StreamCreateAuth::HeadersProvider { + headers_callback, + user_data: SendPtr::new(user_data), + }, + c_opts, + ) + .await + }); - let base = sdk_ref - .stream_builder() - .table(table_name_str) - .headers_provider(headers_provider); - let mut builder = match record_type { - RecordType::Proto => { - let desc = descriptor_proto.ok_or_else(|| { - "Proto descriptor is required for Proto record type".to_string() - })?; - base.compiled_proto(desc) + match res { + Ok(stream_ptr) => { + write_success_result(result); + stream_ptr + } + Err(err) => { + if !result.is_null() { + unsafe { + *result = CResult::error(err); + } + } + ptr::null_mut() + } + } + }) +} + +/// Create a stream with a custom headers provider callback on a background task. +/// +/// Returns `true` once the request has been validated and scheduled. The +/// callback is invoked exactly once with either a non-null stream pointer and a +/// success result, or a null stream pointer and a failure result. The SDK +/// handle must remain valid until the callback runs. +#[no_mangle] +pub extern "C" fn zerobus_sdk_create_stream_with_headers_provider_async( + sdk: *mut CZerobusSdk, + table_name: *const c_char, + descriptor_proto_bytes: *const u8, + descriptor_proto_len: usize, + headers_callback: HeadersProviderCallback, + user_data: *mut std::ffi::c_void, + options: *const CStreamConfigurationOptions, + callback: CreateStreamAsyncCallback, + callback_user_data: *mut std::ffi::c_void, + result: *mut CResult, +) -> bool { + ffi_guard(result, false, move || { + if let Err(msg) = validate_sdk_ptr(sdk) { + write_error_result(result, msg, false); + return false; + } + + let table_name_str = match unsafe { c_str_to_string(table_name) } { + Ok(s) => s, + Err(e) => { + write_error_result(result, e, false); + return false; + } + }; + + let descriptor_proto = if !descriptor_proto_bytes.is_null() && descriptor_proto_len > 0 { + let bytes = unsafe { + std::slice::from_raw_parts(descriptor_proto_bytes, descriptor_proto_len) + }; + match prost_types::DescriptorProto::decode(bytes) { + Ok(desc) => Some(desc), + Err(e) => { + write_error_result(result, &e.to_string(), false); + return false; } - RecordType::Json => base.json(), - RecordType::Unspecified => return Err("Record type is not specified".to_string()), + } + } else { + None + }; + + let c_opts = if !options.is_null() { + Some(unsafe { *options }) + } else { + None + }; + + let sdk_ptr = SendPtr::new(sdk); + let stream_user_data = SendPtr::new(user_data); + let callback_user_data = SendPtr::new(callback_user_data); + RUNTIME.spawn(async move { + let callback_result = match validate_sdk_ptr(sdk_ptr.get()) { + Ok(sdk_ref) => match build_stream_from_parts( + sdk_ref, + table_name_str, + descriptor_proto, + StreamCreateAuth::HeadersProvider { + headers_callback, + user_data: stream_user_data, + }, + c_opts, + ) + .await + { + Ok(stream_ptr) => { + invoke_create_stream_async_callback( + callback, + stream_ptr, + CResult::success(), + callback_user_data.get(), + ); + return; + } + Err(err) => CResult::error(err), + }, + Err(msg) => CResult { + success: false, + error_message: CString::new(msg) + .unwrap_or_else(|_| CString::new("SDK pointer is invalid").unwrap()) + .into_raw(), + is_retryable: false, + }, }; - if let Some(c) = c_opts { - builder = apply_c_stream_options(builder, c); + + invoke_create_stream_async_callback( + callback, + ptr::null_mut(), + callback_result, + callback_user_data.get(), + ); + }); + + write_success_result(result); + true + }) +} + +/// Recreate a stream from an existing stream +/// This is used for recovery scenarios where the stream needs to be re-established +#[no_mangle] +pub extern "C" fn zerobus_sdk_recreate_stream( + sdk: *mut CZerobusSdk, + stream: *mut CZerobusStream, + result: *mut CResult, +) -> *mut CZerobusStream { + ffi_guard(result, ptr::null_mut(), move || { + let sdk_ref = match validate_sdk_ptr(sdk) { + Ok(s) => s, + Err(msg) => { + write_error_result(result, msg, false); + return ptr::null_mut(); } + }; - let stream = builder.build().await.map_err(|e| e.to_string())?; + let stream_ref = match validate_stream_ptr(stream) { + Ok(s) => s, + Err(msg) => { + write_error_result(result, msg, false); + return ptr::null_mut(); + } + }; + + let res = RUNTIME.block_on(async { + let new_stream = sdk_ref.recreate_stream(stream_ref).await?; - let arc = Arc::new(stream); - Ok::<*mut CZerobusStream, String>(Arc::into_raw(arc) as *mut CZerobusStream) + let arc = Arc::new(new_stream); + Ok::<*mut CZerobusStream, ZerobusError>(Arc::into_raw(arc) as *mut CZerobusStream) }); match res { @@ -207,13 +691,76 @@ pub extern "C" fn zerobus_sdk_create_stream_with_headers_provider( stream_ptr } Err(err) => { - write_error_result(result, &err, false); + if !result.is_null() { + unsafe { + *result = CResult::error(err); + } + } ptr::null_mut() } } }) } +/// Recreate a stream from an existing stream on a background task. +#[no_mangle] +pub extern "C" fn zerobus_sdk_recreate_stream_async( + sdk: *mut CZerobusSdk, + stream: *mut CZerobusStream, + callback: CreateStreamAsyncCallback, + user_data: *mut std::ffi::c_void, + result: *mut CResult, +) -> bool { + ffi_guard(result, false, move || { + if let Err(msg) = validate_sdk_ptr(sdk) { + write_error_result(result, msg, false); + return false; + } + if let Err(msg) = validate_stream_ptr(stream) { + write_error_result(result, msg, false); + return false; + } + + let sdk_ptr = SendPtr::new(sdk); + let stream_arc = unsafe { clone_stream_arc(stream) }; + let callback_user_data = SendPtr::new(user_data); + RUNTIME.spawn(async move { + let callback_result = match validate_sdk_ptr(sdk_ptr.get()) { + Ok(sdk_ref) => match sdk_ref.recreate_stream(stream_arc.as_ref()).await { + Ok(new_stream) => { + let stream_ptr = Arc::into_raw(Arc::new(new_stream)) as *mut CZerobusStream; + invoke_create_stream_async_callback( + callback, + stream_ptr, + CResult::success(), + callback_user_data.get(), + ); + return; + } + Err(err) => CResult::error(err), + }, + Err(msg) => CResult { + success: false, + error_message: CString::new(msg) + .unwrap_or_else(|_| CString::new("SDK pointer is invalid").unwrap()) + .into_raw(), + is_retryable: false, + }, + }; + + invoke_create_stream_async_callback( + callback, + ptr::null_mut(), + callback_result, + callback_user_data.get(), + ); + }); + + write_success_result(result); + true + }) +} + /// Free a stream instance #[no_mangle] pub extern "C" fn zerobus_stream_free(stream: *mut CZerobusStream) { @@ -278,6 +825,54 @@ pub extern "C" fn zerobus_stream_ingest_proto_record( }) } +/// Ingest a protobuf record on a background task and report the assigned offset via callback. +#[no_mangle] +pub extern "C" fn zerobus_stream_ingest_proto_record_async( + stream: *mut CZerobusStream, + data: *const u8, + data_len: usize, + callback: OffsetAsyncCallback, + user_data: *mut std::ffi::c_void, + result: *mut CResult, +) -> bool { + ffi_guard(result, false, move || { + if data.is_null() { + write_error_result(result, "Invalid data pointer", false); + return false; + } + if let Err(msg) = validate_stream_ptr(stream) { + write_error_result(result, msg, false); + return false; + } + + let data_vec = unsafe { std::slice::from_raw_parts(data, data_len) }.to_vec(); + let stream_arc = unsafe { clone_stream_arc(stream) }; + let callback_user_data = SendPtr::new(user_data); + RUNTIME.spawn(async move { + match stream_arc + .ingest_record_offset(EncodedRecord::Proto(data_vec)) + .await + { + Ok(offset) => invoke_offset_async_callback( + callback, + offset, + CResult::success(), + callback_user_data.get(), + ), + Err(err) => invoke_offset_async_callback( + callback, + -1, + CResult::error(err), + callback_user_data.get(), + ), + } + }); + + write_success_result(result); + true + }) +} + /// Ingest a JSON record /// Returns the offset directly /// Returns -1 on error @@ -327,6 +922,53 @@ pub extern "C" fn zerobus_stream_ingest_json_record( }) } +/// Ingest a JSON record on a background task and report the assigned offset via callback. +#[no_mangle] +pub extern "C" fn zerobus_stream_ingest_json_record_async( + stream: *mut CZerobusStream, + json_data: *const c_char, + callback: OffsetAsyncCallback, + user_data: *mut std::ffi::c_void, + result: *mut CResult, +) -> bool { + ffi_guard(result, false, move || { + if let Err(msg) = validate_stream_ptr(stream) { + write_error_result(result, msg, false); + return false; + } + + let json_str = match unsafe { c_str_to_string(json_data) } { + Ok(s) => s, + Err(e) => { + write_error_result(result, e, false); + return false; + } + }; + + let stream_arc = unsafe { clone_stream_arc(stream) }; + let callback_user_data = SendPtr::new(user_data); + RUNTIME.spawn(async move { + match stream_arc.ingest_record_offset(EncodedRecord::Json(json_str)).await { + Ok(offset) => invoke_offset_async_callback( + callback, + offset, + CResult::success(), + callback_user_data.get(), + ), + Err(err) => invoke_offset_async_callback( + callback, + -1, + CResult::error(err), + callback_user_data.get(), + ), + } + }); + + write_success_result(result); + true + }) +} + /// Ingest a batch of protobuf records /// Returns the offset of the last record in the batch, or -1 on error /// Returns -2 if batch is empty @@ -400,6 +1042,83 @@ pub extern "C" fn zerobus_stream_ingest_proto_records( }) } +/// Ingest a batch of protobuf records on a background task and report the last offset via callback. +#[no_mangle] +pub extern "C" fn zerobus_stream_ingest_proto_records_async( + stream: *mut CZerobusStream, + records: *const *const u8, + record_lens: *const usize, + num_records: usize, + callback: OffsetAsyncCallback, + user_data: *mut std::ffi::c_void, + result: *mut CResult, +) -> bool { + ffi_guard(result, false, move || { + if records.is_null() || record_lens.is_null() { + write_error_result(result, "Invalid records pointer", false); + return false; + } + + let callback_user_data = SendPtr::new(user_data); + if num_records == 0 { + RUNTIME.spawn(async move { + invoke_offset_async_callback( + callback, + -2, + CResult::success(), + callback_user_data.get(), + ); + }); + write_success_result(result); + return true; + } + + if let Err(msg) = validate_stream_ptr(stream) { + write_error_result(result, msg, false); + return false; + } + + let records_vec: Vec> = unsafe { + let records_slice = std::slice::from_raw_parts(records, num_records); + let lens_slice = std::slice::from_raw_parts(record_lens, num_records); + records_slice + .iter() + .zip(lens_slice.iter()) + .map(|(ptr, len)| std::slice::from_raw_parts(*ptr, *len).to_vec()) + .collect() + }; + + let stream_arc = unsafe { clone_stream_arc(stream) }; + RUNTIME.spawn(async move { + let payloads: Vec = + records_vec.into_iter().map(EncodedRecord::Proto).collect(); + match stream_arc.ingest_records_offset(payloads).await { + Ok(Some(offset)) => invoke_offset_async_callback( + callback, + offset, + CResult::success(), + callback_user_data.get(), + ), + Ok(None) => invoke_offset_async_callback( + callback, + -2, + CResult::success(), + callback_user_data.get(), + ), + Err(err) => invoke_offset_async_callback( + callback, + -1, + CResult::error(err), + callback_user_data.get(), + ), + } + }); + + write_success_result(result); + true + }) +} + /// Ingest a batch of JSON records /// Returns the offset of the last record in the batch, or -1 on error /// Returns -2 if batch is empty @@ -471,6 +1190,84 @@ pub extern "C" fn zerobus_stream_ingest_json_records( }) } +/// Ingest a batch of JSON records on a background task and report the last offset via callback. +#[no_mangle] +pub extern "C" fn zerobus_stream_ingest_json_records_async( + stream: *mut CZerobusStream, + json_records: *const *const c_char, + num_records: usize, + callback: OffsetAsyncCallback, + user_data: *mut std::ffi::c_void, + result: *mut CResult, +) -> bool { + ffi_guard(result, false, move || { + if json_records.is_null() { + write_error_result(result, "Invalid records pointer", false); + return false; + } + + let callback_user_data = SendPtr::new(user_data); + if num_records == 0 { + RUNTIME.spawn(async move { + invoke_offset_async_callback( + callback, + -2, + CResult::success(), + callback_user_data.get(), + ); + }); + write_success_result(result); + return true; + } + + if let Err(msg) = validate_stream_ptr(stream) { + write_error_result(result, msg, false); + return false; + } + + let json_vec: Result, _> = unsafe { + let json_slice = std::slice::from_raw_parts(json_records, num_records); + json_slice.iter().map(|ptr| c_str_to_string(*ptr)).collect() + }; + let json_vec = match json_vec { + Ok(v) => v, + Err(e) => { + write_error_result(result, e, false); + return false; + } + }; + + let stream_arc = unsafe { clone_stream_arc(stream) }; + RUNTIME.spawn(async move { + let payloads: Vec = + json_vec.into_iter().map(EncodedRecord::Json).collect(); + match stream_arc.ingest_records_offset(payloads).await { + Ok(Some(offset)) => invoke_offset_async_callback( + callback, + offset, + CResult::success(), + callback_user_data.get(), + ), + Ok(None) => invoke_offset_async_callback( + callback, + -2, + CResult::success(), + callback_user_data.get(), + ), + Err(err) => invoke_offset_async_callback( + callback, + -1, + CResult::error(err), + callback_user_data.get(), + ), + } + }); + + write_success_result(result); + true + }) +} + /// Clones the `Arc` from a raw `CZerobusStream` pointer without /// consuming the pointer. The caller retains ownership of the original pointer; /// the returned `Arc` will keep the stream alive until it is dropped. @@ -702,6 +1499,45 @@ pub extern "C" fn zerobus_stream_wait_for_offset( }) } +/// Wait for an offset on a background task and report completion via callback. +#[no_mangle] +pub extern "C" fn zerobus_stream_wait_for_offset_async( + stream: *mut CZerobusStream, + offset: i64, + callback: BoolAsyncCallback, + user_data: *mut std::ffi::c_void, + result: *mut CResult, +) -> bool { + ffi_guard(result, false, move || { + if let Err(msg) = validate_stream_ptr(stream) { + write_error_result(result, msg, false); + return false; + } + + let stream_arc = unsafe { clone_stream_arc(stream) }; + let callback_user_data = SendPtr::new(user_data); + RUNTIME.spawn(async move { + match stream_arc.wait_for_offset(offset).await { + Ok(()) => invoke_bool_async_callback( + callback, + true, + CResult::success(), + callback_user_data.get(), + ), + Err(err) => invoke_bool_async_callback( + callback, + false, + CResult::error(err), + callback_user_data.get(), + ), + } + }); + + write_success_result(result); + true + }) +} + /// Flush all pending records #[no_mangle] pub extern "C" fn zerobus_stream_flush(stream: *mut CZerobusStream, result: *mut CResult) -> bool { @@ -733,6 +1569,44 @@ pub extern "C" fn zerobus_stream_flush(stream: *mut CZerobusStream, result: *mut }) } +/// Flush all pending records on a background task and report completion via callback. +#[no_mangle] +pub extern "C" fn zerobus_stream_flush_async( + stream: *mut CZerobusStream, + callback: BoolAsyncCallback, + user_data: *mut std::ffi::c_void, + result: *mut CResult, +) -> bool { + ffi_guard(result, false, move || { + if let Err(msg) = validate_stream_ptr(stream) { + write_error_result(result, msg, false); + return false; + } + + let stream_arc = unsafe { clone_stream_arc(stream) }; + let callback_user_data = SendPtr::new(user_data); + RUNTIME.spawn(async move { + match stream_arc.flush().await { + Ok(_) => invoke_bool_async_callback( + callback, + true, + CResult::success(), + callback_user_data.get(), + ), + Err(err) => invoke_bool_async_callback( + callback, + false, + CResult::error(err), + callback_user_data.get(), + ), + } + }); + + write_success_result(result); + true + }) +} + /// Get unacknowledged records from a closed stream /// Returns a CRecordArray that must be freed with zerobus_free_record_array #[no_mangle] @@ -760,44 +1634,8 @@ pub extern "C" fn zerobus_stream_get_unacked_records( match records_res { Ok(records_iter) => { - // Collect into Vec - let records_vec: Vec = records_iter.collect(); - let len = records_vec.len(); - - // Convert to CRecords - let mut c_records: Vec = records_vec - .into_iter() - .map(|record| match record { - EncodedRecord::Proto(data) => { - let data_len = data.len(); - let data_ptr = Box::into_raw(data.into_boxed_slice()) as *mut u8; - CRecord { - is_json: false, - data: data_ptr, - data_len, - } - } - EncodedRecord::Json(json_str) => { - let bytes = json_str.into_bytes(); - let data_len = bytes.len(); - let data_ptr = Box::into_raw(bytes.into_boxed_slice()) as *mut u8; - CRecord { - is_json: true, - data: data_ptr, - data_len, - } - } - }) - .collect(); - - let records_ptr = c_records.as_mut_ptr(); - std::mem::forget(c_records); // Don't drop, Go will call free - write_success_result(result); - CRecordArray { - records: records_ptr, - len, - } + encoded_records_to_c_array(records_iter.collect()) } Err(err) => { if !result.is_null() { @@ -814,6 +1652,47 @@ pub extern "C" fn zerobus_stream_get_unacked_records( }) } +/// Get unacknowledged records from a closed stream on a background task. +#[no_mangle] +pub extern "C" fn zerobus_stream_get_unacked_records_async( + stream: *mut CZerobusStream, + callback: RecordArrayAsyncCallback, + user_data: *mut std::ffi::c_void, + result: *mut CResult, +) -> bool { + ffi_guard(result, false, move || { + if let Err(msg) = validate_stream_ptr(stream) { + write_error_result(result, msg, false); + return false; + } + + let stream_arc = unsafe { clone_stream_arc(stream) }; + let callback_user_data = SendPtr::new(user_data); + RUNTIME.spawn(async move { + match stream_arc.get_unacked_records().await { + Ok(records_iter) => invoke_record_array_async_callback( + callback, + encoded_records_to_c_array(records_iter.collect()), + CResult::success(), + callback_user_data.get(), + ), + Err(err) => invoke_record_array_async_callback( + callback, + CRecordArray { + records: ptr::null_mut(), + len: 0, + }, + CResult::error(err), + callback_user_data.get(), + ), + } + }); + + write_success_result(result); + true + }) +} + /// Free a CRecordArray returned by zerobus_stream_get_unacked_records #[no_mangle] pub extern "C" fn zerobus_free_record_array(array: CRecordArray) { @@ -864,6 +1743,73 @@ pub extern "C" fn zerobus_stream_close(stream: *mut CZerobusStream, result: *mut }) } +/// Close the stream gracefully on a background task. +#[no_mangle] +pub extern "C" fn zerobus_stream_close_async( + stream: *mut CZerobusStream, + callback: BoolAsyncCallback, + user_data: *mut std::ffi::c_void, + result: *mut CResult, +) -> bool { + ffi_guard(result, false, move || { + if let Err(msg) = validate_stream_ptr_mut(stream) { + write_error_result(result, msg, false); + return false; + } + + let stream_ptr = SendPtr::new(stream); + let callback_user_data = SendPtr::new(user_data); + RUNTIME.spawn(async move { + let callback_result = match validate_stream_ptr_mut(stream_ptr.get()) { + Ok(stream_ref) => match stream_ref.close().await { + Ok(_) => { + invoke_bool_async_callback( + callback, + true, + CResult::success(), + callback_user_data.get(), + ); + return; + } + Err(err) => CResult::error(err), + }, + Err(msg) => CResult { + success: false, + error_message: CString::new(msg) + .unwrap_or_else(|_| CString::new("Stream pointer is invalid").unwrap()) + .into_raw(), + is_retryable: false, + }, + }; + + invoke_bool_async_callback( + callback, + false, + callback_result, + callback_user_data.get(), + ); + }); + + write_success_result(result); + true + }) +} + +/// Returns whether the stream has been closed. +#[no_mangle] +pub extern "C" fn zerobus_stream_is_closed(stream: *mut CZerobusStream) -> bool { + // No CResult out-param; on a caught panic return `true` (treat as closed), + // matching the answer for an invalid handle. + ffi_guard( + ptr::null_mut(), + true, + move || match validate_stream_ptr(stream) { + Ok(s) => s.is_closed(), + Err(_) => true, + }, + ) +} + /// Free error message string #[no_mangle] pub extern "C" fn zerobus_free_error_message(message: *mut c_char) { diff --git a/rust/ffi/src/tests.rs b/rust/ffi/src/tests.rs index a4cc6d3a..0b810c5c 100644 --- a/rust/ffi/src/tests.rs +++ b/rust/ffi/src/tests.rs @@ -6,12 +6,20 @@ mod tests { zerobus_get_default_config, zerobus_sdk_builder_application_name, zerobus_sdk_builder_build, zerobus_sdk_builder_disable_tls, zerobus_sdk_builder_endpoint, zerobus_sdk_builder_free, zerobus_sdk_builder_new, zerobus_sdk_builder_sdk_identifier, - zerobus_sdk_builder_unity_catalog_url, zerobus_sdk_free, CHeaders, CRecordArray, CResult, - CallbackHeadersProvider, RecordType, ZerobusError, + zerobus_sdk_builder_unity_catalog_url, zerobus_sdk_create_stream, + zerobus_sdk_create_stream_async, zerobus_sdk_create_stream_with_headers_provider_async, + zerobus_sdk_recreate_stream_async, + zerobus_sdk_free, CHeaders, CRecordArray, CResult, CallbackHeadersProvider, RecordType, + ZerobusError, zerobus_stream_close_async, zerobus_stream_flush_async, + zerobus_stream_get_unacked_records_async, zerobus_stream_ingest_json_record_async, + zerobus_stream_ingest_json_records_async, zerobus_stream_ingest_proto_record_async, + zerobus_stream_ingest_proto_records_async, zerobus_stream_wait_for_offset_async, }; use databricks_zerobus_ingest_sdk::HeadersProvider; use std::ffi::{CStr, CString}; use std::ptr; + use std::sync::mpsc; + use std::time::Duration; // Helper for c_str_to_string since it's private unsafe fn test_c_str_to_string( @@ -341,6 +349,470 @@ mod tests { assert_eq!(config.record_type, 1); // Proto } + #[test] + fn test_create_stream_retryable_error_sets_retryable_flag() { + // Build an SDK that points to an unreachable local endpoint so stream + // creation fails with a retryable transport/setup error. + let endpoint_c = CString::new("http://127.0.0.1:1").unwrap(); + let uc_c = CString::new("http://127.0.0.1:1").unwrap(); + let builder = zerobus_sdk_builder_new(); + zerobus_sdk_builder_endpoint(builder, endpoint_c.as_ptr()); + zerobus_sdk_builder_unity_catalog_url(builder, uc_c.as_ptr()); + zerobus_sdk_builder_disable_tls(builder); + + let mut build_result = CResult { + success: false, + error_message: ptr::null_mut(), + is_retryable: false, + }; + let sdk = zerobus_sdk_builder_build(builder, &mut build_result as *mut CResult); + assert!(build_result.success, "SDK build should succeed"); + assert!(!sdk.is_null(), "SDK pointer should be non-null"); + + let table = CString::new("main.default.events").unwrap(); + let client_id = CString::new("client-id").unwrap(); + let client_secret = CString::new("client-secret").unwrap(); + let mut opts = zerobus_get_default_config(); + opts.record_type = 2; // JSON stream: no proto descriptor required. + + let mut create_result = CResult { + success: true, + error_message: ptr::null_mut(), + is_retryable: false, + }; + + let stream = zerobus_sdk_create_stream( + sdk, + table.as_ptr(), + ptr::null(), + 0, + client_id.as_ptr(), + client_secret.as_ptr(), + &opts as *const _, + &mut create_result as *mut CResult, + ); + + assert!(stream.is_null(), "create_stream should fail"); + assert!(!create_result.success, "result should indicate failure"); + assert!( + create_result.is_retryable, + "retryable create failures must set is_retryable=true" + ); + + zerobus_free_error_message(create_result.error_message); + zerobus_sdk_free(sdk); + } + + #[test] + fn test_create_stream_async_reports_completion_error_via_callback() { + extern "C" fn create_callback( + stream: *mut crate::CZerobusStream, + result: *const CResult, + user_data: *mut std::ffi::c_void, + ) { + let sender = unsafe { &*(user_data as *const mpsc::Sender<(bool, bool, bool, String)>) }; + let result_ref = unsafe { &*result }; + let error_message = if result_ref.error_message.is_null() { + String::new() + } else { + unsafe { CStr::from_ptr(result_ref.error_message) } + .to_string_lossy() + .into_owned() + }; + sender + .send(( + stream.is_null(), + result_ref.success, + result_ref.is_retryable, + error_message, + )) + .unwrap(); + } + + let endpoint_c = CString::new("http://127.0.0.1:1").unwrap(); + let uc_c = CString::new("http://127.0.0.1:1").unwrap(); + let builder = zerobus_sdk_builder_new(); + zerobus_sdk_builder_endpoint(builder, endpoint_c.as_ptr()); + zerobus_sdk_builder_unity_catalog_url(builder, uc_c.as_ptr()); + zerobus_sdk_builder_disable_tls(builder); + + let mut build_result = CResult { + success: false, + error_message: ptr::null_mut(), + is_retryable: false, + }; + let sdk = zerobus_sdk_builder_build(builder, &mut build_result as *mut CResult); + assert!(build_result.success, "SDK build should succeed"); + assert!(!sdk.is_null(), "SDK pointer should be non-null"); + + let table = CString::new("main.default.events").unwrap(); + let client_id = CString::new("client-id").unwrap(); + let client_secret = CString::new("client-secret").unwrap(); + let mut create_result = CResult { + success: false, + error_message: ptr::null_mut(), + is_retryable: false, + }; + let (sender, receiver): ( + mpsc::Sender<(bool, bool, bool, String)>, + mpsc::Receiver<(bool, bool, bool, String)>, + ) = mpsc::channel(); + let sender = Box::new(sender); + + let started = zerobus_sdk_create_stream_async( + sdk, + table.as_ptr(), + ptr::null(), + 0, + client_id.as_ptr(), + client_secret.as_ptr(), + ptr::null(), + create_callback, + Box::as_ref(&sender) as *const _ as *mut std::ffi::c_void, + &mut create_result as *mut CResult, + ); + + assert!(started, "create_stream_async should schedule the task"); + assert!(create_result.success, "scheduling result should succeed"); + + let (stream_is_null, callback_success, callback_retryable, callback_message) = receiver + .recv_timeout(Duration::from_secs(2)) + .expect("callback should be invoked"); + assert!(stream_is_null, "callback should receive a null stream on failure"); + assert!(!callback_success, "callback result should indicate failure"); + assert!( + !callback_retryable, + "missing proto descriptor should be reported as a non-retryable error" + ); + assert!( + callback_message.contains("Proto descriptor is required for Proto record type"), + "callback should surface the create error" + ); + + drop(sender); + zerobus_sdk_free(sdk); + } + + #[test] + fn test_create_stream_with_headers_provider_async_reports_completion_error_via_callback() { + extern "C" fn headers_callback(_user_data: *mut std::ffi::c_void) -> CHeaders { + CHeaders { + headers: ptr::null_mut(), + count: 0, + error_message: ptr::null_mut(), + } + } + + extern "C" fn create_callback( + stream: *mut crate::CZerobusStream, + result: *const CResult, + user_data: *mut std::ffi::c_void, + ) { + let sender = unsafe { &*(user_data as *const mpsc::Sender<(bool, bool, bool, String)>) }; + let result_ref = unsafe { &*result }; + let error_message = if result_ref.error_message.is_null() { + String::new() + } else { + unsafe { CStr::from_ptr(result_ref.error_message) } + .to_string_lossy() + .into_owned() + }; + sender + .send(( + stream.is_null(), + result_ref.success, + result_ref.is_retryable, + error_message, + )) + .unwrap(); + } + + let endpoint_c = CString::new("http://127.0.0.1:1").unwrap(); + let uc_c = CString::new("http://127.0.0.1:1").unwrap(); + let builder = zerobus_sdk_builder_new(); + zerobus_sdk_builder_endpoint(builder, endpoint_c.as_ptr()); + zerobus_sdk_builder_unity_catalog_url(builder, uc_c.as_ptr()); + zerobus_sdk_builder_disable_tls(builder); + + let mut build_result = CResult { + success: false, + error_message: ptr::null_mut(), + is_retryable: false, + }; + let sdk = zerobus_sdk_builder_build(builder, &mut build_result as *mut CResult); + assert!(build_result.success, "SDK build should succeed"); + assert!(!sdk.is_null(), "SDK pointer should be non-null"); + + let table = CString::new("main.default.events").unwrap(); + let mut create_result = CResult { + success: false, + error_message: ptr::null_mut(), + is_retryable: false, + }; + let (sender, receiver): ( + mpsc::Sender<(bool, bool, bool, String)>, + mpsc::Receiver<(bool, bool, bool, String)>, + ) = mpsc::channel(); + let sender = Box::new(sender); + + let started = zerobus_sdk_create_stream_with_headers_provider_async( + sdk, + table.as_ptr(), + ptr::null(), + 0, + headers_callback, + ptr::null_mut(), + ptr::null(), + create_callback, + Box::as_ref(&sender) as *const _ as *mut std::ffi::c_void, + &mut create_result as *mut CResult, + ); + + assert!( + started, + "create_stream_with_headers_provider_async should schedule the task" + ); + assert!(create_result.success, "scheduling result should succeed"); + + let (stream_is_null, callback_success, callback_retryable, callback_message) = receiver + .recv_timeout(Duration::from_secs(2)) + .expect("callback should be invoked"); + assert!(stream_is_null, "callback should receive a null stream on failure"); + assert!(!callback_success, "callback result should indicate failure"); + assert!( + !callback_retryable, + "missing proto descriptor should be reported as a non-retryable error" + ); + assert!( + callback_message.contains("Proto descriptor is required for Proto record type"), + "callback should surface the create error" + ); + + drop(sender); + zerobus_sdk_free(sdk); + } + + #[test] + fn test_async_overloads_fail_fast_on_invalid_input() { + extern "C" fn stream_cb( + _stream: *mut crate::CZerobusStream, + _result: *const CResult, + _user_data: *mut std::ffi::c_void, + ) { + } + + extern "C" fn offset_cb( + _offset: i64, + _result: *const CResult, + _user_data: *mut std::ffi::c_void, + ) { + } + + extern "C" fn bool_cb( + _value: bool, + _result: *const CResult, + _user_data: *mut std::ffi::c_void, + ) { + } + + extern "C" fn records_cb( + _records: CRecordArray, + _result: *const CResult, + _user_data: *mut std::ffi::c_void, + ) { + } + + extern "C" fn headers_cb(_user_data: *mut std::ffi::c_void) -> CHeaders { + CHeaders { + headers: ptr::null_mut(), + count: 0, + error_message: ptr::null_mut(), + } + } + + let mut r = CResult { + success: true, + error_message: ptr::null_mut(), + is_retryable: false, + }; + + let table = CString::new("main.default.events").unwrap(); + let client_id = CString::new("cid").unwrap(); + let client_secret = CString::new("secret").unwrap(); + let json = CString::new("{}").unwrap(); + + assert!(!zerobus_sdk_create_stream_async( + ptr::null_mut(), + table.as_ptr(), + ptr::null(), + 0, + client_id.as_ptr(), + client_secret.as_ptr(), + ptr::null(), + stream_cb, + ptr::null_mut(), + &mut r as *mut CResult, + )); + assert!(!r.success); + zerobus_free_error_message(r.error_message); + + r = CResult { + success: true, + error_message: ptr::null_mut(), + is_retryable: false, + }; + assert!(!zerobus_sdk_create_stream_with_headers_provider_async( + ptr::null_mut(), + table.as_ptr(), + ptr::null(), + 0, + headers_cb, + ptr::null_mut(), + ptr::null(), + stream_cb, + ptr::null_mut(), + &mut r as *mut CResult, + )); + assert!(!r.success); + zerobus_free_error_message(r.error_message); + + r = CResult { + success: true, + error_message: ptr::null_mut(), + is_retryable: false, + }; + assert!(!zerobus_sdk_recreate_stream_async( + ptr::null_mut(), + ptr::null_mut(), + stream_cb, + ptr::null_mut(), + &mut r as *mut CResult, + )); + assert!(!r.success); + zerobus_free_error_message(r.error_message); + + r = CResult { + success: true, + error_message: ptr::null_mut(), + is_retryable: false, + }; + assert!(!zerobus_stream_ingest_proto_record_async( + ptr::null_mut(), + ptr::null(), + 0, + offset_cb, + ptr::null_mut(), + &mut r as *mut CResult, + )); + assert!(!r.success); + zerobus_free_error_message(r.error_message); + + r = CResult { + success: true, + error_message: ptr::null_mut(), + is_retryable: false, + }; + assert!(!zerobus_stream_ingest_json_record_async( + ptr::null_mut(), + json.as_ptr(), + offset_cb, + ptr::null_mut(), + &mut r as *mut CResult, + )); + assert!(!r.success); + zerobus_free_error_message(r.error_message); + + r = CResult { + success: true, + error_message: ptr::null_mut(), + is_retryable: false, + }; + assert!(!zerobus_stream_ingest_proto_records_async( + ptr::null_mut(), + ptr::null(), + ptr::null(), + 1, + offset_cb, + ptr::null_mut(), + &mut r as *mut CResult, + )); + assert!(!r.success); + zerobus_free_error_message(r.error_message); + + r = CResult { + success: true, + error_message: ptr::null_mut(), + is_retryable: false, + }; + assert!(!zerobus_stream_ingest_json_records_async( + ptr::null_mut(), + ptr::null(), + 1, + offset_cb, + ptr::null_mut(), + &mut r as *mut CResult, + )); + assert!(!r.success); + zerobus_free_error_message(r.error_message); + + r = CResult { + success: true, + error_message: ptr::null_mut(), + is_retryable: false, + }; + assert!(!zerobus_stream_wait_for_offset_async( + ptr::null_mut(), + 1, + bool_cb, + ptr::null_mut(), + &mut r as *mut CResult, + )); + assert!(!r.success); + zerobus_free_error_message(r.error_message); + + r = CResult { + success: true, + error_message: ptr::null_mut(), + is_retryable: false, + }; + assert!(!zerobus_stream_flush_async( + ptr::null_mut(), + bool_cb, + ptr::null_mut(), + &mut r as *mut CResult, + )); + assert!(!r.success); + zerobus_free_error_message(r.error_message); + + r = CResult { + success: true, + error_message: ptr::null_mut(), + is_retryable: false, + }; + assert!(!zerobus_stream_get_unacked_records_async( + ptr::null_mut(), + records_cb, + ptr::null_mut(), + &mut r as *mut CResult, + )); + assert!(!r.success); + zerobus_free_error_message(r.error_message); + + r = CResult { + success: true, + error_message: ptr::null_mut(), + is_retryable: false, + }; + assert!(!zerobus_stream_close_async( + ptr::null_mut(), + bool_cb, + ptr::null_mut(), + &mut r as *mut CResult, + )); + assert!(!r.success); + zerobus_free_error_message(r.error_message); + } + // ======================================================================== // C String Conversion Tests // ======================================================================== diff --git a/rust/ffi/zerobus.h b/rust/ffi/zerobus.h index 60a147db..aaf84015 100644 --- a/rust/ffi/zerobus.h +++ b/rust/ffi/zerobus.h @@ -149,6 +149,36 @@ typedef struct CStreamConfigurationOptions { void *ack_user_data; } CStreamConfigurationOptions; +/** + * Function pointer type for async stream creation completion. + * + * `stream` is non-null on success and null on failure. `result` points to a + * `CResult` valid only for the duration of the call; copy any error text during + * the callback if you need to retain it. + * + * Invoked from a background task, so it must be thread-safe and must not + * unwind across the FFI boundary. + */ +typedef void (*CreateStreamAsyncCallback)(struct CZerobusStream *stream, + const struct CResult *result, + void *user_data); + +/** + * Function pointer type for async offset-returning operations. + * + * `result` points to a `CResult` valid only for the duration of the call; copy + * any error text during the callback if you need to retain it. + */ +typedef void (*OffsetAsyncCallback)(int64_t offset, const struct CResult *result, void *user_data); + +/** + * Function pointer type for async bool-returning operations. + * + * `result` points to a `CResult` valid only for the duration of the call; copy + * any error text during the callback if you need to retain it. + */ +typedef void (*BoolAsyncCallback)(bool value, const struct CResult *result, void *user_data); + /** * Represents a single record (either Proto or JSON) */ @@ -166,6 +196,18 @@ typedef struct CRecordArray { uintptr_t len; } CRecordArray; +/** + * Function pointer type for async `CRecordArray`-returning operations. + * + * On success, ownership of `records` transfers to the callback recipient, who + * must free it with `zerobus_free_record_array`. `result` points to a `CResult` + * valid only for the duration of the call; copy any error text during the + * callback if you need to retain it. + */ +typedef void (*RecordArrayAsyncCallback)(struct CRecordArray records, + const struct CResult *result, + void *user_data); + /** * Opaque handle to a table's protobuf schema: its serialized descriptor plus a * prepared encoder. C code only ever holds a pointer to it; the backing @@ -403,6 +445,25 @@ struct CZerobusStream *zerobus_sdk_create_stream(struct CZerobusSdk *sdk, const struct CStreamConfigurationOptions *options, struct CResult *result); +/** + * Create a stream with OAuth authentication on a background task. + * + * Returns `true` once the request has been validated and scheduled. The + * callback is invoked exactly once with either a non-null stream pointer and a + * success result, or a null stream pointer and a failure result. The SDK + * handle must remain valid until the callback runs. + */ +bool zerobus_sdk_create_stream_async(struct CZerobusSdk *sdk, + const char *table_name, + const uint8_t *descriptor_proto_bytes, + uintptr_t descriptor_proto_len, + const char *client_id, + const char *client_secret, + const struct CStreamConfigurationOptions *options, + CreateStreamAsyncCallback callback, + void *user_data, + struct CResult *result); + /** * Create a stream with a custom headers provider callback * This allows you to provide custom authentication headers via a Go callback function @@ -416,6 +477,42 @@ struct CZerobusStream *zerobus_sdk_create_stream_with_headers_provider(struct CZ const struct CStreamConfigurationOptions *options, struct CResult *result); +/** + * Create a stream with a custom headers provider callback on a background task. + * + * Returns `true` once the request has been validated and scheduled. The + * callback is invoked exactly once with either a non-null stream pointer and a + * success result, or a null stream pointer and a failure result. The SDK + * handle must remain valid until the callback runs. + */ +bool zerobus_sdk_create_stream_with_headers_provider_async(struct CZerobusSdk *sdk, + const char *table_name, + const uint8_t *descriptor_proto_bytes, + uintptr_t descriptor_proto_len, + HeadersProviderCallback headers_callback, + void *user_data, + const struct CStreamConfigurationOptions *options, + CreateStreamAsyncCallback callback, + void *callback_user_data, + struct CResult *result); + +/** + * Recreate a stream from an existing stream + * This is used for recovery scenarios where the stream needs to be re-established + */ +struct CZerobusStream *zerobus_sdk_recreate_stream(struct CZerobusSdk *sdk, + struct CZerobusStream *stream, + struct CResult *result); + +/** + * Recreate a stream from an existing stream on a background task. + */ +bool zerobus_sdk_recreate_stream_async(struct CZerobusSdk *sdk, + struct CZerobusStream *stream, + CreateStreamAsyncCallback callback, + void *user_data, + struct CResult *result); + /** * Free a stream instance */ @@ -431,6 +528,16 @@ int64_t zerobus_stream_ingest_proto_record(struct CZerobusStream *stream, uintptr_t data_len, struct CResult *result); +/** + * Ingest a protobuf record on a background task and report the assigned offset via callback. + */ +bool zerobus_stream_ingest_proto_record_async(struct CZerobusStream *stream, + const uint8_t *data, + uintptr_t data_len, + OffsetAsyncCallback callback, + void *user_data, + struct CResult *result); + /** * Ingest a JSON record * Returns the offset directly @@ -440,6 +547,15 @@ int64_t zerobus_stream_ingest_json_record(struct CZerobusStream *stream, const char *json_data, struct CResult *result); +/** + * Ingest a JSON record on a background task and report the assigned offset via callback. + */ +bool zerobus_stream_ingest_json_record_async(struct CZerobusStream *stream, + const char *json_data, + OffsetAsyncCallback callback, + void *user_data, + struct CResult *result); + /** * Ingest a batch of protobuf records * Returns the offset of the last record in the batch, or -1 on error @@ -451,6 +567,17 @@ int64_t zerobus_stream_ingest_proto_records(struct CZerobusStream *stream, uintptr_t num_records, struct CResult *result); +/** + * Ingest a batch of protobuf records on a background task and report the last offset via callback. + */ +bool zerobus_stream_ingest_proto_records_async(struct CZerobusStream *stream, + const uint8_t *const *records, + const uintptr_t *record_lens, + uintptr_t num_records, + OffsetAsyncCallback callback, + void *user_data, + struct CResult *result); + /** * Ingest a batch of JSON records * Returns the offset of the last record in the batch, or -1 on error @@ -461,6 +588,16 @@ int64_t zerobus_stream_ingest_json_records(struct CZerobusStream *stream, uintptr_t num_records, struct CResult *result); +/** + * Ingest a batch of JSON records on a background task and report the last offset via callback. + */ +bool zerobus_stream_ingest_json_records_async(struct CZerobusStream *stream, + const char *const *json_records, + uintptr_t num_records, + OffsetAsyncCallback callback, + void *user_data, + struct CResult *result); + /** * Ingest a protobuf record without waiting for the record to be queued (fire-and-forget). * @@ -524,11 +661,28 @@ bool zerobus_stream_wait_for_offset(struct CZerobusStream *stream, int64_t offset, struct CResult *result); +/** + * Wait for an offset on a background task and report completion via callback. + */ +bool zerobus_stream_wait_for_offset_async(struct CZerobusStream *stream, + int64_t offset, + BoolAsyncCallback callback, + void *user_data, + struct CResult *result); + /** * Flush all pending records */ bool zerobus_stream_flush(struct CZerobusStream *stream, struct CResult *result); +/** + * Flush all pending records on a background task and report completion via callback. + */ +bool zerobus_stream_flush_async(struct CZerobusStream *stream, + BoolAsyncCallback callback, + void *user_data, + struct CResult *result); + /** * Get unacknowledged records from a closed stream * Returns a CRecordArray that must be freed with zerobus_free_record_array @@ -536,6 +690,14 @@ bool zerobus_stream_flush(struct CZerobusStream *stream, struct CResult *result) struct CRecordArray zerobus_stream_get_unacked_records(struct CZerobusStream *stream, struct CResult *result); +/** + * Get unacknowledged records from a closed stream on a background task. + */ +bool zerobus_stream_get_unacked_records_async(struct CZerobusStream *stream, + RecordArrayAsyncCallback callback, + void *user_data, + struct CResult *result); + /** * Free a CRecordArray returned by zerobus_stream_get_unacked_records */ @@ -546,6 +708,19 @@ void zerobus_free_record_array(struct CRecordArray array); */ bool zerobus_stream_close(struct CZerobusStream *stream, struct CResult *result); +/** + * Close the stream gracefully on a background task. + */ +bool zerobus_stream_close_async(struct CZerobusStream *stream, + BoolAsyncCallback callback, + void *user_data, + struct CResult *result); + +/** + * Returns whether the stream has been closed. + */ +bool zerobus_stream_is_closed(struct CZerobusStream *stream); + /** * Free error message string */